import frida
import sys
import os

JS_SCRIPT = '''
setTimeout(function () {{
  Java.perform(function () {{
      const File = Java.use("java.io.File");
      const Bitmap = Java.use("android.graphics.Bitmap");
      const BitmapCompressFormat = Java.use("android.graphics.Bitmap$CompressFormat");
      const BitmapConfig = Java.use("android.graphics.Bitmap$Config");
      const ByteArrayOutputStream = Java.use("java.io.ByteArrayOutputStream");

      const bitmap = Bitmap.createBitmap(100, 100, BitmapConfig.ARGB_8888.value);
      /* snipped -- render pixels to bitmap */

      // output bitmap to byte array
      const stream = ByteArrayOutputStream.$new();
      const saved = bitmap.compress(BitmapCompressFormat.PNG.value, 100, stream);
      console.log("[*] Compressed as PNG:", saved);
      const byteArray = stream.toByteArray();
      console.log("[*] Byte array length:", byteArray.length);

      send({{ type: "bitmap", page: pageNum }}, byteArray);
      stream.close();
  }});
}}, 1000);
'''

def on_message(message, data):
    if message["type"] == "send" and message["payload"].get("type") == "bitmap":
        page = message["payload"].get("page")
        with open(OUTPUT_FILENAME, "wb") as f:
            f.write(data)
        print(f"[+] Saved page {page} as {OUTPUT_FILENAME}")
    else:
        print(f"[?] Unknown message: {message}")

def main():
    device = frida.get_usb_device(timeout=5)
    session = device.attach(pid)
    script = session.create_script(agent_code)
    script.on("message", on_message)
    script.load()
    device.resume(pid)

if __name__ == "__main__":
    main()