#!/usr/bin/python3 -I
"""Usage: instant-hotspot [--auto] [--creds DIR] [--print] | selftest

Asks a same-account iPhone to start Personal Hotspot over Rapport (BLE L2CAP
PSM 0x81), then joins it with NetworkManager. --auto repeats that whenever the
machine has no network. Credentials: rapport_selfidentity.bin and
rapport_selfidentity_gena.bin of an Apple device on the same account.
"""
import argparse
import asyncio
import os
import socket
import struct
import subprocess
import sys
import tempfile
import time

PSM = 0x81
PV_START, PV_NEXT, E_OPACK = 5, 6, 8
APPLE = 0x004C


# OPACK (Apple binary plist-like encoding), the subset Rapport uses here
def opack_pack(v):
    if v is None:
        return b'\x04'
    if isinstance(v, bool):
        return b'\x01' if v else b'\x02'
    if isinstance(v, int):
        if 0 <= v < 0x28:
            return bytes([v + 8])
        for tag, n in ((0x30, 1), (0x31, 2), (0x32, 4), (0x33, 8)):
            if v < 1 << (8 * n):
                return bytes([tag]) + v.to_bytes(n, 'little')
    if isinstance(v, float):
        return b'\x36' + struct.pack('<d', v)
    if isinstance(v, str):
        return _sized(v.encode(), 0x40, 0x61)
    if isinstance(v, bytes):
        return _sized(v, 0x70, 0x91)
    if isinstance(v, list):
        body = b''.join(map(opack_pack, v))
        return bytes([0xD0 + min(len(v), 15)]) + body + (b'\x03' if len(v) >= 15 else b'')
    if isinstance(v, dict):
        body = b''.join(opack_pack(k) + opack_pack(x) for k, x in v.items())
        return bytes([0xE0 + min(len(v), 15)]) + body + (b'\x03' if len(v) >= 15 else b'')
    raise TypeError(type(v))


def _sized(b, short, long):
    if len(b) <= 0x20:
        return bytes([short + len(b)]) + b
    for i, n in enumerate((1, 2, 4)):
        if len(b) < 1 << (8 * n):
            return bytes([long + i]) + len(b).to_bytes(n, 'little') + b
    raise ValueError('too long')


def opack_unpack(b):
    objs = []

    def one(i):
        t = b[i]
        if t in (1, 2, 4):
            return {1: True, 2: False, 4: None}[t], i + 1
        if 0x08 <= t <= 0x2F:
            return t - 8, i + 1
        if 0xA0 <= t <= 0xC0:
            return objs[t - 0xA0], i + 1
        if 0xC1 <= t <= 0xC4:
            n = 1 << (t - 0xC1)
            return objs[int.from_bytes(b[i + 1:i + 1 + n], 'little')], i + 1 + n
        if t == 0x05:
            v, j = b[i + 1:i + 17], i + 17
        elif t == 0x06:
            v, j = struct.unpack('<d', b[i + 1:i + 9])[0], i + 9
        elif t == 0x35:
            v, j = struct.unpack('<f', b[i + 1:i + 5])[0], i + 5
        elif t == 0x36:
            v, j = struct.unpack('<d', b[i + 1:i + 9])[0], i + 9
        elif 0x30 <= t <= 0x33:
            n = 1 << (t - 0x30)
            v, j = int.from_bytes(b[i + 1:i + 1 + n], 'little'), i + 1 + n
        elif 0x40 <= t <= 0x60:
            v, j = b[i + 1:i + 1 + t - 0x40].decode(), i + 1 + t - 0x40
        elif 0x61 <= t <= 0x64:
            n = (1, 2, 3, 4)[t - 0x61]
            ln = int.from_bytes(b[i + 1:i + 1 + n], 'little')
            v, j = b[i + 1 + n:i + 1 + n + ln].decode(), i + 1 + n + ln
        elif 0x70 <= t <= 0x90:
            v, j = b[i + 1:i + 1 + t - 0x70], i + 1 + t - 0x70
        elif 0x91 <= t <= 0x94:
            n = 1 << (t - 0x91)
            ln = int.from_bytes(b[i + 1:i + 1 + n], 'little')
            v, j = b[i + 1 + n:i + 1 + n + ln], i + 1 + n + ln
        elif 0xD0 <= t <= 0xDF or 0xE0 <= t <= 0xEF:
            is_dict = t >= 0xE0
            count, j = t & 0xF, i + 1
            items = []
            while count == 15 and b[j] != 3 or count < 15 and len(items) < count * (2 if is_dict else 1):
                x, j = one(j)
                items.append(x)
            if count == 15:
                j += 1
            v = dict(zip(items[::2], items[1::2])) if is_dict else items
            objs.append(v)
            return v, j
        else:
            raise ValueError(f'opack tag {t:#x}')
        objs.append(v)
        return v, j

    return one(0)[0]


def tlv_pack(d):
    out = b''
    for k, v in d.items():
        for i in range(0, max(len(v), 1), 255):
            out += bytes([k, len(v[i:i + 255])]) + v[i:i + 255]
    return out


def tlv_unpack(b):
    d, i, last = {}, 0, None
    while i + 2 <= len(b):
        k, n = b[i], b[i + 1]
        v = b[i + 2:i + 2 + n]
        d[k] = d[k] + v if k == last and k in d else v
        last, i = k, i + 2 + n
    return d


class Rapport:
    def __init__(self, addr):
        self.s = socket.socket(socket.AF_BLUETOOTH, socket.SOCK_SEQPACKET, socket.BTPROTO_L2CAP)
        # LE source type, else the kernel sends basic-mode frames without the LE CoC SDU header
        self.s.bind((socket.BDADDR_ANY, 0, 0, socket.BDADDR_LE_PUBLIC))
        self.s.settimeout(15)
        self.s.connect((addr, PSM, 0, socket.BDADDR_LE_RANDOM))
        self.buf = b''
        self.tx = self.rx = None
        self.ntx = self.nrx = 0

    def send(self, ftype, obj):
        payload = opack_pack(obj)
        hdr = bytes([ftype]) + (len(payload) + (16 if self.tx else 0)).to_bytes(3, 'big')
        if self.tx:
            payload = self.tx.encrypt(self.ntx.to_bytes(12, 'little'), payload, hdr)
            self.ntx += 1
        self.s.sendall(hdr + payload)

    def _read(self, n):
        while len(self.buf) < n:
            c = self.s.recv(65536)
            if not c:
                raise ConnectionError('phone closed the channel')
            self.buf += c
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def recv(self):
        hdr = self._read(4)
        body = self._read(int.from_bytes(hdr[1:], 'big'))
        if self.rx and body:
            body = self.rx.decrypt(self.nrx.to_bytes(12, 'little'), body, hdr)
            self.nrx += 1
        return opack_unpack(body) if body else {}

    def pair_verify(self, edsk):
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
        from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
        from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
        from cryptography.hazmat.primitives.kdf.hkdf import HKDF

        def hkdf(salt, info, ikm):
            return HKDF(hashes.SHA512(), 32, salt, info).derive(ikm)

        eph = X25519PrivateKey.generate()
        pub = eph.public_key().public_bytes_raw()
        self.send(PV_START, {'_pd': tlv_pack({6: b'\x01', 3: pub, 0x19: b'\x01'})})
        m2 = tlv_unpack(self.recv()['_pd'])
        if 7 in m2:
            raise PermissionError(f'pair-verify M2 error {m2[7].hex()}')
        spub = m2[3]
        shared = eph.exchange(X25519PublicKey.from_public_bytes(spub))
        k = ChaCha20Poly1305(hkdf(b'Pair-Verify-Encrypt-Salt', b'Pair-Verify-Encrypt-Info', shared))
        k.decrypt(b'\0\0\0\0PV-Msg02', m2[5], None)
        # same-account mode: identity is resolved from the signature, no identifier TLV
        sig = Ed25519PrivateKey.from_private_bytes(edsk).sign(pub + spub)
        enc = k.encrypt(b'\0\0\0\0PV-Msg03', tlv_pack({10: sig}), None)
        self.send(PV_NEXT, {'_pd': tlv_pack({6: b'\x03', 5: enc})})
        m4 = tlv_unpack(self.recv().get('_pd', b''))
        if 7 in m4:
            raise PermissionError(f'phone rejected this identity (M4 error {m4[7].hex()})')
        self.tx = ChaCha20Poly1305(hkdf(b'', b'ClientEncrypt-main', shared))
        self.rx = ChaCha20Poly1305(hkdf(b'', b'ServerEncrypt-main', shared))

    def request(self, rid, content):
        self.send(E_OPACK, {'_i': rid, '_x': 1, '_t': 2, '_c': content})
        while True:
            r = self.recv()
            if r.get('_t') == 3 and r.get('_x') == 1:
                if '_ec' in r or '_em' in r:
                    raise RuntimeError(f"phone error: {r.get('_em') or r.get('_ec')}")
                return r.get('_c', {})


def load_creds(d):
    b = open(os.path.join(d, 'rapport_selfidentity.bin'), 'rb').read()
    ident = opack_unpack(b)
    gena = {}
    try:
        gena = opack_unpack(open(os.path.join(d, 'rapport_selfidentity_gena.bin'), 'rb').read())
    except OSError:
        pass
    return ident['edSK'], ident.get('dIRK'), gena


def siphash24(key, msg):
    m64 = (1 << 64) - 1

    def rotl(x, b):
        return ((x << b) | (x >> (64 - b))) & m64

    k0, k1 = int.from_bytes(key[:8], 'little'), int.from_bytes(key[8:], 'little')
    v = [k0 ^ 0x736f6d6570736575, k1 ^ 0x646f72616e646f6d, k0 ^ 0x6c7967656e657261, k1 ^ 0x7465646279746573]

    def rnd():
        v[0] = (v[0] + v[1]) & m64; v[1] = rotl(v[1], 13) ^ v[0]; v[0] = rotl(v[0], 32)
        v[2] = (v[2] + v[3]) & m64; v[3] = rotl(v[3], 16) ^ v[2]
        v[0] = (v[0] + v[3]) & m64; v[3] = rotl(v[3], 21) ^ v[0]
        v[2] = (v[2] + v[1]) & m64; v[1] = rotl(v[1], 17) ^ v[2]; v[2] = rotl(v[2], 32)

    n = len(msg) // 8 * 8
    blocks = [int.from_bytes(msg[i:i + 8], 'little') for i in range(0, n, 8)]
    blocks.append((len(msg) << 56) | int.from_bytes(msg[n:].ljust(8, b'\0'), 'little') & ((1 << 56) - 1))
    for b in blocks:
        v[3] ^= b; rnd(); rnd(); v[0] ^= b
    v[2] ^= 0xff
    for _ in range(4):
        rnd()
    return v[0] ^ v[1] ^ v[2] ^ v[3]


def nearby_info(adv):
    md = adv.manufacturer_data.get(APPLE, b'')
    i = 0
    while i + 2 <= len(md):
        t, n = md[i], md[i + 1]
        v = md[i + 2:i + 2 + n]
        if t == 0x10 and len(v) >= 2:
            return v
        i += 2 + n
    return None


def own_tag(dirk, addr):
    h = siphash24(dirk, bytes.fromhex(addr.replace(':', '')))
    return (h & ((1 << 48) - 1)).to_bytes(6, 'big')[:3]


def drop_classic_phone_links():
    # an open BR/EDR link to the phone makes it refuse the LE connection
    out = subprocess.run(['bluetoothctl', 'devices', 'Connected'], capture_output=True, text=True).stdout
    for line in out.splitlines():
        addr = line.split()[1] if line.startswith('Device ') else None
        if addr and 'Icon: phone' in subprocess.run(['bluetoothctl', 'info', addr], capture_output=True, text=True).stdout:
            subprocess.run(['bluetoothctl', 'disconnect', addr], capture_output=True)


async def fetch_hotspot(creds):
    from bleak import BleakScanner

    edsk, dirk, gena = load_creds(creds)
    drop_classic_phone_links()
    seen = {}

    def talk(addr):
        r = Rapport(addr)
        r.pair_verify(edsk)
        return r.request('com.sharing.hotspot.request', {
            'reqT': 1, 'sID': int.from_bytes(os.urandom(4), 'little'),
            'altDSID': gena.get('accountAltDSID', ''), 'aplID': gena.get('accountID', ''),
            'devnm': socket.gethostname(), 'icld': 1, 'proVer': '15.2', '5ghz': True})

    def seen_cb(d, a):
        v = nearby_info(a)
        if not v or a.rssi < -85:
            return
        if dirk and v[1] & 0x10 and len(v) >= 5 and v[2:5] == own_tag(dirk, d.address):
            return  # the Apple device whose identity we use
        seen[d.address] = (bool(v[1] & 0x02), a.rssi, d)

    # connect while discovery runs: BlueZ then sends a direct LE connect instead of waiting on the accept list
    async with BleakScanner(seen_cb):
        await asyncio.sleep(4)
        if not seen:
            await asyncio.sleep(6)
        if not seen:
            raise LookupError('no Apple phone nearby')
        last = None
        for hotspot, rssi, dev in sorted(seen.values(), key=lambda x: (x[0], x[1]), reverse=True)[:3]:
            try:
                return await connect_and_talk(dev, talk)
            except Exception as e:  # dbus, socket, auth: try the next candidate
                last = e
        raise ConnectionError(f'no phone answered: {last}')


async def connect_and_talk(dev, talk):
    from dbus_next import BusType
    from dbus_next.aio import MessageBus
    # plain Device1.Connect: the phone drops centrals that run GATT discovery (BleakClient does)
    bus = await MessageBus(bus_type=BusType.SYSTEM).connect()
    path = dev.details['path']
    intro = await bus.introspect('org.bluez', path)
    d1 = bus.get_proxy_object('org.bluez', path, intro).get_interface('org.bluez.Device1')
    await asyncio.wait_for(d1.call_connect(), 20)
    try:
        return await asyncio.get_running_loop().run_in_executor(None, talk, dev.address)
    finally:
        try:
            await d1.call_disconnect()
        except Exception:
            pass
        bus.disconnect()


def nm(*args, check=True):
    return subprocess.run(['nmcli', *args], capture_output=True, text=True, check=check).stdout


def join(ssid, password):
    # brcmfmac cannot associate to a new AP while AWDL runs; it is fine to restart once joined
    awdl = subprocess.run(['systemctl', 'is-active', '--quiet', 'awdld'], check=False).returncode == 0
    if awdl:
        subprocess.run(['systemctl', 'stop', 'awdld'], check=False)
    try:
        return _join(ssid, password)
    finally:
        if awdl:
            subprocess.run(['systemctl', 'start', 'awdld'], check=False)


def _join(ssid, password):
    profiles = [l.rsplit(':', 1) for l in nm('-t', '-f', 'NAME,TYPE', 'connection', 'show').splitlines()]
    names = [n for n, t in profiles if t == '802-11-wireless' and nm('-g', '802-11-wireless.ssid', 'connection', 'show', n, check=False).strip() == ssid]
    if names:
        name = names[0]
    else:
        name = ssid
        nm('connection', 'add', 'type', 'wifi', 'con-name', name, 'ssid', ssid,
           'wifi-sec.key-mgmt', 'wpa-psk', 'wifi-sec.psk-flags', '2', 'connection.metered', 'yes')
    fd, path = tempfile.mkstemp(dir=os.environ.get('XDG_RUNTIME_DIR'))
    try:
        os.write(fd, f'802-11-wireless-security.psk:{password}\n'.encode())
        os.close(fd)
        for _ in range(8):
            nm('device', 'wifi', 'rescan', check=False)
            time.sleep(3)
            if subprocess.run(['nmcli', 'connection', 'up', name, 'passwd-file', path], capture_output=True).returncode == 0:
                return name
        raise ConnectionError(f'could not join {ssid}')
    finally:
        os.unlink(path)


def online():
    return nm('networking', 'connectivity', 'check', check=False).strip() in ('full', 'limited', 'portal')


def run(creds, show):
    info = asyncio.run(fetch_hotspot(creds))
    print(f"hotspot {info['name']} channel {info.get('chan')}" + (f" password {info['pass']}" if show else ''), flush=True)
    print('joined', join(info['name'], info['pass']), flush=True)


def selftest():
    sample = {'_i': 'com.sharing.hotspot.request', '_x': 1, '_t': 2, 'n': 300, 'b': b'\x01' * 40,
              'l': [1, 'x', None, True], 's': 'y' * 40}
    assert opack_unpack(opack_pack(sample)) == sample
    assert opack_pack({'_pd': b''}) == b'\xe1\x43_pd\x70'
    assert siphash24(bytes(range(16)), bytes(range(15))) == 0xa129ca6149be45e5
    assert tlv_unpack(tlv_pack({5: b'a' * 300, 6: b'\x01'})) == {5: b'a' * 300, 6: b'\x01'}
    assert opack_unpack(bytes.fromhex('e2 41 61 41 62 41 63 a1')) == {'a': 'b', 'c': 'b'}
    print('ok')


def main():
    if sys.argv[1:] == ['selftest']:
        return selftest()
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument('--creds', default=os.path.join(os.environ.get('XDG_CONFIG_HOME') or os.path.expanduser('~/.config'), 'instant-hotspot'))
    ap.add_argument('--auto', action='store_true', help='join the hotspot whenever offline')
    ap.add_argument('--print', dest='show', action='store_true', help='also print the password')
    a = ap.parse_args()
    if not a.auto:
        try:
            run(a.creds, a.show)
        except (LookupError, ConnectionError, PermissionError, RuntimeError, TimeoutError, OSError) as e:
            sys.exit(f'instant-hotspot: {e}')
        return
    while True:
        if not online():
            try:
                run(a.creds, False)
            except Exception as e:
                print(f'instant-hotspot: {e}', flush=True)
        time.sleep(20)


if __name__ == '__main__':
    main()
