diff --git a/src/code.py b/src/code.py index 1abdf05..b67ec09 100644 --- a/src/code.py +++ b/src/code.py @@ -9,12 +9,11 @@ import socketpool import wifi from adafruit_httpserver import POST, FileResponse, Request, Server -from ducky import run_script_file +from ducky import run_boot_script from api import handle - -async def main(): +async def setup_server(): """ Begin a wifi access point defined by the SSID and PASSWORD environment variables. @@ -45,19 +44,16 @@ def api(request: Request): server.serve_forever(str(wifi.radio.ipv4_address_ap)) -async def run_boot_script(): - """ - If a script with the name 'boot.dd' exists, - run it without user interaction on boot. - """ - boot_script = 'payloads/boot.dd' - if os.path.exists(boot_script): - run_script_file(boot_script) +async def main(): + await asyncio.gather( + run_boot_script(), + setup_server() + ) if __name__ == "__main__": try: - asyncio.gather(main(), run_boot_script()) + asyncio.run(main()) # For some reason, wifi.stop_ap is not implemented. except NotImplementedError: microcontroller.reset() diff --git a/src/ducky.py b/src/ducky.py index 9babdac..65c31b2 100644 --- a/src/ducky.py +++ b/src/ducky.py @@ -118,3 +118,15 @@ def run_script_file(path: str): run_script(handle.read()) except OSError as error: warn(f"unable to open file {path}: {error}") + + +async def run_boot_script(): + """ + If a script with the name 'boot.dd' exists, + run it without user interaction on boot. + """ + try: + with open("payloads/boot.dd", "r", encoding="utf-8") as handle: + run_script(handle.read()) + except OSError: + info("boot script does not exist, skipping its execution")