diff --git a/umbra/umbra.py b/umbra/umbra.py index 4fa9d2f..f116d1c 100755 --- a/umbra/umbra.py +++ b/umbra/umbra.py @@ -4,28 +4,56 @@ import os,sys,argparse, urllib2 import websocket import thread import time +import uuid +import logging +import threading +logging.basicConfig(level=logging.DEBUG) -def on_message(ws, message): - message = loads(message) - if "method" in message.keys() and message["method"] == "Network.requestWillBeSent": - print message +class Umbra: + def __init__(self, port): + self.cmd_id = 0 + self.chrome_debug_port = port + self.launch_tab_socket = self.get_websocket(self.on_open) + self.launch_tab_socket.run_forever() - -def on_error(ws, error): - print error - -def on_close(ws): - print "### closed ###" - -def on_open(ws): - cmd = {} - cmd['id'] = 1001 - cmd['method'] = "Network.enable" - ws.send(dumps(cmd)) - cmd['id'] = 1002 - cmd['method'] = "Runtime.evaluate" - cmd["params"] = { "expression" : "document.location = 'http://archive.org'"} - ws.send(dumps(cmd)) + def get_websocket(self, on_open, url=None): + debug_info = loads(urllib2.urlopen("http://localhost:%s/json" % self.chrome_debug_port).read()) + if url: #Polling for the data url we used to initialize the window + while not filter(lambda x: x['url'] == url, debug_info): + debug_info = loads(urllib2.urlopen("http://localhost:%s/json" % self.chrome_debug_port).read()) + time.sleep(0.5) + debug_info = filter(lambda x: x['url'] == url, debug_info) + return_socket = websocket.WebSocketApp(debug_info[0]['webSocketDebuggerUrl'], on_message = self.on_message) + return_socket.on_open = on_open + return return_socket + + def on_message(self, ws, message): + message = loads(message) + if "method" in message.keys() and message["method"] == "Network.requestWillBeSent": + pass #print message + + def on_open(self, ws): + self.fetch_url("http://archive.org") + time.sleep(10) + sys.exit(0) + + def send_command(self,tab=None, **kwargs): + if not tab: + tab = self.launch_tab_socket + command = {} + command.update(kwargs) + self.cmd_id += 1 + command['id'] = self.cmd_id + tab.send(dumps(command)) + + def fetch_url(self, url): + new_page = 'data:text/html;charset=utf-8,%s' % str(uuid.uuid4()) + self.send_command(method="Runtime.evaluate", params={"expression":"window.open('%s');" % new_page}) + def on_open(ws): + self.send_command(tab=ws, method="Network.enable") + self.send_command(tab=ws, method="Runtime.evaluate", params={"expression":"document.location = '%s';" % url}) + socket = self.get_websocket(on_open, new_page) + threading.Thread(target=socket.run_forever).start() class Chrome(): def __init__(self, port): @@ -55,12 +83,4 @@ if __name__ == "__main__": help='Port to have invoked chrome listen on for debugging connections') args = arg_parser.parse_args(args=sys.argv[1:]) with Chrome(args.port): - debug_info = loads(urllib2.urlopen("http://localhost:%s/json" % args.port).read()) - url = debug_info[0]['webSocketDebuggerUrl'] - ws = websocket.WebSocketApp(url, - on_message = on_message, - on_error = on_error, - on_close = on_close) - ws.on_open = on_open - ws.run_forever() - + Umbra(args.port)