58 lines
2 KiB
Python
58 lines
2 KiB
Python
|
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()
|