#!/bin/python3

import os
from http.server import HTTPServer, SimpleHTTPRequestHandler

class Custom404Handler(SimpleHTTPRequestHandler):
    def send_error(self, code, message=None, explain=None):
        if code == 404:
            # Serve custom 404.html if it exists
            path = os.path.join(self.directory, "404.html")
            if os.path.exists(path):
                self.send_response(404)
                self.send_header("Content-type", "text/html")
                self.end_headers()
                with open(path, "rb") as f:
                    self.wfile.write(f.read())
                return
        super().send_error(code, message, explain)

    def do_GET(self):
        # Try to serve file, otherwise call send_error(404)
        if self.path == "/":
            return super().do_GET()
        else:
            # Remove query string, fragments
            file_path = self.translate_path(self.path.split('?',1)[0].split('#',1)[0])
            if os.path.exists(file_path) and not os.path.isdir(file_path):
                return super().do_GET()
            else:
                self.send_error(404)

if __name__ == '__main__':
    import sys
    import os

    # Set working directory
    directory = os.getcwd()
    if len(sys.argv) > 1:
        directory = sys.argv[1]
    handler = Custom404Handler
    handler.directory = directory

    server = HTTPServer(("0.0.0.0", 8080), handler)
    print(f"Serving on http://localhost:8080 from {directory}")
    server.serve_forever()