#!/usr/bin/env python3
"""aviation-relay.py — a tiny local relay for the aviation tools on geekychris.github.io.

NOTE: the gitwebsite app (and `gw serve`) now has this built in — its *companion* gateway on
127.0.0.1:7317 relays the same things, and the tool uses it first when the app is running. This
script stays as the fallback for when the app is not running.

Web pages can only read another server's data when that server sends CORS headers, and an https page
may not fetch plain-http addresses on your network (mixed content). Public ADS-B and aviation-weather
APIs send no CORS headers, and a Raspberry Pi receiver serves plain http. This script runs on YOUR
computer, listens on http://127.0.0.1:8099 (browsers treat localhost as a secure origin, so an https
page may call it), fetches what the tools ask for, and adds the CORS headers. Standard library only.

  python3 aviation-relay.py                                   # public feeds + weather only
  python3 aviation-relay.py --receiver http://piaware.local:8080 --beast piaware.local:30005 --raw piaware.local:30002

Routes (GET only):
  /adsblol/...   -> https://api.adsb.lol/...            (ODbL; be modest)
  /adsbfi/...    -> https://opendata.adsb.fi/...        (personal, non-commercial; max 1 request/s)
  /opensky/...   -> https://opensky-network.org/...     (anonymous: 400 credits/day)
  /awc/...       -> https://aviationweather.gov/...     (METAR/TAF Data API)
  /receiver/...  -> <--receiver>/...                    (e.g. /receiver/data/aircraft.json)
  /beast  (WebSocket) -> TCP <--beast>  binary Beast frames, passed through unchanged
  /raw    (WebSocket) -> TCP <--raw>    AVR text (*8D...;) passed through unchanged
Each upstream response is cached for 2 s so several tabs do not multiply the load.
Only these upstreams are reachable: it is not an open proxy. It binds to 127.0.0.1 unless you pass
--host (do not expose it to the internet).
"""
import argparse, base64, hashlib, json, os, socket, ssl, struct, sys, threading, time, urllib.error, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

VERSION = '1.0'
UA = f'gw-aviation-relay/{VERSION} (+https://geekychris.github.io/tools/flight-tracker/)'
UPSTREAMS = {
    'adsblol': 'https://api.adsb.lol',
    'adsbfi': 'https://opendata.adsb.fi',
    'opensky': 'https://opensky-network.org',
    'awc': 'https://aviationweather.gov',
}
CACHE_S = 2.0
_cache, _lock = {}, threading.Lock()
WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'


def tls_context():
    # python.org builds on macOS ship without root certificates; fall back to certifi or the system bundle.
    ctx = ssl.create_default_context()
    try:
        import certifi
        ctx.load_verify_locations(certifi.where())
    except Exception:
        for f in ('/etc/ssl/cert.pem', '/etc/ssl/certs/ca-certificates.crt'):
            if os.path.exists(f):
                try:
                    ctx.load_verify_locations(f)
                except Exception:
                    pass
    return ctx


TLS = tls_context()


# Only these pages may read through the relay: the site, and local previews (the gitwebsite app).
# Add your own with --allow-origin.
ALLOWED = {'https://geekychris.github.io'}


def origin_ok(o):
    return o in ALLOWED or o.startswith(('http://127.0.0.1:', 'http://localhost:')) or o in ('http://127.0.0.1', 'http://localhost')


def cors(h):
    o = h.headers.get('Origin', '')
    if origin_ok(o):
        h.send_header('Access-Control-Allow-Origin', o)
        h.send_header('Vary', 'Origin')
    h.send_header('Access-Control-Allow-Private-Network', 'true')
    h.send_header('Access-Control-Expose-Headers', 'X-Rate-Limit-Remaining, X-Relay-Upstream, X-Relay-Cache')


class Relay(BaseHTTPRequestHandler):
    server_version = 'aviation-relay/' + VERSION
    protocol_version = 'HTTP/1.1'

    def log_message(self, fmt, *args):
        if not self.server.quiet:
            sys.stderr.write('%s %s\n' % (time.strftime('%H:%M:%S'), fmt % args))

    def do_OPTIONS(self):
        self.send_response(204)
        cors(self)
        self.send_header('Access-Control-Allow-Methods', 'GET, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', '*')
        self.send_header('Access-Control-Max-Age', '600')
        self.send_header('Content-Length', '0')
        self.end_headers()

    def reply(self, code, body, ctype='application/json', extra=None):
        data = body if isinstance(body, bytes) else body.encode()
        self.send_response(code)
        cors(self)
        self.send_header('Content-Type', ctype)
        self.send_header('Cache-Control', 'no-store')
        for k, v in (extra or {}).items():
            self.send_header(k, v)
        self.send_header('Content-Length', str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        path = self.path
        if self.headers.get('Upgrade', '').lower() == 'websocket':
            name = path.split('?')[0].strip('/')
            target = {'beast': self.server.beast, 'raw': self.server.raw}.get(name)
            if not target:
                return self.reply(404, json.dumps({'error': f'no --{name} receiver port configured'}))
            return self.websocket(target)
        first, _, rest = path.lstrip('/').partition('/')
        if first == '' or first == 'status':
            return self.reply(200, json.dumps({
                'relay': 'aviation-relay', 'version': VERSION, 'upstreams': list(UPSTREAMS),
                'receiver': self.server.receiver, 'beast': self.server.beast, 'raw': self.server.raw}))
        base = self.server.receiver if first == 'receiver' else UPSTREAMS.get(first)
        if not base:
            return self.reply(404, json.dumps({'error': 'unknown route', 'routes': list(UPSTREAMS) + ['receiver']}))
        url = base.rstrip('/') + '/' + rest
        now = time.time()
        with _lock:
            hit = _cache.get(url)
        if hit and now - hit[0] < CACHE_S:
            return self.reply(hit[1], hit[2], hit[3], {'X-Relay-Upstream': first, 'X-Relay-Cache': 'hit', **hit[4]})
        try:
            req = urllib.request.Request(url, headers={'User-Agent': UA, 'Accept': '*/*'})
            with urllib.request.urlopen(req, timeout=15, context=TLS if url.startswith('https') else None) as r:
                code, body, ctype = r.status, r.read(), r.headers.get('Content-Type', 'application/octet-stream')
                keep = {k: r.headers[k] for k in ('X-Rate-Limit-Remaining',) if r.headers.get(k)}
        except urllib.error.HTTPError as e:
            code, body, ctype, keep = e.code, e.read() or b'', e.headers.get('Content-Type', 'text/plain'), {}
        except Exception as e:  # network errors, timeouts
            return self.reply(502, json.dumps({'error': f'{type(e).__name__}: {e}', 'upstream': url}))
        with _lock:
            _cache[url] = (now, code, body, ctype, keep)
            for k in [k for k, v in _cache.items() if now - v[0] > 60]:
                del _cache[k]
        self.reply(code, body, ctype, {'X-Relay-Upstream': first, 'X-Relay-Cache': 'miss', **keep})

    # ---------------------------------------------------------------- WebSocket (RFC 6455, server side)
    def websocket(self, target):
        o = self.headers.get('Origin', '')
        if o and not origin_ok(o):   # browsers always send Origin; refuse other web pages
            return self.reply(403, json.dumps({'error': f'origin {o} not allowed (see --allow-origin)'}))
        key = self.headers.get('Sec-WebSocket-Key', '')
        accept = base64.b64encode(hashlib.sha1((key + WS_GUID).encode()).digest()).decode()
        host, _, port = target.rpartition(':')
        try:
            up = socket.create_connection((host, int(port)), timeout=10)
        except Exception as e:
            return self.reply(502, json.dumps({'error': f'cannot connect to {target}: {e}'}))
        up.settimeout(None)
        self.send_response(101, 'Switching Protocols')
        self.send_header('Upgrade', 'websocket')
        self.send_header('Connection', 'Upgrade')
        self.send_header('Sec-WebSocket-Accept', accept)
        self.end_headers()
        self.wfile.flush()
        sock = self.connection
        opcode = 0x2 if target == self.server.beast else 0x1
        alive = [True]

        def frame(op, payload):
            n = len(payload)
            head = bytes([0x80 | op]) + (bytes([n]) if n < 126 else struct.pack('!BH', 126, n) if n < 65536 else struct.pack('!BQ', 127, n))
            sock.sendall(head + payload)

        def from_client():   # read (masked) client frames: answer pings, stop on close
            try:
                while alive[0]:
                    h = self.rfile.read(2)
                    if len(h) < 2:
                        break
                    op, n = h[0] & 0x0F, h[1] & 0x7F
                    if n == 126:
                        n = struct.unpack('!H', self.rfile.read(2))[0]
                    elif n == 127:
                        n = struct.unpack('!Q', self.rfile.read(8))[0]
                    mask = self.rfile.read(4) if h[1] & 0x80 else b'\0\0\0\0'
                    data = bytes(b ^ mask[i % 4] for i, b in enumerate(self.rfile.read(n)))
                    if op == 0x8:
                        break
                    if op == 0x9:
                        frame(0xA, data)
            except Exception:
                pass
            alive[0] = False
            try:
                up.shutdown(socket.SHUT_RDWR)
            except Exception:
                pass

        threading.Thread(target=from_client, daemon=True).start()
        try:
            while alive[0]:
                chunk = up.recv(16384)
                if not chunk:
                    break
                if opcode == 0x1:
                    chunk = chunk.decode('ascii', 'replace').encode()
                frame(opcode, chunk)
        except Exception:
            pass
        alive[0] = False
        try:
            frame(0x8, b'')
        except Exception:
            pass
        up.close()
        self.close_connection = True


def main():
    p = argparse.ArgumentParser(description='Local CORS relay for the geekychris.github.io aviation tools.')
    p.add_argument('--port', type=int, default=8099)
    p.add_argument('--host', default='127.0.0.1', help='address to listen on (default 127.0.0.1 — keep it local)')
    p.add_argument('--receiver', help='your receiver web base, e.g. http://piaware.local:8080 or http://pi.local/tar1090')
    p.add_argument('--beast', help='receiver Beast output host:port, e.g. piaware.local:30005')
    p.add_argument('--raw', help='receiver raw AVR output host:port, e.g. piaware.local:30002')
    p.add_argument('--quiet', action='store_true')
    p.add_argument('--allow-origin', action='append', default=[], help='another web origin allowed to use the relay')
    a = p.parse_args()
    ALLOWED.update(a.allow_origin)
    srv = ThreadingHTTPServer((a.host, a.port), Relay)
    srv.daemon_threads = True
    srv.receiver, srv.beast, srv.raw, srv.quiet = a.receiver, a.beast, a.raw, a.quiet
    print(f'aviation-relay {VERSION} on http://{a.host}:{a.port}/  (Ctrl+C to stop)')
    print('  receiver:', a.receiver or '-', ' beast:', a.beast or '-', ' raw:', a.raw or '-')
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    main()
