PyChromeDevTools.py

#!/usr/bin/python3

import json
import time

import requests
import websocket
from Exceptions import AutoChromeException

TIMEOUT = 1


class GenericElement(object):
    def __init__(self, name, parent):
        self.name = name
        self.parent = parent

    def __getattr__(self, attr):
        func_name = '{}.{}'.format(self.name, attr)

        def generic_function(**args):
            self.parent.pop_messages()
            self.parent.message_counter += 1
            message_id = self.parent.message_counter
            params = {}
            for key in args:
                if args[key] is not None:
                    params[key] = args[key]
            call_obj = {'id': message_id, 'method': func_name, 'params': params}
            data = json.dumps(call_obj)
            print('>>> send: ' + data)
            self.parent.ws.send(data)
            result, _ = self.parent.wait_result(message_id)
            if result is not None:
                if 'error' in result:
                    raise AutoChromeException(result['error']['data'])
                else:
                    return result['result']
            else:
                return None

        return generic_function


class ChromeInterface(object):
    message_counter = 0

    def __init__(self, host='localhost', port=9222, tab=0, timeout=TIMEOUT, auto_connect=True):
        self.host = host
        self.port = port
        self.ws = None
        self.tabs = None
        self.timeout = timeout
        if auto_connect:
            self.connect(tab=tab)

    def get_tabs(self):
        response = requests.get('http://{}:{}/json'.format(self.host, self.port))
        self.tabs = json.loads(response.text)

    def connect(self, tab=0, update_tabs=True):
        if update_tabs or self.tabs is None:
            self.get_tabs()
        wsurl = self.tabs[tab]['webSocketDebuggerUrl']
        self.close()
        self.ws = websocket.create_connection(wsurl)
        self.ws.settimeout(self.timeout)

    def connect_targetID(self, targetID):
        try:
            wsurl = 'ws://{}:{}/devtools/page/{}'.format(self.host, self.port, targetID)
            self.close()
            self.ws = websocket.create_connection(wsurl)
            self.ws.settimeout(self.timeout)
        except:
            wsurl = self.tabs[0]['webSocketDebuggerUrl']
            self.ws = websocket.create_connection(wsurl)
            self.ws.settimeout(self.timeout)

    def close(self):
        if self.ws:
            self.ws.close()

    # Blocking
    def wait_message(self, timeout=None):
        timeout = timeout if timeout is not None else self.timeout
        self.ws.settimeout(timeout)
        try:
            message = self.ws.recv()
        except:
            return None
        finally:
            self.ws.settimeout(self.timeout)
        return json.loads(message)

    # Blocking
    def wait_event(self, event, timeout=None):
        timeout = timeout if timeout is not None else self.timeout
        start_time = time.time()
        messages = []
        matching_message = None
        while True:
            now = time.time()
            if now - start_time > timeout:
                break
            try:
                message = self.ws.recv()
                print('<<< event: ' + message)
                parsed_message = json.loads(message)
                messages.append(parsed_message)
                if 'method' in parsed_message and parsed_message['method'] == event:
                    matching_message = parsed_message['params']
                    break
            except:
                break
        return matching_message, messages

    # Blocking
    def wait_result(self, result_id, timeout=None):
        timeout = timeout if timeout is not None else self.timeout
        start_time = time.time()
        messages = []
        matching_result = None
        while True:
            now = time.time()
            if now - start_time > timeout:
                break
            try:
                message = self.ws.recv()
                print('<<< receive: ' + message)
                parsed_message = json.loads(message)
                messages.append(parsed_message)
                if ('result' in parsed_message or 'error' in parsed_message) and parsed_message['id'] == result_id:
                    matching_result = parsed_message
                    break
            except:
                break
        return matching_result, messages

    # Non Blocking
    def pop_messages(self):
        messages = []
        self.ws.settimeout(0)
        while True:
            try:
                message = self.ws.recv()
                messages.append(json.loads(message))
            except:
                break
        self.ws.settimeout(self.timeout)
        return messages

    def __getattr__(self, attr):
        genericelement = GenericElement(attr, self)
        self.__setattr__(attr, genericelement)
        return genericelement