#!/usr/bin/env python3
from urllib.parse import quote, quote_plus, unquote_plus, parse_qsl, urlencode
from logging import getLogger, DEBUG, debug, info, error
from tempfile import TemporaryDirectory
from collections import OrderedDict
from re import search, sub, findall
from urllib.request import urlopen
from traceback import print_exc
from json import loads, dumps
from subprocess import Popen
from os import access, X_OK
from random import randint
from shutil import which
from shlex import join
from time import sleep

from os.path import dirname, realpath

__import__('sys').path.append(dirname(realpath(__file__)) + '/../..')
from pbtk.utils.transports import GMapsAPIPublic, GMapsAPIPrivate
from pbtk.utils.common import register_extractor, extractor_main

"""
    Protobuf-URL (named internally JsProtoUrl) is the way I call the
    text encoding you can observe in Google Maps-related URLs (both the
    frontend web application, private API requests and public API
    requests).

    This script opens a Chrome instance controlled through the remote
    debugging API [1] to the URL of your choice, searches for and hooks
    JavaScript functions related to Protobuf-URL, and when an URL query
    string is generated through this function and AJAX-called (or browsed
    to through the history API), it generates both reconstructed .protos
    and HTTP endpoint information.
    
    For Chrome debugging protocol's calls documentation, see [2]. As of
    January 2017, using a developement version is advised.
    
    [1] https://developer.chrome.com/devtools/docs/debugger-protocol
    [2] https://chromedevtools.github.io/debugger-protocol-viewer/tot/
"""

browser = (
    which('C:/Program Files (x86)/Google/Chrome/Application/chrome.exe')
    or which('C:/Program Files/Google/Chrome/Application/chrome.exe')
    or which('chromium-browser')
    or which('chromium')
    or which('chrome')
    or which('/Applications/Google Chrome.app/Contents/MacOS/Google Chrome')
    or which('/Applications/Chromium.app/Contents/MacOS/Chromium')
    or which('Google Chrome')
    or which('Chromium')
    or 'google-chrome'
)


@register_extractor(
    name='pburl_extract',
    desc=(
        'Extract and capture Protobuf-URL endpoints from a Chrome instance (http://*)'
        + ' - needs update to work as of 2026'
    ),
    pick_url=True,
    depends={'modules': ['websocket']},
)
def pburl_extract(url):
    global \
        URL, \
        req_id, \
        req_data, \
        sent_msgs, \
        awaiting_srcs, \
        endpoints, \
        sid_to_vars, \
        proto_to_urls
    URL = url

    req_id = 1
    req_data = {}
    sent_msgs = {}
    awaiting_srcs = []
    endpoints = []
    sid_to_vars = {}
    proto_to_urls = OrderedDict()

    port = randint(1024, 32767)
    temp_profile = True

    with TemporaryDirectory() as profile:
        cmd = [browser, '--remote-debugging-port=%d' % port, 'about:blank']
        if temp_profile:
            cmd += [
                '--user-data-dir=' + profile,
                '--remote-allow-origins=*',
                '--no-first-run',
                '--start-maximized',
                '--no-default-browser-check',
            ]

        yield (
            '_progress',
            (
                'Trying to launch browser...\n(requires Chrome or Chromium)',
                None,
            ),
        )

        if which(browser) and access(realpath(which(browser)), X_OK):
            debug('Launching browser: ' + join(cmd))

            chrome = Popen(cmd)  # , stdout=DEVNULL, stderr=DEVNULL
        else:
            chrome = None
            cmd[0] = '$(which google-chrome || which chromium-browser)'
            yield (
                '_info',
                (
                    "Can't launch Chrome or a Chromium-based "
                    + 'browser, please launch yourself: "'
                    + join(cmd)
                    + '", interact with it, and close it'
                ),
            )

        try:
            while True:
                try:
                    tabs = (
                        urlopen('http://localhost:%d/json' % port)
                        .read()
                        .decode('utf8')
                    )
                    info('Received CDP JSON API response: %r' % tabs)
                    tab = next(
                        tab for tab in loads(tabs) if tab['type'] == 'page'
                    )
                    break
                except OSError:
                    sleep(0.1)

            yield (
                '_progress',
                (
                    'Connecting to browser debugging controller...\n(Your activity '
                    + 'from the first tab will be captured, until you close it)',
                    None,
                ),
            )

            from websocket import WebSocketApp

            getLogger('websocket').setLevel(DEBUG)
            WebSocketApp(
                tab['webSocketDebuggerUrl'],
                on_message=on_message,
                on_open=on_open,
                on_close=on_close,
            ).run_forever()

            debug('[i] CDP connection succeeded')

        finally:
            # Make sure that Chromium is killed when quitting by error
            if chrome:
                chrome.terminate()

            # Avoid race condition with it writing to profile directory
            sleep(0.2)

    for proto in list(proto_to_urls):
        # Determine .proto name based on endpoint URL
        pbname = min(proto_to_urls[proto], key=len)
        pbname = sub(r'/contrib/\{contrib\}/[a-z]+/@', '/@', pbname)
        pbname = sub(r'/([^/=]+[/=]|@)\{.+?\}(/\{.+?\})?', '', pbname)
        pbname = pbname.split('/')[3:]
        if pbname[0] in ('maps', 'rt') and len(pbname) > 1:
            pbname.pop(0)
        if pbname[0] == 'preview':
            pbname.pop(0)
        pbname = ''.join(x[0].upper() + x[1:] for x in pbname)
        pbname = sub('[^A-Za-z0-9]', '', pbname)
        pbname = {'Vt': 'VectorTown', 'S': 'Search'}.get(pbname, pbname)
        pbname_, tries = pbname, 0
        while pbname in proto_to_urls:
            tries += 1
            pbname = pbname_ + str(tries)

        # Save generated .proto to ~/.pbtk
        if len(pbname) == 1:
            yield (
                pbname + '.proto',
                proto.replace('message Top', 'message ' + pbname).replace(
                    'Top', '.' + pbname
                ),
            )
        else:
            yield pbname + '.proto', proto.replace('Top', pbname)
        proto_to_urls[pbname] = proto_to_urls[proto]
        del proto_to_urls[proto]

        info(
            '[i] Extracted Protobuf declaration: %r => %r'
            % (pbname, proto_to_urls[pbname])
        )

    # Save endpoint to ~/.pbtk
    for sample in endpoints:
        yield (
            next(k for k, v in proto_to_urls.items() if sample['url'] in v)
            + '.sample',
            sample,
        )


def on_open(ws):
    send(ws, 'Runtime.enable')
    send(ws, 'Debugger.enable')
    send(ws, 'Network.enable')
    send(ws, 'Page.navigate', {'url': URL})


def on_close(ws, *args):
    debug('[i] CDP connection closed')


def send(ws, call, params=None, data=None):
    global req_id
    req_data[req_id] = (call, data)
    if params:
        ws.send(dumps({'id': req_id, 'method': call, 'params': params}))
    else:
        ws.send(dumps({'id': req_id, 'method': call}))
    req_id += 1


def on_message(ws, msg):
    try:
        global seen_scripts
        msg = loads(msg)

        # debug('[i] Received Websocket message: %r' % msg)

        if 'method' in msg:
            call, msg = msg['method'], msg['params']

            if call == 'Network.requestWillBeSent':
                msg = msg['request']
                if '=!' in msg['url'] or search('&[0-9]', msg['url']):
                    logUrl(msg['url'])

            elif call == 'Runtime.executionContextCreated':
                seen_scripts = set()
                send(
                    ws,
                    'Runtime.evaluate',
                    {
                        'expression': """
    var wrap = function(orig) {
        return function() {
            var new_path = arguments[2];
            if(new_path[0] == '/') {
                new_path = location.origin + new_path;
            }
            console.log(JSON.stringify(['__URL', new_path]));
            return orig.apply(this, arguments);
        };
    };
    history.replaceState = wrap(history.replaceState);"""
                    },
                )

            elif call == 'Debugger.scriptParsed':
                myId = (msg['url'], msg['startLine'], msg['startColumn'])
                if myId not in seen_scripts and not (
                    msg['startLine'] == msg['endLine']
                    and msg['endColumn'] - msg['startColumn'] < 100
                ):
                    seen_scripts.add(myId)
                    awaiting_srcs.append(msg['scriptId'])
                    send(ws, 'Debugger.pause')
                    send(
                        ws,
                        'Debugger.getScriptSource',
                        {'scriptId': msg['scriptId']},
                        (
                            msg['scriptId'],
                            '\n' * msg['startLine'] + ' ' * msg['startColumn'],
                        ),
                    )

            elif call == 'Debugger.paused':
                loc = msg['callFrames'][0]['location']

                if (
                    msg['hitBreakpoints']
                    and msg['reason'] == 'other'
                    and msg['callFrames']
                ):
                    send(ws, 'Debugger.pause')
                    messageSpecStringVar, nestedMessagesSpecArrayVar = (
                        sid_to_vars[loc['scriptId']]
                    )

                    # Converting the Protobuf structure object into JSON
                    # isn't possible since it has circular references, and
                    # accessing variables individually through the API is
                    # very slow, so conversion to text is done on the
                    # JavaScript side.

                    send(
                        ws,
                        'Debugger.evaluateOnCallFrame',
                        {
                            'callFrameId': msg['callFrames'][0]['callFrameId'],
                            'expression': r"""
    (function() {
        var objToName = new WeakMap();
        var specStringToName = {};

        var typeArray = {'d': 'double', 'f': 'float', 'i': 'int32', 'j': 'int64', 'u': 'uint32', 'v': 'uint64', 'x': 'fixed32', 'y': 'fixed64', 'g': 'sfixed32', 'h': 'sfixed64', 'n': 'sint32', 'o': 'sint64', 'e': 'enum', 's': 'string', 'z': 'string (base64)', 'B': 'bytes', 'b': 'bool', 'm': 'message'};
        var namerGen = function* () {
            for(var len = 0; ; len++)
                yield* (function* namer2(len) {
                    for(var char of 'abcdefghijklmnopqrstuvwxyz')
                        for(var suffix of len ? namer2(len - 1) : [''])
                            yield char + suffix;
                })(len);
        };
        var parseMsg = function(msg, tab, name) {
            var namer = namerGen();
            var namer_msg = namerGen();

            var text = `${tab}message ${name.split('.').pop()} {\n`;
            tab += '    ';

            // console.log(JSON.stringify(['__INFO', typeof msg]));

            if(typeof msg == 'string') {
                specStringToName[msg] = name;

                var message_spec_str = msg;
                var nested_messages_spec_array = null;
            }
            else {
                objToName.set(msg, name);

                var message_spec_str = msg.%s;
                var nested_messages_spec_array = msg.%s;
            }

            var new_types_alias = {
                a: "B",
                k: "j",
                p: "o",
                w: "v",
                q: "y",
                r: "h"
            };

            var field_number = 0;

            var nested_message_index = 0;

            while(message_spec_str.length) { // Parse a string such as: "MMmemmswm11mmibbb18mbmkmImi"

                var field_number_regex = message_spec_str.match(/^(\d+)(.+)/);

                // Consume the number indicating the next field number, if present

                if(field_number_regex) {

                    var [_, field_number, message_spec_str] = field_number_regex;

                    field_number = parseInt(field_number, 10);
                }
                else {
                    field_number++;
                }

                // Consume the next letter

                var field_letter = message_spec_str[0];
                message_spec_str = message_spec_str.substr(1);

                var label = field_letter.toUpperCase() == field_letter ? 'repeated ': '';

                field_letter = field_letter.toLowerCase();

                if(new_types_alias[field_letter]) {

                    field_letter = new_types_alias[field_letter];
                }
                var type = typeArray[field_letter];

                var comment = '';
                if(type == 'enum') {
                    comment = ' // enum';
                    type = 'int32';
                }

                if(type == 'message') {

                    if(typeof nested_messages_spec_array[nested_message_index] == 'string') {

                        type = specStringToName[nested_messages_spec_array[nested_message_index]];
                    }
                    else {

                        type = objToName.get(nested_messages_spec_array[nested_message_index]);
                    }
                    if(!type) {
                        type = namer_msg.next().value;
                        type = `${name}.${type[0].toUpperCase()}${type.slice(1)}`;
                        text += parseMsg(nested_messages_spec_array[nested_message_index], tab, type);
                    }
                    if(type.split('.').slice(0, -1).join('.') == name) {
                        type = type.split('.').pop();
                    }

                    nested_message_index++;
                }

                text += `${tab}${label}${type} ${namer.next().value} = ${field_number};${comment}\n`;
            }

            return `${text}${tab.slice(0, -4)}}\n`;
        };
        try {
            console.log(JSON.stringify(['__HOOK', [parseMsg(b, '', 'Top'), c.join('')[0] === '!' ? c.join('') : c.join('&').replace(/'/g, '%%27')]]));
        }
        catch(e) {
            console.log(JSON.stringify(['__ERR', e.stack]));
        }
    })();"""
                            % (
                                messageSpecStringVar,
                                nestedMessagesSpecArrayVar,
                            ),
                        },
                    )

                if not awaiting_srcs:
                    send(ws, 'Debugger.resume')

            elif call == 'Runtime.consoleAPICalled':
                msg = msg['args'][0]['value']
                if '__URL' in msg:
                    logUrl(loads(msg)[1])
                elif '__HOOK' in msg:
                    proto, protoMsg = loads(msg)[1]

                    proto = 'syntax = "proto3";\n\n' + proto

                    if protoMsg:
                        sent_msgs[protoMsg] = proto
                        sent_msgs[protoMsg.replace(' ', '+')] = proto
                        sent_msgs[quote(protoMsg, safe='~()*!.')] = proto
                        sent_msgs[quote(protoMsg, safe="~()*!.'")] = proto
                        sent_msgs[quote(protoMsg, safe="~()*!.':")] = proto
                        sent_msgs[quote_plus(protoMsg, safe='~()*!.')] = proto
                        sent_msgs[quote_plus(protoMsg, safe="~()*!.'")] = proto
                        sent_msgs[quote_plus(protoMsg, safe="~()*!.':")] = (
                            proto
                        )

        elif 'error' in msg:
            req_id, err = msg['id'], msg['error']
            call, data = req_data.pop(req_id)
            if call in ('Debugger.resume', 'Debugger.evaluateOnCallFrame'):
                return
            if call == 'Debugger.getScriptSource':
                sid = err['message'].split(': ')[1]
                awaiting_srcs.remove(sid)
                if not awaiting_srcs:
                    send(ws, 'Debugger.resume')
                return
            error('[Error] %r: %r' % (call, err))
            exit()

        else:
            req_id, msg = msg['id'], msg['result']
            call, data = req_data.pop(req_id)

            if call == 'Debugger.getScriptSource':
                sid, padding = data
                src = padding + msg['scriptSource']

                targets = [
                    '0);c.length=a;return c.join("")}',
                    '0);return c.join("")}',
                    'return c.join("&").replace(',
                ]
                for target in targets:
                    if target in src:
                        messageSpecStringVar, nestedMessagesSpecArrayVar = (
                            search(
                                r'=a:\(this\.[\w$]+=a\.([\w$]+),this\.[\w$]+=a\.([\w$]+)\);(?:a=)?this\.',
                                src,
                            ).groups()
                        )
                        before = src.split(target)[0]
                        sid_to_vars[sid] = (
                            messageSpecStringVar,
                            nestedMessagesSpecArrayVar,
                        )

                        send(
                            ws,
                            'Debugger.setBreakpoint',
                            {
                                'location': {
                                    'scriptId': sid,
                                    'lineNumber': before.count('\n'),
                                    'columnNumber': len(
                                        before.split('\n')[-1]
                                    ),
                                }
                            },
                        )
                        info('[i] Script successfully hooked 🥳 🎉')

                awaiting_srcs.remove(sid)
                if not awaiting_srcs:
                    send(ws, 'Debugger.resume')
    except Exception:
        print_exc()


def logUrl(url):
    url = sub('//www.google.[a-z]+', '//www.google.com', url)
    if '=!' in url:
        datas = findall(r'(?<==)\!\d[^/?&]+', url)
    else:
        datas = [i.strip('?&') for i in findall(r'\?(?:\d[^&]+&?)+', url)]
    if not datas:
        return
    for data in sorted(datas, key=len, reverse=True):
        if data in sent_msgs:
            proto = sent_msgs[data]

            url = url.replace('/pb=', '?pb=')
            if '?' not in url:
                url += '?'
            url, qs = url.split('?')

            if '=!' in url + '?' + qs:
                qsl = OrderedDict(parse_qsl(qs, True))
                for k, v in findall(
                    '/(space/|place/|search/|contrib/|[^/]+=|@)([^/]*)', url
                ) + findall('/(dir/)([^/]*/[^/]*)', url):
                    k2 = k.strip('/=') if k != '@' else 'coords'
                    if '/' not in v:
                        qsl[k2] = unquote_plus(v)
                        url = url.replace(k + v, k + '{' + k2 + '}')
                    else:
                        qsl[k2 + '1'], qsl[k2 + '2'] = map(
                            unquote_plus, v.split('/')
                        )
                        url = url.replace(k + v, k + '{%s1}/{%s2}' % (k2, k2))
                qs = urlencode(qsl)

                pb_param = next(
                    k for k, v in qsl.items() if v == unquote_plus(data)
                )
                endpoints.append(
                    {
                        'transport': 'pburl_private',
                        'proto_path': '',
                        'proto_msg': '',
                        'url': url,
                        'pb_param': pb_param,
                        'samples': [
                            GMapsAPIPrivate(pb_param, url).serialize_sample(qs)
                        ],
                    }
                )

            else:
                endpoints.append(
                    {
                        'transport': 'pburl_public',
                        'proto_path': '',
                        'proto_msg': '',
                        'url': url,
                        'samples': [
                            GMapsAPIPublic(None, url).serialize_sample(qs)
                        ],
                    }
                )

            info('[i] Captured URL: %r - %r - %r' % (url, qs, data))
            if proto not in proto_to_urls:
                proto_to_urls[proto] = set()
            proto_to_urls[proto].add(url)
            return

    info(
        '[i] Did not capture URL: %r - %r - across => %r'
        % (url, data, sent_msgs)
    )


def main():
    getLogger().setLevel(DEBUG)

    extractor_main('pburl_extract')


if __name__ == '__main__':
    main()
