#!/usr/bin/env python3
"""parts-relay.py — an optional local relay for the parts-inventory tool 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.

The Digi-Key, Mouser and Nexar (Octopart) APIs answer browsers directly (they send CORS headers), so
the tool can call them without this script. Run it when you would rather your API secrets never sit in
the browser: give them to the relay (environment variables or a 0600 file) and it signs the requests
for the page, or when a network or embedded web view blocks the direct calls.

  python3 parts-relay.py                       # forwards; the page sends its own keys
  DIGIKEY_CLIENT_ID=… DIGIKEY_CLIENT_SECRET=… MOUSER_API_KEY=… python3 parts-relay.py
  python3 parts-relay.py --secrets ~/.config/parts-relay.json   # {"digikey_client_id": …, …}

Then, in the tool: Data → Supplier APIs → Route through the local relay (http://127.0.0.1:8098).

Routes (GET and POST):
  /digikey/...   -> https://api.digikey.com/...       (OAuth token + Product Information v4)
  /mouser/...    -> https://api.mouser.com/...        (Search API)
  /nexar/...     -> https://api.nexar.com/...         (GraphQL)
  /nexar-id/...  -> https://identity.nexar.com/...    (Nexar OAuth token)
  /status        -> which secrets the relay holds (never the secrets themselves)

Security: binds 127.0.0.1; answers only pages from https://geekychris.github.io and local previews
(http://127.0.0.1:* / http://localhost:*) plus any --allow-origin; refuses other Host headers (DNS
rebinding); reaches only the four hosts above (not an open proxy); forwards only the headers those
APIs need (never cookies); bodies are limited to 64 KB. Standard library only.
"""
import argparse, json, os, re, ssl, sys, threading, time, urllib.error, urllib.parse, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

VERSION = '1.0'
UA = f'gw-parts-relay/{VERSION} (+https://geekychris.github.io/tools/parts-inventory/)'
UPSTREAMS = {
    'digikey': 'https://api.digikey.com',
    'mouser': 'https://api.mouser.com',
    'nexar': 'https://api.nexar.com',
    'nexar-id': 'https://identity.nexar.com',
}
FORWARD_HEADERS = ('authorization', 'content-type', 'accept', 'x-digikey-client-id', 'x-digikey-locale-site', 'x-digikey-locale-language',
                   'x-digikey-locale-currency', 'x-digikey-locale-shiptocountry', 'x-digikey-customer-id')
EXPOSE = 'X-RateLimit-Remaining, X-RateLimit-Limit, X-BurstLimit-Remaining, Retry-After, X-Relay-Upstream'
MAX_BODY = 64 * 1024
SAFE_PATH = re.compile(r"^[A-Za-z0-9._~!$&'()*+,;=:@%/?-]*$")
ALLOWED = {'https://geekychris.github.io'}
urlopen = urllib.request.urlopen   # tests replace this


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


def origin_ok(o):
    return o in ALLOWED or bool(re.match(r'^http://(127\.0\.0\.1|localhost)(:\d+)?$', o or ''))


def upstream_url(route, rest):
    """The upstream URL for /<route>/<rest>, or None when it is not allowed. The host is fixed per route;
    nothing in `rest` can change it (no '..', no '//', no '@', no scheme, only URL-safe characters)."""
    base = UPSTREAMS.get(route)
    if not base or not SAFE_PATH.match(rest or ''):
        return None
    path = (rest or '').split('?')[0]
    if path.startswith('/') or '//' in path or '@' in path or '\\' in path:
        return None
    segs = urllib.parse.unquote(path).split('/')
    if any(s in ('..', '.') for s in segs) or '://' in urllib.parse.unquote(rest or ''):
        return None
    url = base + '/' + (rest or '')
    if urllib.parse.urlsplit(url).hostname != urllib.parse.urlsplit(base).hostname:
        return None
    return url


class Secrets:
    """API secrets held by the relay (optional). Tokens are fetched and cached here."""

    def __init__(self, d):
        self.dk_id, self.dk_secret = d.get('digikey_client_id'), d.get('digikey_client_secret')
        self.mouser_key = d.get('mouser_api_key')
        self.nx_id, self.nx_secret = d.get('nexar_client_id'), d.get('nexar_client_secret')
        self.tokens, self.lock = {}, threading.Lock()

    def holds(self):
        return {'digikey': bool(self.dk_id and self.dk_secret), 'mouser': bool(self.mouser_key), 'nexar': bool(self.nx_id and self.nx_secret)}

    def token(self, which):
        with self.lock:
            t = self.tokens.get(which)
            if t and t[1] > time.time() + 60:
                return t[0]
        if which == 'digikey':
            url, cid, sec, extra = UPSTREAMS['digikey'] + '/v1/oauth2/token', self.dk_id, self.dk_secret, ''
        else:
            url, cid, sec, extra = UPSTREAMS['nexar-id'] + '/connect/token', self.nx_id, self.nx_secret, '&scope=supply.domain'
        body = f'client_id={urllib.parse.quote(cid)}&client_secret={urllib.parse.quote(sec)}&grant_type=client_credentials{extra}'.encode()
        req = urllib.request.Request(url, data=body, method='POST', headers={'Content-Type': 'application/x-www-form-urlencoded', 'Accept': 'application/json', 'User-Agent': UA})
        with urlopen(req, timeout=20, context=TLS) as r:
            j = json.loads(r.read() or b'{}')
        tok = j.get('access_token')
        if not tok:
            raise RuntimeError(f'{which}: no access token in the response')
        with self.lock:
            self.tokens[which] = (tok, time.time() + float(j.get('expires_in') or 599))
        return tok


class Relay(BaseHTTPRequestHandler):
    server_version = 'parts-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 cors(self):
        o = self.headers.get('Origin', '')
        if origin_ok(o):
            self.send_header('Access-Control-Allow-Origin', o)
            self.send_header('Vary', 'Origin')
            self.send_header('Access-Control-Allow-Private-Network', 'true')
            self.send_header('Access-Control-Expose-Headers', EXPOSE)

    def reply(self, code, body, ctype='application/json', extra=None):
        data = body if isinstance(body, bytes) else body.encode()
        self.send_response(code)
        self.cors()
        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 refuse(self):
        """Origin and Host checks, before anything else happens. Returns True when refused."""
        host = (self.headers.get('Host') or '').lower()
        port = self.server.server_address[1]
        if host not in (f'127.0.0.1:{port}', f'localhost:{port}', f'[::1]:{port}'):
            self.reply(403, json.dumps({'error': f'Host {host or "(none)"} not allowed: use http://127.0.0.1:{port}/'}))
            return True
        o = self.headers.get('Origin')
        if o is not None and not origin_ok(o):
            self.reply(403, json.dumps({'error': f'origin {o} not allowed (see --allow-origin)'}))
            return True
        if self.path.startswith(('http:', 'https:', '//')) or not self.path.startswith('/'):
            self.reply(400, json.dumps({'error': 'absolute URLs are not accepted: this is not an open proxy'}))
            return True
        return False

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

    def do_GET(self):
        self.forward('GET')

    def do_POST(self):
        self.forward('POST')

    def do_PUT(self):
        self.reply(405, json.dumps({'error': 'GET and POST only'}), extra={'Allow': 'GET, POST, OPTIONS'})

    do_DELETE = do_PATCH = do_PUT

    def read_body(self):
        """The request body (Content-Length or chunked), None when empty, False when over MAX_BODY."""
        if 'chunked' in (self.headers.get('Transfer-Encoding') or '').lower():
            out = b''
            while True:
                size = int((self.rfile.readline(64).split(b';')[0].strip() or b'0'), 16)
                if size == 0:
                    while self.rfile.readline(1024) not in (b'\r\n', b'\n', b''):
                        pass
                    break
                if len(out) + size > MAX_BODY:
                    self.close_connection = True
                    return False
                out += self.rfile.read(size)
                self.rfile.readline(8)
            return out or None
        n = int(self.headers.get('Content-Length') or 0)
        if n > MAX_BODY:
            self.close_connection = True
            return False
        return self.rfile.read(n) if n else None

    def forward(self, method):
        if self.refuse():
            return
        first, _, rest = self.path.lstrip('/').partition('/')
        first = first.split('?')[0]
        if first in ('', 'status'):
            return self.reply(200, json.dumps({'relay': 'parts-relay', 'version': VERSION, 'upstreams': UPSTREAMS, 'holds': self.server.secrets.holds()}))
        url = upstream_url(first, rest)
        if not url:
            return self.reply(404 if first not in UPSTREAMS else 400, json.dumps({'error': 'unknown or unsafe route', 'routes': list(UPSTREAMS)}))
        body = self.read_body()
        if body is False:
            return self.reply(413, json.dumps({'error': 'request body too large'}))
        headers = {k: v for k, v in self.headers.items() if k.lower() in FORWARD_HEADERS}
        headers['User-Agent'] = UA
        sec = self.server.secrets
        try:
            if first == 'digikey' and sec.holds()['digikey'] and not rest.startswith('v1/oauth2/'):
                headers = {k: v for k, v in headers.items() if k.lower() not in ('authorization', 'x-digikey-client-id')}
                headers['Authorization'] = 'Bearer ' + sec.token('digikey')
                headers['X-DIGIKEY-Client-Id'] = sec.dk_id
            if first == 'mouser' and sec.holds()['mouser'] and 'apikey=' not in url.lower():
                url += ('&' if '?' in url else '?') + 'apiKey=' + urllib.parse.quote(sec.mouser_key)
            if first == 'nexar' and sec.holds()['nexar'] and 'authorization' not in {k.lower() for k in headers}:
                headers['Authorization'] = 'Bearer ' + sec.token('nexar')
        except Exception as e:
            return self.reply(502, json.dumps({'error': f'could not sign in with the relay\'s secrets: {type(e).__name__}: {e}'}))
        req = urllib.request.Request(url, data=body, method=method, headers=headers)
        try:
            with urlopen(req, timeout=25, context=TLS) as r:
                code, data, ctype = r.status, r.read(), r.headers.get('Content-Type', 'application/json')
                keep = {k: r.headers[k] for k in ('X-RateLimit-Remaining', 'X-RateLimit-Limit', 'X-BurstLimit-Remaining', 'Retry-After') if r.headers.get(k)}
        except urllib.error.HTTPError as e:
            code, data, ctype = e.code, e.read() or b'', e.headers.get('Content-Type', 'application/json')
            keep = {k: e.headers[k] for k in ('Retry-After', 'X-RateLimit-Remaining') if e.headers.get(k)}
        except Exception as e:  # network errors, timeouts
            return self.reply(502, json.dumps({'error': f'{type(e).__name__}: {e}', 'upstream': UPSTREAMS[first]}))
        self.reply(code, data, ctype, {'X-Relay-Upstream': first, **keep})


def load_secrets(path):
    d = {
        'digikey_client_id': os.environ.get('DIGIKEY_CLIENT_ID'), 'digikey_client_secret': os.environ.get('DIGIKEY_CLIENT_SECRET'),
        'mouser_api_key': os.environ.get('MOUSER_API_KEY'),
        'nexar_client_id': os.environ.get('NEXAR_CLIENT_ID'), 'nexar_client_secret': os.environ.get('NEXAR_CLIENT_SECRET'),
    }
    if path:
        p = os.path.expanduser(path)
        if os.name == 'posix' and os.stat(p).st_mode & 0o077:
            sys.stderr.write(f'warning: {p} is readable by other users; chmod 600 it\n')
        with open(p) as f:
            for k, v in json.load(f).items():
                if k in d and v:
                    d[k] = v
    return d


def main(argv=None):
    p = argparse.ArgumentParser(description='Local relay for the geekychris.github.io parts-inventory tool (Digi-Key, Mouser, Nexar).')
    p.add_argument('--port', type=int, default=8098)
    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('--secrets', help='JSON file with digikey_client_id, digikey_client_secret, mouser_api_key, nexar_client_id, nexar_client_secret')
    p.add_argument('--allow-origin', action='append', default=[], help='another web origin allowed to use the relay')
    p.add_argument('--quiet', action='store_true')
    a = p.parse_args(argv)
    ALLOWED.update(a.allow_origin)
    srv = ThreadingHTTPServer((a.host, a.port), Relay)
    srv.daemon_threads = True
    srv.quiet = a.quiet
    srv.secrets = Secrets(load_secrets(a.secrets))
    held = [k for k, v in srv.secrets.holds().items() if v]
    print(f'parts-relay {VERSION} on http://{a.host}:{srv.server_address[1]}/  (Ctrl+C to stop)', flush=True)
    print('  holds secrets for:', ', '.join(held) or 'nothing (the page sends its own keys)', flush=True)
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    main()