<commit_before>#!/usr/bin/env python
# vim: set sw=4 et:

from json import dumps, loads
from itertools import count
import os,sys,argparse, urllib.request, urllib.error, urllib.parse
import websocket
import time
import uuid
import logging
import threading
import subprocess
import signal
from kombu import Connection, Exchange, Queue

class UmbraWorker:
    logger = logging.getLogger('umbra.UmbraWorker')

    def __init__(self, umbra, chrome_port=9222, chrome_exe='chromium-browser', chrome_wait=10, client_id='request'):
        self.command_id = count(1)
        self.lock = threading.Lock()
        self.umbra = umbra
        self.chrome_port = chrome_port
        self.chrome_exe = chrome_exe
        self.chrome_wait = chrome_wait
        self.client_id = client_id
        self.page_done = threading.Event()
        self.idle_timer = None

    def browse_page(self, url, url_metadata):
        with self.lock:
            self.url = url
            self.url_metadata = url_metadata
            with Chrome(self.chrome_port, self.chrome_exe, self.chrome_wait) as websocket_url:
                websock = websocket.WebSocketApp(websocket_url,
                        on_open=self.visit_page, on_message=self.handle_message)
                websock_thread = threading.Thread(target=websock.run_forever)
                websock_thread.start()

                self.page_done.clear()
                self._reset_idle_timer()
                while not self.page_done.is_set():
                    time.sleep(0.5)

                websock.close()
                self.idle_timer = None

    def _reset_idle_timer(self):
        if self.idle_timer:
            self.idle_timer.cancel()
        self.idle_timer = threading.Timer(60, self.page_done.set)
        self.idle_timer.start()

    def visit_page(self, websock):
        msg = dumps(dict(method="Network.enable", id=next(self.command_id)))
        self.logger.debug('sending message to {}: {}'.format(websock, msg))
        websock.send(msg)

        msg = dumps(dict(method="Page.enable", id=next(self.command_id)))
        self.logger.debug('sending message to {}: {}'.format(websock, msg))
        websock.send(msg)

        msg = dumps(dict(method="Page.navigate", id=next(self.command_id), params={"url": self.url}))
        self.logger.debug('sending message to {}: {}'.format(websock, msg))
        websock.send(msg)

    def send_request_to_amqp(self, chrome_msg):
        payload = chrome_msg['params']['request']
        payload['parentUrl'] = self.url
        payload['parentUrlMetadata'] = self.url_metadata
        self.logger.debug('sending to amqp exchange={} routing_key={} payload={}'.format(self.umbra.umbra_exchange, self.client_id, payload))
        with self.umbra.producer_lock:
            self.umbra.producer.publish(payload,
                    exchange=self.umbra.umbra_exchange,
                    routing_key=self.client_id)

    def handle_message(self, websock, message):
        self.logger.debug("handling message from websocket {} - {}".format(websock, message[:95]))
        self._reset_idle_timer()
        message = loads(message)
        if "method" in message.keys() and message["method"] == "Network.requestWillBeSent":
            self.send_request_to_amqp(message)
        elif "method" in message.keys() and message["method"] == "Page.loadEventFired":
            self.logger.debug("got Page.loadEventFired, starting behaviors for {}".format(self.url))
            from umbra import behaviors
            behaviors.execute(self.url, websock, self.command_id)            

    def get_message_handler(self, url, url_metadata, command_id):
        this_watchdog = self.watchdog(command_id)
        def handle_message(ws, message):
            this_watchdog.send(ws)
            message = loads(message)
            if "method" in message.keys() and message["method"] == "Network.requestWillBeSent":
                to_send = {}
                to_send.update(message['params']['request'])
                to_send.update(dict(parentUrl=url,parentUrlMetadata=url_metadata))
                self.logger.debug('sending to amqp: {}'.format(to_send))
                with self.producer_lock:
                    self.producer.publish(to_send,
                            routing_key='request',
                            exchange=self.umbra_exchange)
        return handle_message

class Umbra:
    logger = logging.getLogger('umbra.Umbra')

    def __init__(self, amqp_url, chrome_exe, browser_wait):
        self.amqp_url = amqp_url
        self.chrome_exe = chrome_exe
        self.browser_wait = browser_wait
        self.producer = None
        self.producer_lock = None
        self.workers = {}
        self.workers_lock = threading.Lock()
        self.amqp_thread = threading.Thread(target=self.consume_amqp)
        self.amqp_stop = threading.Event()
        self.amqp_thread.start()

    def shutdown(self):
        self.logger.info("shutting down amqp consumer {}".format(self.amqp_url))
        self.amqp_stop.set()
        self.amqp_thread.join()

    def consume_amqp(self):
        self.umbra_exchange = Exchange(name='umbra', type='direct', durable=True)
        url_queue = Queue('urls', routing_key='url', exchange=self.umbra_exchange)
        self.logger.info("connecting to amqp {} at {}".format(repr(self.umbra_exchange), self.amqp_url))
        with Connection(self.amqp_url) as conn:
            self.producer = conn.Producer(serializer='json')
            self.producer_lock = threading.Lock()
            with conn.Consumer(url_queue, callbacks=[self.fetch_url]) as consumer:
                import socket
                while not self.amqp_stop.is_set():
                    try:
                        conn.drain_events(timeout=0.5)
                    except socket.timeout:
                        pass

    def fetch_url(self, body, message):
        client_id = body['clientId']
        with self.workers_lock:
            if not client_id in self.workers:
                port = 9222 + len(self.workers)
                t = UmbraWorker(umbra=self, chrome_port=port,
                        chrome_exe=self.chrome_exe,
                        chrome_wait=self.browser_wait,
                        client_id=client_id)
                self.workers[client_id] = t

        def browse_page_async():
            self.logger.info('client_id={} body={}'.format(client_id, body))
            self.workers[client_id].browse_page(body['url'], body['metadata'])

        threading.Thread(target=browse_page_async).start()

        message.ack()


class Chrome:
    logger = logging.getLogger('umbra.Chrome')

    def __init__(self, port, executable, browser_wait):
        self.port = port
        self.executable = executable
        self.browser_wait = browser_wait

    def fetch_debugging_json():
        raw_json = urllib.request.urlopen("http://localhost:%s/json" % self.port).read()
        json = raw_json.decode('utf-8')
        return loads(json)

    # returns websocket url to chrome window with about:blank loaded
    def __enter__(self):
        chrome_args = [self.executable, "--temp-profile",
                "--remote-debugging-port=%s" % self.port,
                "--disable-web-sockets", "--disable-cache",
                "--window-size=1100,900", "--enable-logging",
                "--homepage=about:blank", "about:blank"]
        self.logger.info("running {}".format(chrome_args))
        self.chrome_process = subprocess.Popen(chrome_args, start_new_session=True)
        self.logger.info("chrome running, pid {}".format(self.chrome_process.pid))
        start = time.time()

        json_url = "http://localhost:%s/json" % self.port<commit_msg>create temp dir for user profile rather than rely on --temp-profile<commit_after>#!/usr/bin/env python
# vim: set sw=4 et:

from json import dumps, loads
from itertools import count
import os,sys,argparse, urllib.request, urllib.error, urllib.parse
import websocket
import time
import uuid
import logging
import threading
import subprocess
import signal
from kombu import Connection, Exchange, Queue
import tempfile

class UmbraWorker:
    logger = logging.getLogger('umbra.UmbraWorker')

    def __init__(self, umbra, chrome_port=9222, chrome_exe='chromium-browser', chrome_wait=10, client_id='request'):
        self.command_id = count(1)
        self.lock = threading.Lock()
        self.umbra = umbra
        self.chrome_port = chrome_port
        self.chrome_exe = chrome_exe
        self.chrome_wait = chrome_wait
        self.client_id = client_id
        self.page_done = threading.Event()
        self.idle_timer = None

    def browse_page(self, url, url_metadata):
        with self.lock:
            self.url = url
            self.url_metadata = url_metadata
            with tempfile.TemporaryDirectory() as user_data_dir:
                with Chrome(self.chrome_port, self.chrome_exe, self.chrome_wait, user_data_dir) as websocket_url:
                    websock = websocket.WebSocketApp(websocket_url,
                            on_open=self.visit_page, on_message=self.handle_message)
                    websock_thread = threading.Thread(target=websock.run_forever)
                    websock_thread.start()

                    self.page_done.clear()
                    self._reset_idle_timer()
                    while not self.page_done.is_set():
                        time.sleep(0.5)

                    websock.close()
                    self.idle_timer = None

    def _reset_idle_timer(self):
        if self.idle_timer:
            self.idle_timer.cancel()
        self.idle_timer = threading.Timer(60, self.page_done.set)
        self.idle_timer.start()

    def visit_page(self, websock):
        msg = dumps(dict(method="Network.enable", id=next(self.command_id)))
        self.logger.debug('sending message to {}: {}'.format(websock, msg))
        websock.send(msg)

        msg = dumps(dict(method="Page.enable", id=next(self.command_id)))
        self.logger.debug('sending message to {}: {}'.format(websock, msg))
        websock.send(msg)

        msg = dumps(dict(method="Page.navigate", id=next(self.command_id), params={"url": self.url}))
        self.logger.debug('sending message to {}: {}'.format(websock, msg))
        websock.send(msg)

    def send_request_to_amqp(self, chrome_msg):
        payload = chrome_msg['params']['request']
        payload['parentUrl'] = self.url
        payload['parentUrlMetadata'] = self.url_metadata
        self.logger.debug('sending to amqp exchange={} routing_key={} payload={}'.format(self.umbra.umbra_exchange, self.client_id, payload))
        with self.umbra.producer_lock:
            self.umbra.producer.publish(payload,
                    exchange=self.umbra.umbra_exchange,
                    routing_key=self.client_id)

    def handle_message(self, websock, message):
        self.logger.debug("handling message from websocket {} - {}".format(websock, message[:95]))
        self._reset_idle_timer()
        message = loads(message)
        if "method" in message.keys() and message["method"] == "Network.requestWillBeSent":
            self.send_request_to_amqp(message)
        elif "method" in message.keys() and message["method"] == "Page.loadEventFired":
            self.logger.debug("got Page.loadEventFired, starting behaviors for {}".format(self.url))
            from umbra import behaviors
            behaviors.execute(self.url, websock, self.command_id)

    def get_message_handler(self, url, url_metadata, command_id):
        this_watchdog = self.watchdog(command_id)
        def handle_message(ws, message):
            this_watchdog.send(ws)
            message = loads(message)
            if "method" in message.keys() and message["method"] == "Network.requestWillBeSent":
                to_send = {}
                to_send.update(message['params']['request'])
                to_send.update(dict(parentUrl=url,parentUrlMetadata=url_metadata))
                self.logger.debug('sending to amqp: {}'.format(to_send))
                with self.producer_lock:
                    self.producer.publish(to_send,
                            routing_key='request',
                            exchange=self.umbra_exchange)
        return handle_message

class Umbra:
    logger = logging.getLogger('umbra.Umbra')

    def __init__(self, amqp_url, chrome_exe, browser_wait):
        self.amqp_url = amqp_url
        self.chrome_exe = chrome_exe
        self.browser_wait = browser_wait
        self.producer = None
        self.producer_lock = None
        self.workers = {}
        self.workers_lock = threading.Lock()
        self.amqp_thread = threading.Thread(target=self.consume_amqp)
        self.amqp_stop = threading.Event()
        self.amqp_thread.start()

    def shutdown(self):
        self.logger.info("shutting down amqp consumer {}".format(self.amqp_url))
        self.amqp_stop.set()
        self.amqp_thread.join()

    def consume_amqp(self):
        self.umbra_exchange = Exchange(name='umbra', type='direct', durable=True)
        url_queue = Queue('urls', routing_key='url', exchange=self.umbra_exchange)
        self.logger.info("connecting to amqp {} at {}".format(repr(self.umbra_exchange), self.amqp_url))
        with Connection(self.amqp_url) as conn:
            self.producer = conn.Producer(serializer='json')
            self.producer_lock = threading.Lock()
            with conn.Consumer(url_queue, callbacks=[self.fetch_url]) as consumer:
                import socket
                while not self.amqp_stop.is_set():
                    try:
                        conn.drain_events(timeout=0.5)
                    except socket.timeout:
                        pass

    def fetch_url(self, body, message):
        client_id = body['clientId']
        with self.workers_lock:
            if not client_id in self.workers:
                port = 9222 + len(self.workers)
                t = UmbraWorker(umbra=self, chrome_port=port,
                        chrome_exe=self.chrome_exe,
                        chrome_wait=self.browser_wait,
                        client_id=client_id)
                self.workers[client_id] = t

        def browse_page_async():
            self.logger.info('client_id={} body={}'.format(client_id, body))
            self.workers[client_id].browse_page(body['url'], body['metadata'])

        threading.Thread(target=browse_page_async).start()

        message.ack()


class Chrome:
    logger = logging.getLogger('umbra.Chrome')

    def __init__(self, port, executable, browser_wait, user_data_dir):
        self.port = port
        self.executable = executable
        self.browser_wait = browser_wait
        self.user_data_dir = user_data_dir

    def fetch_debugging_json():
        raw_json = urllib.request.urlopen("http://localhost:%s/json" % self.port).read()
        json = raw_json.decode('utf-8')
        return loads(json)

    # returns websocket url to chrome window with about:blank loaded
    def __enter__(self):
        chrome_args = [self.executable,
                "--user-data-dir={}".format(self.user_data_dir),
                "--remote-debugging-port=%s" % self.port,
                "--disable-web-sockets", "--disable-cache",
                "--window-size=1100,900", "--enable-logging",
                "--homepage=about:blank", "about:blank"]
        self.logger.info("running {}".format(chrome_args))
        self.chrome_process = subprocess.Popen(chrome_args, start_new_session=True)
        self.logger.info("chrome running, pid {}".format(self.chrome_process.pid))
        start = time.time()

        json_url = "http://localhost:%s/json" % self.port
{'fraction_non_alphanumeric': 0.11472897912277341, 'fraction_numerical': 0.003575304858583924, 'mean_word_length': 2.3993055555555554, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2524468', 'n_tokens_mistral': 4763, 'n_tokens_neox': 4545, 'n_words': 833}