#!/usr/bin/env python3
"""Checks for the INBOX push feature (inbox_server.py + inbox_push.py), on a scratch store, no network:
the server runs with INBOX_PUSH_FAKE=1 (sends are recorded in push_fake.jsonl, endpoints ending -gone answer 410,
-missing 404, -down 0) and a short INBOX_PUSH_GAP. Every recorded push is decrypted with the fake device's key.

  python3 _tools/inbox_push_test.py [--keep]      (prints one row per check; exit 1 on a failure)
"""
import base64, hashlib, hmac, json, os, shutil, socket, struct, subprocess, sys, tempfile, time
from urllib.request import Request, urlopen
from urllib.error import HTTPError

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

TOOLS = os.path.dirname(os.path.abspath(__file__))
SERVER = os.path.join(TOOLS, 'inbox_server.py')
CLI = os.path.join(TOOLS, 'inbox.sh')
GAP = 8
ROWS = []


def row(name, ok, info=''):
    ROWS.append((name, ok))
    print('%-4s %-58s %s' % ('PASS' if ok else 'FAIL', name, info))


def b64u(b):
    return base64.urlsafe_b64encode(b).rstrip(b'=').decode()


def _h(k, m):
    return hmac.new(k, m, hashlib.sha256).digest()


class Device:
    """A fake browser subscription with real keys, so pushes to it can be decrypted."""
    def __init__(self, endpoint):
        self.endpoint = endpoint
        self.priv = ec.generate_private_key(ec.SECP256R1())
        self.pub = self.priv.public_key().public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint)
        self.auth = os.urandom(16)

    def sub(self):
        return {'endpoint': self.endpoint, 'keys': {'p256dh': b64u(self.pub), 'auth': b64u(self.auth)}}

    def decrypt(self, body):
        salt, rs, idlen = body[:16], struct.unpack('!I', body[16:20])[0], body[20]
        as_pub = body[21:21 + idlen]
        ct = body[21 + idlen:]
        assert rs == 4096 and idlen == 65
        ecdh = self.priv.exchange(ec.ECDH(), ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), as_pub))
        ikm = _h(_h(self.auth, ecdh), b'WebPush: info\x00' + self.pub + as_pub + b'\x01')
        prk = _h(salt, ikm)
        cek = _h(prk, b'Content-Encoding: aes128gcm\x00\x01')[:16]
        nonce = _h(prk, b'Content-Encoding: nonce\x00\x01')[:12]
        pt = AESGCM(cek).decrypt(nonce, ct, None).rstrip(b'\x00')
        assert pt.endswith(b'\x02')
        return json.loads(pt[:-1])


def free_port():
    s = socket.socket()
    s.bind(('127.0.0.1', 0))
    p = s.getsockname()[1]
    s.close()
    return p


def main():
    keep = '--keep' in sys.argv
    d = tempfile.mkdtemp(prefix='inbox-push-test-')
    env = dict(os.environ, INBOX_DIR=d, INBOX_PUSH_FAKE='1', INBOX_PUSH_GAP=str(GAP))
    env.pop('INBOX_AS', None)
    tok = b64u(os.urandom(12))
    with open(os.path.join(d, 'token'), 'w') as f:
        f.write(tok)
    fake = os.path.join(d, 'push_fake.jsonl')

    def cli(*args, stdin=''):
        r = subprocess.run([sys.executable, SERVER] + list(args), input=stdin, env=env, capture_output=True, text=True, timeout=30)
        return r

    def sh(*args):
        e = dict(env, INBOX_LOCAL_DIR=d)
        return subprocess.run(['bash', CLI] + list(args), env=e, capture_output=True, text=True, timeout=30)

    r = cli('push-init')
    pub = json.loads(r.stdout or '{}').get('publicKey', '')
    pem = open(os.path.join(d, 'vapid_private.pem')).read()
    row('push-init makes the VAPID key, prints only the public key',
        r.returncode == 0 and len(pub) == 87 and 'PRIVATE' not in r.stdout and 'PRIVATE' in pem)
    row('vapid_private.pem is 0600', oct(os.stat(os.path.join(d, 'vapid_private.pem')).st_mode & 0o777) == '0o600')
    r2 = cli('push-init')
    row('push-init again keeps the same key', json.loads(r2.stdout).get('publicKey') == pub)
    row('RFC 8291 example + VAPID JWT self test',
        subprocess.run([sys.executable, os.path.join(TOOLS, 'inbox_push.py'), 'selftest'], capture_output=True).returncode == 0)

    port = free_port()
    srv = subprocess.Popen([sys.executable, SERVER, 'serve', str(port)], env=env, stderr=subprocess.DEVNULL)
    base = 'http://127.0.0.1:%d/api/inbox' % port

    def http(method, path, body=None, token=tok):
        h = {'Content-Type': 'application/json'}
        if token is not None:
            h['X-Inbox-Token'] = token
        req = Request(base + path, data=json.dumps(body).encode() if body is not None else None, headers=h, method=method)
        try:
            with urlopen(req, timeout=20) as resp:
                return resp.status, json.loads(resp.read() or b'{}')
        except HTTPError as e:
            return e.code, json.loads(e.read() or b'{}')

    def sends():
        if not os.path.exists(fake):
            return []
        with open(fake) as f:
            return [json.loads(x) for x in f if x.strip()]

    def wait_sends(n, t=GAP + 6):
        end = time.time() + t
        while time.time() < end:
            if len(sends()) >= n:
                return sends()
            time.sleep(0.25)
        return sends()

    try:
        for _ in range(50):
            try:
                if http('GET', '')[0] == 200:
                    break
            except OSError:
                time.sleep(0.1)
        # --- auth ---
        a1 = http('POST', '/push/subscribe', {'subscription': Device('https://fcm.googleapis.com/fcm/send/x').sub()}, token=None)
        a2 = http('POST', '/push/subscribe', {'subscription': Device('https://fcm.googleapis.com/fcm/send/x').sub()}, token='wrong')
        a3 = http('GET', '/push/vapid', token=None)
        a4 = http('POST', '/push/test', {}, token='nope')
        row('no / wrong token refused (subscribe, vapid, test)', [x[0] for x in (a1, a2, a3, a4)] == [401, 401, 401, 401])
        row('a refused subscribe stored nothing', not json.load(open(os.path.join(d, 'push.json')))['subs']
            if os.path.exists(os.path.join(d, 'push.json')) else True)
        st, v = http('GET', '/push/vapid')
        row('GET /push/vapid -> the public key', st == 200 and v.get('enabled') is True and v.get('publicKey') == pub)
        # --- the old endpoints unchanged ---
        st, g = http('GET', '')
        st2, m = http('POST', '', {'text': 'hello from an old page'})
        st3, g2 = http('GET', '?rev=%d' % (g['rev'] + 1))
        row('old endpoints unchanged (GET, POST {text}, ?rev)', st == 200 and 'messages' in g and st2 == 200 and m.get('topic') == 'general'
            and m.get('to') == 'kills' and g2.get('unchanged') is True)
        # --- subscribe ---
        phone = Device('https://web.push.apple.com/QK-phone')
        desk = Device('https://fcm.googleapis.com/fcm/send/desk')
        s1 = http('POST', '/push/subscribe', {'subscription': phone.sub()})
        s2 = http('POST', '/push/subscribe', {'subscription': desk.sub()})
        s3 = http('POST', '/push/subscribe', {'subscription': phone.sub()})
        row('subscribe two devices (again = no duplicate)', s1[0] == 200 and s2[1].get('devices') == 2 and s3[1].get('devices') == 2)
        bad1 = http('POST', '/push/subscribe', {'subscription': dict(Device('https://evil.example/fcm.googleapis.com').sub())})
        bad2 = http('POST', '/push/subscribe', {'subscription': {'endpoint': 'https://fcm.googleapis.com/fcm/send/z', 'keys': {'p256dh': 'AAAA', 'auth': 'BBBB'}}})
        bad3 = http('POST', '/push/subscribe', {'subscription': dict(Device('http://fcm.googleapis.com/fcm/send/z').sub())})
        row('bad endpoint host / http / bad keys refused (400)', [bad1[0], bad2[0], bad3[0]] == [400, 400, 400])
        row('push.json is 0600', oct(os.stat(os.path.join(d, 'push.json')).st_mode & 0o777) == '0o600')
        # --- status never pushes ---
        n0 = len(sends())
        r = sh('status', 'building the shop page', '--to', 'live')
        r2 = sh('status', 'kills status', '--to', 'kills')
        time.sleep(3)
        q = json.load(open(os.path.join(d, 'push.json')))['queue']
        row('status lines never push', r.returncode == 0 and r2.returncode == 0 and len(sends()) == n0 and not q)
        # --- a post pushes, decrypts, right title/body/url ---
        r = sh('post', '--to', 'live', '--topic', 'Shop prices', 'The coin packs are live on the test build: 100 / 550 / 1200. ' * 3)
        out = wait_sends(n0 + 2)
        new = out[n0:]
        got = {x['url']: x for x in new}
        p_phone = phone.decrypt(bytes.fromhex(got[phone.endpoint]['body'])) if phone.endpoint in got else {}
        p_desk = desk.decrypt(bytes.fromhex(got[desk.endpoint]['body'])) if desk.endpoint in got else {}
        row('a post pushes to both devices (decrypts)', r.returncode == 0 and len(new) == 2 and p_phone == p_desk and bool(p_phone))
        row('title LANE + topic, body <= 120 chars, url = the topic',
            p_phone.get('title') == 'LIVE · Shop prices' and 0 < len(p_phone.get('body', '')) <= 120
            and p_phone.get('body', '').startswith('The coin packs') and p_phone.get('url') == 'inbox.html?topic=shop-prices',
            repr(p_phone.get('title')))
        hd = got.get(phone.endpoint, {}).get('headers', {})
        row('VAPID + aes128gcm headers', hd.get('Content-Encoding') == 'aes128gcm' and hd.get('Authorization', '').startswith('vapid t=')
            and ('k=' + pub) in hd.get('Authorization', '') and hd.get('TTL') == '86400')
        # --- coalescing: 3 more LIVE inside the window -> one "3 new from LIVE"; KILLS is its own lane ---
        t_live = time.time()
        n1 = len(sends())
        for i, tp in enumerate(['Shop prices', 'Saves', 'Saves']):
            sh('post', '--to', 'live', '--topic', tp, 'update %d' % i)
        r = sh('ask', '--to', 'kills', '--topic', 'Kill films', 'Which ending for THE X?', '--choices', 'A|B')
        out = wait_sends(n1 + 2, t=3)
        k_new = [phone.decrypt(bytes.fromhex(x['body'])) for x in out[n1:] if x['url'] == phone.endpoint]
        row('KILLS pushes at once while LIVE waits (lanes apart)', len(k_new) == 1 and k_new[0]['title'] == 'KILLS asks · Kill films'
            and k_new[0]['url'] == 'inbox.html?topic=kill-films&qs=1' and time.time() - t_live < GAP,
            k_new[0]['title'] if k_new else 'none')
        out = wait_sends(n1 + 4, t=GAP + 5)
        l_new = [phone.decrypt(bytes.fromhex(x['body'])) for x in out[n1:] if x['url'] == phone.endpoint and x is not None]
        l_new = [x for x in l_new if x['lane'] == 'live']
        row('3 LIVE posts in one window -> ONE push "3 new from LIVE"', len(l_new) == 1 and l_new[0]['n'] == 3
            and l_new[0]['body'].startswith('3 new from LIVE') and l_new[0]['url'] == 'inbox.html?topic=all',
            l_new[0]['body'] if l_new else 'none')
        sent_t = json.load(open(os.path.join(d, 'push.json')))['sent']
        row('no LIVE push before the gap', bool(l_new) and sent_t.get('live', 0) - t_live >= GAP - 2.5, '%.1fs' % (sent_t.get('live', 0) - t_live))
        # --- reply pushes; the token is never in a push ---
        time.sleep(GAP)
        n2 = len(sends())
        mid = m['id']
        r = sh('reply', str(mid), 'Got it. The key is %s and the sha is 4f9a1c2be8d7a6f5e4d3c2b1a0f9e8d7c6b5a4f3' % tok, '--to', 'kills')
        out = wait_sends(n2 + 2)
        rp = [phone.decrypt(bytes.fromhex(x['body'])) for x in out[n2:] if x['url'] == phone.endpoint]
        row('a reply pushes (KILLS · General)', r.returncode == 0 and len(rp) == 1 and rp[0]['title'] == 'KILLS · General')
        row('no token / key-like strings in a push', bool(rp) and tok not in json.dumps(rp) and '4f9a1c2be8d7' not in rp[0]['body'],
            rp[0]['body'] if rp else '')
        row('badge = open questions', bool(rp) and rp[0].get('badge') == 1)
        # --- pull / answers do not push ---
        n3 = len(sends())
        sh('pull', '--to', 'kills')
        sh('answers', '--to', 'live')
        http('POST', '', {'text': 'user message', 'to': 'both'})
        time.sleep(3)
        row('user messages, pull, answers never push', len(sends()) == n3)
        # --- prune on 410 / 404, keep a network error ---
        gone = Device('https://fcm.googleapis.com/fcm/send/old-gone')
        miss = Device('https://web.push.apple.com/old-missing')
        down = Device('https://fcm.googleapis.com/fcm/send/flaky-down')
        for dv in (gone, miss, down):
            http('POST', '/push/subscribe', {'subscription': dv.sub()})
        time.sleep(GAP)
        n4 = len(sends())
        sh('post', '--to', 'kills', '--topic', 'General', 'prune check')
        wait_sends(n4 + 5)
        time.sleep(0.6)
        subs = json.load(open(os.path.join(d, 'push.json')))['subs']
        eps = {s['endpoint']: s for s in subs}
        row('404 / 410 subscriptions pruned, the others kept', gone.endpoint not in eps and miss.endpoint not in eps
            and phone.endpoint in eps and desk.endpoint in eps and down.endpoint in eps, '%d left' % len(subs))
        row('a network error is kept and counted', eps.get(down.endpoint, {}).get('fails') == 1)
        # --- test + unsubscribe ---
        n5 = len(sends())
        st, t1 = http('POST', '/push/test', {'endpoint': desk.endpoint})
        out = wait_sends(n5 + 1, t=2)
        tp = desk.decrypt(bytes.fromhex(out[-1]['body'])) if len(out) > n5 else {}
        row('SEND TEST: one push to that device only', st == 200 and t1.get('sent') == 1 and len(out) == n5 + 1
            and tp.get('lane') == 'test' and out[-1]['url'] == desk.endpoint)
        st, u = http('POST', '/push/unsubscribe', {'endpoint': desk.endpoint})
        st2, t2 = http('POST', '/push/test', {'endpoint': desk.endpoint})
        row('unsubscribe removes it; a test to it is refused', st == 200 and u.get('removed') == 1 and st2 == 400)
        r = sh('push-status')
        ps = json.loads(r.stdout or '{}')
        row('inbox.sh push-status (no keys / endpoints printed)', r.returncode == 0 and ps.get('enabled') is True and ps.get('devices') == 2
            and 'p256dh' not in r.stdout and 'QK-phone' not in r.stdout and 'PRIVATE' not in r.stdout)
        # --- push off (no key) leaves the inbox working ---
        os.rename(os.path.join(d, 'vapid_private.pem'), os.path.join(d, 'vapid_off.pem'))
        st, v = http('GET', '/push/vapid')
        r = sh('post', '--to', 'live', '--topic', 'General', 'with push off')
        st2, g = http('GET', '')
        row('no VAPID key: push off, inbox still works', st == 200 and v.get('enabled') is False and r.returncode == 0
            and g['messages'][-1]['text'] == 'with push off')
        os.rename(os.path.join(d, 'vapid_off.pem'), os.path.join(d, 'vapid_private.pem'))
    finally:
        srv.terminate()
        srv.wait(5)
        if not keep:
            shutil.rmtree(d, ignore_errors=True)
        else:
            print('kept', d)
    bad = [n for n, ok in ROWS if not ok]
    print('%d/%d passed' % (len(ROWS) - len(bad), len(ROWS)))
    sys.exit(1 if bad else 0)


if __name__ == '__main__':
    main()
