diff options
Diffstat (limited to 'backend/injector.py')
| -rw-r--r-- | backend/injector.py | 196 |
1 files changed, 188 insertions, 8 deletions
diff --git a/backend/injector.py b/backend/injector.py index 82436adf..de914d17 100644 --- a/backend/injector.py +++ b/backend/injector.py @@ -1,7 +1,7 @@ -#Injector code from https://github.com/SteamDeckHomebrew/steamdeck-ui-inject. More info on how it works there. +# Injector code from https://github.com/SteamDeckHomebrew/steamdeck-ui-inject. More info on how it works there. from asyncio import sleep -from logging import debug, getLogger +from logging import getLogger from traceback import format_exc from aiohttp import ClientSession @@ -11,7 +11,10 @@ BASE_ADDRESS = "http://localhost:8080" logger = getLogger("Injector") + class Tab: + cmd_id = 0 + def __init__(self, res) -> None: self.title = res["title"] self.id = res["id"] @@ -24,14 +27,24 @@ class Tab: self.client = ClientSession() self.websocket = await self.client.ws_connect(self.ws_url) + async def close_websocket(self): + await self.client.close() + async def listen_for_message(self): async for message in self.websocket: - yield message + data = message.json() + yield data async def _send_devtools_cmd(self, dc, receive=True): if self.websocket: + self.cmd_id += 1 + dc["id"] = self.cmd_id await self.websocket.send_json(dc) - return (await self.websocket.receive_json()) if receive else None + if receive: + async for msg in self.listen_for_message(): + if "id" in msg and msg["id"] == dc["id"]: + return msg + return None raise RuntimeError("Websocket not opened") async def evaluate_js(self, js, run_async=False, manage_socket=True, get_result=True): @@ -39,7 +52,6 @@ class Tab: await self.open_websocket() res = await self._send_devtools_cmd({ - "id": 1, "method": "Runtime.evaluate", "params": { "expression": js, @@ -49,9 +61,172 @@ class Tab: }, get_result) if manage_socket: - await self.client.close() + await self.close_websocket() return res + async def enable(self): + """ + Enables page domain notifications. + """ + await self._send_devtools_cmd({ + "method": "Page.enable", + }, False) + + async def disable(self): + """ + Disables page domain notifications. + """ + await self._send_devtools_cmd({ + "method": "Page.disable", + }, False) + + async def reload_and_evaluate(self, js, manage_socket=True): + """ + Reloads the current tab, with JS to run on load via debugger + """ + if manage_socket: + await self.open_websocket() + + await self._send_devtools_cmd({ + "method": "Debugger.enable" + }, True) + + await self._send_devtools_cmd({ + "method": "Runtime.evaluate", + "params": { + "expression": "location.reload();", + "userGesture": True, + "awaitPromise": False + } + }, False) + + breakpoint_res = await self._send_devtools_cmd({ + "method": "Debugger.setInstrumentationBreakpoint", + "params": { + "instrumentation": "beforeScriptExecution" + } + }, True) + + logger.info(breakpoint_res) + + # Page finishes loading when breakpoint hits + + for x in range(20): + # this works around 1/5 of the time, so just send it 8 times. + # the js accounts for being injected multiple times allowing only one instance to run at a time anyway + await self._send_devtools_cmd({ + "method": "Runtime.evaluate", + "params": { + "expression": js, + "userGesture": True, + "awaitPromise": False + } + }, False) + + await self._send_devtools_cmd({ + "method": "Debugger.removeBreakpoint", + "params": { + "breakpointId": breakpoint_res["result"]["breakpointId"] + } + }, False) + + for x in range(4): + await self._send_devtools_cmd({ + "method": "Debugger.resume" + }, False) + + await self._send_devtools_cmd({ + "method": "Debugger.disable" + }, True) + + if manage_socket: + await self.close_websocket() + return + + async def add_script_to_evaluate_on_new_document(self, js, add_dom_wrapper=True, manage_socket=True, get_result=True): + """ + How the underlying call functions is not particularly clear from the devtools docs, so stealing puppeteer's description: + + Adds a function which would be invoked in one of the following scenarios: + * whenever the page is navigated + * whenever the child frame is attached or navigated. In this case, the + function is invoked in the context of the newly attached frame. + + The function is invoked after the document was created but before any of + its scripts were run. This is useful to amend the JavaScript environment, + e.g. to seed `Math.random`. + + Parameters + ---------- + js : str + The script to evaluate on new document + add_dom_wrapper : bool + True to wrap the script in a wait for the 'DOMContentLoaded' event. + DOM will usually not exist when this execution happens, + so it is necessary to delay til DOM is loaded if you are modifying it + manage_socket : bool + True to have this function handle opening/closing the websocket for this tab + get_result : bool + True to wait for the result of this call + + Returns + ------- + int or None + The identifier of the script added, used to remove it later. + (see remove_script_to_evaluate_on_new_document below) + None is returned if `get_result` is False + """ + + wrappedjs = """ + function scriptFunc() { + {js} + } + if (document.readyState === 'loading') { + addEventListener('DOMContentLoaded', () => { + scriptFunc(); + }); + } else { + scriptFunc(); + } + """.format(js=js) if add_dom_wrapper else js + + if manage_socket: + await self.open_websocket() + + res = await self._send_devtools_cmd({ + "method": "Page.addScriptToEvaluateOnNewDocument", + "params": { + "source": wrappedjs + } + }, get_result) + + if manage_socket: + await self.close_websocket() + return res + + async def remove_script_to_evaluate_on_new_document(self, script_id, manage_socket=True): + """ + Removes a script from a page that was added with `add_script_to_evaluate_on_new_document` + + Parameters + ---------- + script_id : int + The identifier of the script to remove (returned from `add_script_to_evaluate_on_new_document`) + """ + + if manage_socket: + await self.open_websocket() + + res = await self._send_devtools_cmd({ + "method": "Page.removeScriptToEvaluateOnNewDocument", + "params": { + "identifier": script_id + } + }, False) + + if manage_socket: + await self.close_websocket() + async def get_steam_resource(self, url): res = await self.evaluate_js(f'(async function test() {{ return await (await fetch("{url}")).text() }})()', True) return res["result"]["result"]["value"] @@ -59,6 +234,7 @@ class Tab: def __repr__(self): return self.title + async def get_tabs(): async with ClientSession() as web: res = {} @@ -80,6 +256,7 @@ async def get_tabs(): else: raise Exception(f"/json did not return 200. {await res.text()}") + async def get_tab(tab_name): tabs = await get_tabs() tab = next((i for i in tabs if i.title == tab_name), None) @@ -87,11 +264,13 @@ async def get_tab(tab_name): raise ValueError(f"Tab {tab_name} not found") return tab + async def inject_to_tab(tab_name, js, run_async=False): tab = await get_tab(tab_name) return await tab.evaluate_js(js, run_async) + async def tab_has_global_var(tab_name, var_name): try: tab = await get_tab(tab_name) @@ -104,14 +283,15 @@ async def tab_has_global_var(tab_name, var_name): return res["result"]["result"]["value"] + async def tab_has_element(tab_name, element_name): try: tab = await get_tab(tab_name) except ValueError: return False res = await tab.evaluate_js(f"document.getElementById('{element_name}') != null", False) - + if not "result" in res or not "result" in res["result"] or not "value" in res["result"]["result"]: return False - return res["result"]["result"]["value"]
\ No newline at end of file + return res["result"]["result"]["value"] |