# SPDX-FileCopyrightText: 2025 CyberRex
# SPDX-License-Identifier: MIT

import requests
import os
import argparse

class CFSite:
    def __init__(self, url, plan):
        self.url = url
        self.plan = plan

parser = argparse.ArgumentParser(description='Check Cloudflare Edge Locations and post to Misskey.')
parser.add_argument('--dry-run', action='store_true', help='Run without posting to Misskey')
args = parser.parse_args()

if os.path.exists('config.py'):
    import config
else:
    print("config.py not found")
    exit(1)

test_sites = [
    CFSite('misskey.io', 'Enterprise'),
    CFSite('misskey.systems', 'Pro'),
    CFSite('ohtrmi.cbrx.io', 'Free'),
]

result_txt = []
result_txt.append('Cloudflare Edge Location Check')
result_txt.append('')

# read from /cdn-cgi/trace
for site in test_sites:
    try:
        response = requests.get(f"https://{site.url}/cdn-cgi/trace", timeout=5)
        response.raise_for_status()
        trace_info = response.text
        colo_line = next((line for line in trace_info.splitlines() if line.startswith('colo=')), None)
        if colo_line:
            colo = colo_line.split('=')[1]
            if colo in config.expected_edge_locations:
                result_txt.append(f'- {site.url} ({site.plan}) : {colo}✅️')
            else:
                result_txt.append(f'- {site.url} ({site.plan}) : {colo}⚠')
        else:
            print(f"Could not find 'colo' in the trace info for {site.url}")
    except requests.RequestException as e:
        print(f"Error accessing {site.url}: {e}")

txt = '\n'.join(result_txt)

# post to misskey
payload = {
    'text': txt,
    'visibility': 'home'
}

if not args.dry_run:
    requests.post(f"{config.misskey_url}/api/notes/create", headers={'Authorization': 'Bearer ' + config.misskey_token}, json=payload)
else:
    print(txt)