#!/usr/bin/env python3
"""Driver: starts the HTTP server, launches the ASan chrome binary,
injects a user-activated call to window.startPOC() over CDP, and tees
chrome stderr (where ASan reports go) to stdout."""
import json
import os
import shutil
import signal
import subprocess
import sys
import threading
import time
import urllib.request

import websocket  # pip install websocket-client

HERE   = os.path.dirname(os.path.abspath(__file__))
SRC    = os.environ.get('CHROMIUM_SRC', '/chromium/chromium/src')
CHROME = os.environ.get('CHROME', os.path.join(SRC, 'out', 'asan', 'chrome'))
PORT   = int(os.environ.get('PORT', '28765'))
CDP    = int(os.environ.get('CDP_PORT', '29222'))
PROFILE = '/tmp/chrome-poc-profile'


def start_server():
    return subprocess.Popen(
        [sys.executable, os.path.join(HERE, 'server.py')],
        env={**os.environ, 'PORT': str(PORT), 'CHROMIUM_SRC': SRC},
        stdout=sys.stderr, stderr=sys.stderr)


def tee(stream, sink, prefix=b''):
    for line in iter(stream.readline, b''):
        sink.buffer.write(prefix + line)
        sink.flush()


def cdp_connect():
    for _ in range(200):
        try:
            data = urllib.request.urlopen(
                f'http://127.0.0.1:{CDP}/json/list', timeout=1).read()
            for t in json.loads(data):
                if t.get('type') == 'page' and 'poc.html' in t.get('url', ''):
                    return websocket.create_connection(
                        t['webSocketDebuggerUrl'], timeout=60)
        except Exception:
            pass
        time.sleep(0.2)
    raise RuntimeError('CDP not reachable')


def cdp(ws, _id, method, params=None):
    ws.send(json.dumps({'id': _id, 'method': method,
                        'params': params or {}}))
    while True:
        msg = json.loads(ws.recv())
        if msg.get('id') == _id:
            return msg


def main():
    if os.path.exists(PROFILE):
        shutil.rmtree(PROFILE, ignore_errors=True)

    srv = start_server()
    time.sleep(0.5)

    flags = [
        CHROME,
        '--no-first-run',
        '--no-default-browser-check',
        '--use-gl=angle',
        '--use-angle=swiftshader',
        '--ignore-gpu-blocklist',
        '--enable-unsafe-swiftshader',
        '--in-process-gpu',
        '--user-data-dir=' + PROFILE,
        '--enable-blink-features=MojoJS,MojoJSTest',
        '--js-flags=--expose-gc',
        '--no-sandbox',
        '--disable-crash-reporter',
        '--disable-breakpad',
        '--disable-in-process-stack-traces',
        '--remote-debugging-port=%d' % CDP,
        '--remote-allow-origins=*',
        '--headless=new',
        '--enable-logging=stderr',
        '--v=0',
        f'http://127.0.0.1:{PORT}/poc.html',
    ]
    sym = os.path.join(SRC, 'third_party', 'llvm-build', 'Release+Asserts',
                       'bin', 'llvm-symbolizer')
    env = {**os.environ,
           'ASAN_OPTIONS':
               'handle_segv=1:handle_sigbus=1:handle_abort=1:'
               'detect_odr_violation=0:symbolize=1:'
               'external_symbolizer_path=' + sym,
           'ASAN_SYMBOLIZER_PATH': sym}
    print('launching:', ' '.join(flags), file=sys.stderr)
    chrome = subprocess.Popen(
        flags, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    threading.Thread(target=tee, args=(chrome.stdout, sys.stdout),
                     daemon=True).start()

    try:
        ws = cdp_connect()
        cdp(ws, 1, 'Runtime.enable')

        # Wait for the module to load.
        for _ in range(300):
            r = cdp(ws, 4, 'Runtime.evaluate',
                    {'expression': 'typeof window.startPOC',
                     'returnByValue': True})
            if r.get('result', {}).get('result', {}).get('value') == 'function':
                break
            time.sleep(0.1)
        else:
            print('startPOC never defined', file=sys.stderr)

        print('>>> injecting user-activated startPOC()', file=sys.stderr)
        cdp(ws, 5, 'Runtime.evaluate',
            {'expression': 'window.startPOC()',
             'userGesture': True, 'awaitPromise': False})

        # Wait for completion or crash.
        for _ in range(600):
            if chrome.poll() is not None:
                break
            try:
                r = cdp(ws, 6, 'Runtime.evaluate',
                        {'expression': 'document.title',
                         'returnByValue': True})
                if r.get('result', {}).get('result', {}).get('value') == 'DONE':
                    time.sleep(0.5)
                    break
            except Exception:
                break
            time.sleep(0.2)
    finally:
        try:
            chrome.send_signal(signal.SIGTERM)
            chrome.wait(timeout=10)
        except Exception:
            chrome.kill()
        srv.terminate()


if __name__ == '__main__':
    main()