221 lines
8.9 KiB
Python
221 lines
8.9 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
import argparse
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
import urllib
|
||
|
|
|
||
|
|
import aiofiles.os
|
||
|
|
import magic
|
||
|
|
import numpy as np
|
||
|
|
import requests
|
||
|
|
from PIL import Image
|
||
|
|
from nio import AsyncClient, AsyncClientConfig, LoginResponse, UploadResponse
|
||
|
|
from urllib3.exceptions import InsecureRequestWarning
|
||
|
|
|
||
|
|
import nagios
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser(description='')
|
||
|
|
parser.add_argument('--user', required=True, help='User ID for the bot.')
|
||
|
|
parser.add_argument('--pw', required=True, help='Password for the bot.')
|
||
|
|
parser.add_argument('--hs', required=True, help='Homeserver of the bot.')
|
||
|
|
parser.add_argument('--admin-endpoint', required=True, help='Admin endpoint that will be called to purge media for this user.')
|
||
|
|
parser.add_argument('--room', required=True, help='The room the bot should send its test messages in.')
|
||
|
|
parser.add_argument('--media-cdn-domain', required=True, help='The domain to make sure it redirects to.')
|
||
|
|
parser.add_argument('--auth-file', help="File to cache the bot's login details to.")
|
||
|
|
parser.add_argument('--timeout', type=float, default=90, help='Request timeout limit.')
|
||
|
|
parser.add_argument('--warn', type=float, default=2.0, help='Manually set warn level.')
|
||
|
|
parser.add_argument('--crit', type=float, default=2.5, help='Manually set critical level.')
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
CONFIG_FILE = args.auth_file
|
||
|
|
|
||
|
|
|
||
|
|
def verify_media_header(header: str, header_dict: dict, good_value: str = None, warn_value: str = None, critical_value: str = None):
|
||
|
|
"""
|
||
|
|
If you don't specify good_value, warn_value, or critical_value then the header will only be checked for existience.
|
||
|
|
"""
|
||
|
|
# Convert everything to strings to prevent any wierdness
|
||
|
|
header_value = str(header_dict.get(header))
|
||
|
|
good_value = str(good_value)
|
||
|
|
warn_value = str(warn_value)
|
||
|
|
critical_value = str(critical_value)
|
||
|
|
if not header_value:
|
||
|
|
return f'CRITICAL: missing header "{header}"', nagios.CRITICAL
|
||
|
|
elif good_value and header_value == good_value:
|
||
|
|
return f'OK: {header}: "{header_value}"', nagios.OK
|
||
|
|
elif warn_value and header_value == warn_value:
|
||
|
|
return f'WARN: {header}: "{header_value}"', nagios.WARNING
|
||
|
|
elif critical_value and header_value == critical_value:
|
||
|
|
return f'CRITICAL: {header}: "{header_value}"', nagios.CRITICAL
|
||
|
|
return f'OK: {header} is present with value "{header_value}"', nagios.OK
|
||
|
|
|
||
|
|
|
||
|
|
def write_details_to_disk(resp: LoginResponse, homeserver) -> None:
|
||
|
|
"""Writes the required login details to disk so we can log in later without
|
||
|
|
using a password.
|
||
|
|
Arguments:
|
||
|
|
resp {LoginResponse} -- the successful client login response.
|
||
|
|
homeserver -- URL of homeserver, e.g. "https://matrix.example.org"
|
||
|
|
"""
|
||
|
|
# open the config file in write-mode
|
||
|
|
with open(CONFIG_FILE, "w") as f:
|
||
|
|
# write the login details to disk
|
||
|
|
json.dump({"homeserver": homeserver, # e.g. "https://matrix.example.org"
|
||
|
|
"user_id": resp.user_id, # e.g. "@user:example.org"
|
||
|
|
"device_id": resp.device_id, # device ID, 10 uppercase letters
|
||
|
|
"access_token": resp.access_token, # cryptogr. access token
|
||
|
|
}, f, )
|
||
|
|
|
||
|
|
|
||
|
|
async def send_image(client, room_id, image):
|
||
|
|
"""Send image to room.
|
||
|
|
Arguments:
|
||
|
|
---------
|
||
|
|
client : Client
|
||
|
|
room_id : str
|
||
|
|
image : str, file name of image
|
||
|
|
This is a working example for a JPG image.
|
||
|
|
"content": {
|
||
|
|
"body": "someimage.jpg",
|
||
|
|
"info": {
|
||
|
|
"size": 5420,
|
||
|
|
"mimetype": "image/jpeg",
|
||
|
|
"thumbnail_info": {
|
||
|
|
"w": 100,
|
||
|
|
"h": 100,
|
||
|
|
"mimetype": "image/jpeg",
|
||
|
|
"size": 2106
|
||
|
|
},
|
||
|
|
"w": 100,
|
||
|
|
"h": 100,
|
||
|
|
"thumbnail_url": "mxc://example.com/SomeStrangeThumbnailUriKey"
|
||
|
|
},
|
||
|
|
"msgtype": "m.image",
|
||
|
|
"url": "mxc://example.com/SomeStrangeUriKey"
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
mime_type = magic.from_file(image, mime=True) # e.g. "image/jpeg"
|
||
|
|
if not mime_type.startswith("image/"):
|
||
|
|
print(f'UNKNOWN: wrong mime type "{mime_type}"')
|
||
|
|
sys.exit(nagios.UNKNOWN)
|
||
|
|
|
||
|
|
im = Image.open(image)
|
||
|
|
(width, height) = im.size # im.size returns (width,height) tuple
|
||
|
|
|
||
|
|
# first do an upload of image, then send URI of upload to room
|
||
|
|
file_stat = await aiofiles.os.stat(image)
|
||
|
|
async with aiofiles.open(image, "r+b") as f:
|
||
|
|
resp, maybe_keys = await client.upload(f, content_type=mime_type, # image/jpeg
|
||
|
|
filename=os.path.basename(image), filesize=file_stat.st_size, )
|
||
|
|
if not isinstance(resp, UploadResponse):
|
||
|
|
print(f'UNKNOWN: failed to upload image "{resp}"')
|
||
|
|
sys.exit(nagios.UNKNOWN)
|
||
|
|
|
||
|
|
content = {"body": os.path.basename(image), # descriptive title
|
||
|
|
"info": {"size": file_stat.st_size, "mimetype": mime_type, "thumbnail_info": None, # TODO
|
||
|
|
"w": width, # width in pixel
|
||
|
|
"h": height, # height in pixel
|
||
|
|
"thumbnail_url": None, # TODO
|
||
|
|
}, "msgtype": "m.image", "url": resp.content_uri, }
|
||
|
|
|
||
|
|
try:
|
||
|
|
return await client.room_send(room_id, message_type="m.room.message", content=content)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Image send of file {image} failed.")
|
||
|
|
print(f'UNKNOWN: failed to send image event "{e}"')
|
||
|
|
sys.exit(nagios.UNKNOWN)
|
||
|
|
|
||
|
|
|
||
|
|
async def main() -> None:
|
||
|
|
client = AsyncClient(args.hs, args.user, config=AsyncClientConfig(request_timeout=args.timeout, max_timeout_retry_wait_time=10))
|
||
|
|
if args.auth_file:
|
||
|
|
# If there are no previously-saved credentials, we'll use the password
|
||
|
|
if not os.path.exists(CONFIG_FILE):
|
||
|
|
resp = await client.login(args.pw)
|
||
|
|
|
||
|
|
# check that we logged in successfully
|
||
|
|
if isinstance(resp, LoginResponse):
|
||
|
|
write_details_to_disk(resp, args.hs)
|
||
|
|
else:
|
||
|
|
print(f'UNKNOWN: failed to log in "{resp}"')
|
||
|
|
sys.exit(nagios.UNKNOWN)
|
||
|
|
else:
|
||
|
|
# Otherwise the config file exists, so we'll use the stored credentials
|
||
|
|
with open(CONFIG_FILE, "r") as f:
|
||
|
|
config = json.load(f)
|
||
|
|
client = AsyncClient(config["homeserver"])
|
||
|
|
client.access_token = config["access_token"]
|
||
|
|
client.user_id = config["user_id"]
|
||
|
|
client.device_id = config["device_id"]
|
||
|
|
else:
|
||
|
|
await client.login(args.pw)
|
||
|
|
|
||
|
|
await client.join(args.room)
|
||
|
|
|
||
|
|
# Create a random image
|
||
|
|
imarray = np.random.rand(100, 100, 3) * 255
|
||
|
|
im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
|
||
|
|
_, test_image_path = tempfile.mkstemp()
|
||
|
|
test_image_path = test_image_path + '.png'
|
||
|
|
im.save(test_image_path)
|
||
|
|
|
||
|
|
# Send the image and get the event ID
|
||
|
|
image_event_id = (await send_image(client, args.room, test_image_path)).event_id
|
||
|
|
|
||
|
|
# Get the event
|
||
|
|
image_event = (await client.room_get_event(args.room, image_event_id)).event
|
||
|
|
|
||
|
|
# convert mxc:// to http://
|
||
|
|
target_file_url = await client.mxc_to_http(image_event.url)
|
||
|
|
|
||
|
|
# Check the headers. Ignore the non-async thing here, it doesn't
|
||
|
|
# matter in this situation.
|
||
|
|
headers = dict(requests.head(target_file_url).headers)
|
||
|
|
|
||
|
|
exit_code = nagios.OK
|
||
|
|
|
||
|
|
# Check domain
|
||
|
|
domain = urllib.parse.urlparse(headers['location']).netloc
|
||
|
|
if domain != args.media_cdn_domain:
|
||
|
|
exit_code = nagios.CRITICAL
|
||
|
|
print(f'CRITICAL: media CDN domain is "{domain}"')
|
||
|
|
else:
|
||
|
|
print(f'OK: media CDN domain is "{domain}"')
|
||
|
|
|
||
|
|
results = [verify_media_header('synapse-media-local-status', headers), verify_media_header('synapse-media-s3-status', headers, good_value='200'), verify_media_header('synapse-media-server', headers, good_value='s3'),
|
||
|
|
verify_media_header('Server', headers, good_value='cloudflare')]
|
||
|
|
for header_chk, code in results:
|
||
|
|
if code != nagios.OK:
|
||
|
|
exit_code = code
|
||
|
|
print(header_chk)
|
||
|
|
|
||
|
|
# Clean up
|
||
|
|
await client.room_redact(args.room, image_event_id)
|
||
|
|
os.remove(test_image_path)
|
||
|
|
await client.close()
|
||
|
|
|
||
|
|
requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)
|
||
|
|
try:
|
||
|
|
r = requests.delete(f'{args.admin_endpoint}/_synapse/admin/v1/users/{args.user}/media', headers={'Authorization': f'Bearer {client.access_token}'}, verify=False)
|
||
|
|
if r.status_code != 200:
|
||
|
|
if nagios.WARNING < exit_code:
|
||
|
|
exit_code = nagios.WARNING
|
||
|
|
print(f"WARN: failed to purge media for this user, request failed with '{r.text}'")
|
||
|
|
except Exception as e:
|
||
|
|
if nagios.WARNING < exit_code:
|
||
|
|
exit_code = nagios.WARNING
|
||
|
|
print(f"WARN: failed to purge media for this user '{e}'")
|
||
|
|
|
||
|
|
sys.exit(exit_code)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
try:
|
||
|
|
asyncio.run(main())
|
||
|
|
except Exception as e:
|
||
|
|
print(f'UNKNOWN: exception "{e}"')
|
||
|
|
sys.exit(nagios.UNKNOWN)
|