icinga2-checks/Checks/Matrix Synapse/check_media_cdn.py

221 lines
8.9 KiB
Python
Raw Normal View History

2023-04-21 23:54:16 -06:00
#!/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)