kaiscr.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. #!/usr/bin/python3
  2. import socket
  3. import base64
  4. import json
  5. import sys
  6. class TakeScreenshot:
  7. def __init__(self, host="127.0.0.1", port=6000):
  8. self.screenshot_cmd = '{"type":"screenshotToDataURL","to":"%s"}'
  9. self.listTabs_cmd = '{"to":"root","type":"listTabs"}'
  10. self.substring_cmd='{"type":"substring","start":%d,"end":%d,"to":"%s"}'
  11. self.sock = socket.socket()
  12. self.sock.connect((host, port))
  13. buffer = b""
  14. while not buffer.endswith(b":"):
  15. buffer += self.sock.recv(1)
  16. size = int(buffer.replace(b":", b""))
  17. buffer = b""
  18. while len(buffer) < size:
  19. buffer += self.sock.recv(size)
  20. # Do something with it if you want.
  21. self.sock.send(self.__with_len(self.listTabs_cmd))
  22. buffer = b""
  23. while not buffer.endswith(b":"):
  24. buffer += self.sock.recv(1)
  25. size = int(buffer.replace(b":", b"").decode())
  26. self.deviceActor = json.loads(self.sock.recv(size))["deviceActor"]
  27. def __with_len(self, command):
  28. return str(len(command)).encode() + b":" + command.encode()
  29. def __receive(self, size):
  30. buffer = b""
  31. while len(buffer) < size:
  32. buffer += self.sock.recv(1)
  33. return buffer
  34. def screenshot(self):
  35. cmd = self.__with_len(self.screenshot_cmd % self.deviceActor)
  36. self.sock.send(cmd)
  37. buffer = b""
  38. while not buffer.endswith(b":"):
  39. buffer += self.sock.recv(1)
  40. size = int(buffer.replace(b":", b"").decode())
  41. buffer = self.__receive(size)
  42. image = json.loads(buffer)["value"]
  43. if type(image) == str:
  44. return base64.b64decode(image.split(",")[1])
  45. image_len = image["length"]
  46. actor = image["actor"]
  47. cmd = self.substring_cmd % (0, image_len, actor)
  48. self.sock.send(self.__with_len(cmd))
  49. buffer = b""
  50. while not buffer.endswith(b":"):
  51. buffer += self.sock.recv(1)
  52. buffer = self.__receive(int(buffer.replace(b":", b"")))
  53. image = json.loads(buffer)["substring"].split(",")[1]
  54. image += "=" * (-len(image) % 4)
  55. return base64.b64decode(image)
  56. def close(self):
  57. self.sock.close()
  58. if __name__ == "__main__":
  59. from argparse import ArgumentParser
  60. d = "Take screenshot(s) from a KaiOS/FFOS device"
  61. parser = ArgumentParser(description = d)
  62. parser.add_argument("--host",
  63. metavar = "host",
  64. type = str,
  65. default = "127.0.0.1",
  66. help = "The host to connect to")
  67. parser.add_argument("--port",
  68. metavar = "port",
  69. type = int,
  70. default = 6000,
  71. help = "The port to connect to")
  72. parser.add_argument("--prefix",
  73. metavar = "prefix",
  74. type = str,
  75. default = "out",
  76. help = "The prefix for naming files")
  77. parser.add_argument("--count",
  78. metavar = "count",
  79. type = int,
  80. default = 1,
  81. help = "How many times should I take screenshot? 0 does nothing and negative values take screenshot forever")
  82. args = parser.parse_args()
  83. if args.count == 0:
  84. sys.exit(0)
  85. takeScreenshot = TakeScreenshot(args.host, args.port)
  86. c = 1
  87. try:
  88. while c <= args.count or args.count < 0:
  89. c += 1
  90. with open(args.prefix + str(c) + ".png", "wb") as fp:
  91. fp.write(takeScreenshot.screenshot())
  92. finally:
  93. takeScreenshot.close()