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()