icinga2-checks/check_media_cdn.py

#!/usr/bin/env python3
import argparse
import asyncio
import json
import os
import sys
import tempfile
import urllib

import numpy as np
import requests
from PIL import Image
from nio import AsyncClient, AsyncClientConfig, LoginResponse, RoomSendError
from urllib3.exceptions import InsecureRequestWarning

from checker import nagios
from checker.synapse_client import send_image, write_login_details_to_disk

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()


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


async def main() -> None:
    async def cleanup(client, test_image_path, image_event_id=None):
        global exit_code
        # Clean up
        if image_event_id:
            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}'")

    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(args.auth_file):
            resp = await client.login(args.pw)

            # check that we logged in successfully
            if isinstance(resp, LoginResponse):
                write_login_details_to_disk(resp, args.hs, args.auth_file)
            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(args.auth_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))
    if isinstance(image_event_id, RoomSendError):
        await cleanup(client, test_image_path)
        print(f'UNKNOWN: failed to send message "{image_event_id}"')
        sys.exit(nagios.UNKNOWN)
    image_event_id = image_event_id.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)

    # Make sure we aren't redirected if we're a Synapse server
    test = requests.head(target_file_url, headers={'User-Agent': 'Synapse/1.77.3'}, allow_redirects=False)
    if test.status_code != 200:
        print('CRITICAL: Synapse user-agent redirected with status code', test.status_code)
        exit_code = nagios.CRITICAL
    else:
        print(f'OK: Synapse user-agent not redirected.')

    await cleanup(client, test_image_path, image_event_id=image_event_id)
    sys.exit(exit_code)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception as e:
        print(f'UNKNOWN: exception "{e}"')
        import traceback

        print(traceback.format_exc())
        sys.exit(nagios.UNKNOWN)