From ad6a97c2978ef96537a5475cea397f0b05a9eb81 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sof=C3=ADa=20Aritz?= <sofi@sofiaritz.com>
Date: Thu, 10 Oct 2024 20:16:23 +0200
Subject: [PATCH] Initial commit

---
 .gitignore                  |  9 ++++++
 main.py                     | 51 +++++++++++++++++++++++++++++++++
 screenshots/.gitkeep        |  0
 server.py                   | 57 +++++++++++++++++++++++++++++++++++++
 server_screenshots/.gitkeep |  0
 5 files changed, 117 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 main.py
 create mode 100644 screenshots/.gitkeep
 create mode 100644 server.py
 create mode 100644 server_screenshots/.gitkeep

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..35691af
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,9 @@
+/__pycache__
+
+/screenshots/*
+!/screenshots/.gitkeep
+
+/server_screenshots/*
+!/server_screenshots/.gitkeep
+
+/assets
\ No newline at end of file
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..afcd3d0
--- /dev/null
+++ b/main.py
@@ -0,0 +1,51 @@
+from PIL import Image, ImageTk, ImageGrab
+from datetime import datetime
+from server import start_server
+import requests
+import argparse
+import base64
+import webbrowser
+import tkinter.font
+import pytesseract, clipboard, tkinter as tk
+
+formatted_date = datetime.today().strftime('Screenshot from %Y-%m-%d %H-%M-%S')
+
+parser = argparse.ArgumentParser("screenshot-magic")
+parser.add_argument("--server", help="Run server", type=bool, default=False)
+args = parser.parse_args()
+
+if args.server == True:
+    start_server(addr="0.0.0.0", port=8080)
+
+def store_ocr_to_clipboard():
+    clipboard.copy(ocr)
+    return
+
+def push_image_to_server():
+    im.save(f"./screenshots/{formatted_date}.png")
+    with open(f"./screenshots/{formatted_date}.png", "rb") as file:
+        encoded = base64.b64encode(file.read())
+        res = requests.post(url="http://localhost:8080/", data=encoded)
+        webbrowser.open(res.text)
+    return
+
+im = ImageGrab.grab()
+ocr = pytesseract.image_to_string(im)
+
+im.thumbnail((1080, 720))
+
+window = tk.Tk()
+noto20 = tk.font.Font(family="Noto Sans", size = 20, weight="normal")
+
+im_widget = ImageTk.PhotoImage(im)
+panel = tk.Label(window, image = im_widget)
+panel.pack(side = "bottom", fill = "both", expand = "yes")
+
+ocr_to_clipboard_btn = tk.Button(window, text="Guardar OCR a clipboard", command=store_ocr_to_clipboard, font=noto20)
+push_image_to_server_btn = tk.Button(window, text="Cargar imagen a servidor", command=push_image_to_server, font=noto20)
+
+ocr_to_clipboard_btn.pack()
+push_image_to_server_btn.pack()
+
+window.mainloop()
+im.save(f"./screenshots/{formatted_date}.png")
\ No newline at end of file
diff --git a/screenshots/.gitkeep b/screenshots/.gitkeep
new file mode 100644
index 0000000..e69de29
diff --git a/server.py b/server.py
new file mode 100644
index 0000000..2a5e65c
--- /dev/null
+++ b/server.py
@@ -0,0 +1,57 @@
+from http.server import BaseHTTPRequestHandler, HTTPServer
+import base64
+from hashlib import sha256
+import os
+
+# FIXME: Hardcoded URL
+base_url = "http://localhost:8080"
+
+class Server(HTTPServer):
+    def __init__(self, address, request_handler):
+        super().__init__(address, request_handler)
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    def __init__(self, request, client_address, server_class):
+        self.server_class = server_class
+        super().__init__(request, client_address, server_class)
+
+    def do_GET(self):
+        parts = self.path.split("?", 1)
+        path = parts[0]
+        if path == "/screenshot" and len(parts) == 2:
+            id = parts[1].replace("id=", "")
+            # FIXME: Potential vulnerability
+            with open(f"./server_screenshots/{id}", "rb") as screenshot:
+                self.send_response(200)
+                self.send_header("Content-Type", "image/png")
+                self.send_header("Content-Length", str(os.fstat(screenshot.fileno()).st_size))
+                self.end_headers()
+                self.wfile.write(screenshot.read())
+        else:
+            self.send_response(400)
+            self.end_headers()
+        
+
+    def do_POST(self):
+        length = int(self.headers['Content-length'])
+        decoded = base64.decodebytes(self.rfile.read(length))
+
+        id = sha256(decoded).hexdigest()
+        url = f"{base_url}/screenshot?id={id}.png"
+
+        f = open(f"./server_screenshots/{id}.png", "wb")
+        f.write(decoded)
+        f.close()
+        
+        self.send_response(200)
+        self.send_header("Content-type", "text/plain")
+        self.send_header("Content-Length", str(len(url)))
+        self.end_headers()
+        self.wfile.write(str(url).encode('utf8'))
+
+def start_server(addr, port, server_class=Server, handler_class=RequestHandler):
+    server_address = (addr, port)
+    http_server = server_class(server_address, handler_class)
+    print(f"Starting server on http://{addr}:{port}")
+    http_server.serve_forever()
diff --git a/server_screenshots/.gitkeep b/server_screenshots/.gitkeep
new file mode 100644
index 0000000..e69de29