klemek/stapler
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 10 Wiki Activity

feat: X-Proxy

Browse Source
This commit is contained in:
Klemek
2026-04-20 14:41:27 +02:00
committed by klemek
parent fb70638330
commit 33cfd350a5
8 changed files with 609 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/handlers.py
+73 -20
View File
@@ -47,7 +47,7 @@ class BaseHandler(abc.ABC, http.server.BaseHTTPRequestHandler):
) -> None:
shortmsg, longmsg = self.responses[code]
if message is None:
message = shortmsg # pragma: no cover
message = shortmsg
if explain is None:
explain = longmsg
if "text/" in self._get_header("Accept"):
@@ -104,7 +104,7 @@ class BaseHandler(abc.ABC, http.server.BaseHTTPRequestHandler):
encoded: bytes = body.encode()
self.out_size = len(encoded)
self.send_response(code, message)
self.send_header("Content-type", f"{content_type}; charset=UTF-8")
self.send_header("Content-Type", f"{content_type}; charset=UTF-8")
self.send_header("Content-Length", str(len(encoded)))
for header, value in headers.items():
self.send_header(header, value) # pragma: no cover
@@ -134,14 +134,14 @@ class BaseHandler(abc.ABC, http.server.BaseHTTPRequestHandler):
)
def send_proxy(self, url: str) -> None:
body: bytes | None = None
if self.in_size > 0:
body = self.rfile.read(self.in_size)
headers = dict(self.headers)
headers["Host"] = urllib.parse.urlparse(url).netloc
headers["X-Forwarded-For"] = self.client_address[0]
headers["X-Real-IP"] = self.client_address[0]
try:
body: bytes | None = None
if self.in_size > 0:
body = self.rfile.read(self.in_size)
response: requests.Response = requests.request(
self.command, url, data=body, headers=headers, timeout=240
)
@@ -220,6 +220,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
TOKEN_HEADER = "X-Token" # noqa: S105
HOST_HEADER = "X-Host"
REDIRECT_HEADER = "X-Redirect"
PROXY_HEADER = "X-Proxy"
@typing.override
def __init__(
@@ -241,6 +242,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
self.__token: str | None = None
self.__target_host: str | None = None
self.__target_redirect: str | None = None
self.__target_proxy: str | None = None
super().__init__(*args, directory=params.data_dir, **kwargs, params=params) # ty:ignore[unknown-argument]
@property
@@ -249,6 +251,10 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
self.__token = self._get_header(self.TOKEN_HEADER)
return self.__token
@property
def has_token(self) -> bool:
return len(self.token) > 0
@property
def target_host(self) -> str:
if self.__target_host is None:
@@ -269,28 +275,35 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
def has_target_redirect(self) -> bool:
return len(self.target_redirect) > 0
@property
def target_proxy(self) -> str:
if self.__target_proxy is None:
self.__target_proxy = self._get_header(self.PROXY_HEADER).lower()
return self.__target_proxy
@property
def has_target_proxy(self) -> bool:
return len(self.target_proxy) > 0
@typing.override
def do_HEAD(self) -> None:
self._pre_log_request()
if (
page := self.__get_page(self.path)
) is not None and page.redirect is not None:
return self.send_redirect(page.redirect)
return super().do_HEAD()
if not self._proxy_or_redirect():
super().do_HEAD()
@typing.override
def do_GET(self) -> None:
self._pre_log_request()
if self._proxy_or_redirect():
return None
if self.path == "/" and self.host == self.default_host:
return self.send_basic_body(self.server_signature())
if (
page := self.__get_page(self.path)
) is not None and page.redirect is not None:
return self.send_redirect(page.redirect)
return super().do_GET()
def do_PUT(self) -> None:
self._pre_log_request()
if self._proxy_or_redirect():
return None
if (path := self.__check_update_request()) is None:
return None
if self.has_target_host and not self.__valid_host(self.target_host):
@@ -303,8 +316,15 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
and page.path != path
):
return self.send_error(http.HTTPStatus.FORBIDDEN, "Host already taken")
if self.has_target_proxy and self.has_target_redirect:
return self.send_error(
http.HTTPStatus.BAD_REQUEST,
f"Cannot use {self.PROXY_HEADER} with {self.REDIRECT_HEADER}",
)
if self.has_target_redirect:
self._update_redirect(path)
elif self.has_target_proxy:
self._update_proxy(path)
else:
self._update_extract(path)
if self.has_target_host and self.cert_manager.create_or_update(
@@ -321,21 +341,26 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
def do_DELETE(self) -> None:
self._pre_log_request()
if self._proxy_or_redirect():
return None
if (path := self.__check_update_request()) is None:
return None
return self._update_remove(path)
def do_CONNECT(self) -> None:
self._pre_log_request()
self.send_error(http.HTTPStatus.METHOD_NOT_ALLOWED)
if not self._proxy_or_redirect():
self.send_error(http.HTTPStatus.METHOD_NOT_ALLOWED)
def do_OPTIONS(self) -> None:
self._pre_log_request()
self.send_error(http.HTTPStatus.METHOD_NOT_ALLOWED)
if not self._proxy_or_redirect():
self.send_error(http.HTTPStatus.METHOD_NOT_ALLOWED)
def do_TRACE(self) -> None:
self._pre_log_request()
self.send_error(http.HTTPStatus.METHOD_NOT_ALLOWED)
if not self._proxy_or_redirect():
self.send_error(http.HTTPStatus.METHOD_NOT_ALLOWED)
def _update_extract(self, path: str) -> None:
if self.in_size == 0:
@@ -366,10 +391,22 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
http.HTTPStatus.BAD_REQUEST,
f"No content must be sent with {self.REDIRECT_HEADER}",
)
self.data_dir.empty(path)
self.registry.add(path)
self.token_manager.set_token(path, self.token)
self.registry.set_redirect(path, self.target_redirect)
self.token_manager.set_token(path, self.token)
self.send_status_only(
http.HTTPStatus.CREATED,
f"Resource /{path}/ updated",
)
return None
def _update_proxy(self, path: str) -> None:
if self.in_size > 0:
return self.send_error(
http.HTTPStatus.BAD_REQUEST,
f"No content must be sent with {self.PROXY_HEADER}",
)
self.registry.set_proxy(path, self.target_proxy)
self.token_manager.set_token(path, self.token)
self.send_status_only(
http.HTTPStatus.CREATED,
f"Resource /{path}/ updated",
@@ -391,6 +428,22 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
self.registry.remove(path)
return None
def _proxy_or_redirect(self) -> bool:
if self.has_token:
return False
if (page := self.__get_page(self.path)) is None:
return False
if page.redirect is not None:
self.send_redirect(page.redirect)
return True
if page.proxy is not None:
if self.host == self.default_host:
self.send_proxy(page.proxy + self.path.removeprefix(f"/{page.path}"))
else:
self.send_proxy(page.proxy + self.path)
return True
return False
@typing.override
def list_directory(self, *_: typing.Any, **__: typing.Any) -> None:
"""Disable default directory listing."""
src/page.py
+3
View File
@@ -8,6 +8,7 @@ class Page:
host: str | None = None
token_hash: str | None = None
redirect: str | None = None
proxy: str | None = None
def __repr__(self) -> str:
out = f"/{self.path}/"
@@ -15,6 +16,8 @@ class Page:
out += f" [{self.host}]"
if self.redirect is not None:
out += f" (redirect: {self.redirect})"
elif self.proxy is not None:
out += f" (proxy: {self.proxy})"
elif not self.with_index:
out += " (no index)"
return out
src/registry.py
+16 -2
View File
@@ -18,6 +18,7 @@ class Registry:
HOST_FILE = ".host"
TOKEN_FILE = ".token" # noqa: S105
REDIRECT_FILE = ".redirect"
PROXY_FILE = ".proxy"
def __init__(self, params: Parameters) -> None:
self.logger: logging.Logger = logging.getLogger(self.__class__.__name__)
@@ -39,6 +40,7 @@ class Registry:
self.data_dir.get_file(path, self.HOST_FILE),
self.data_dir.get_file(path, self.TOKEN_FILE),
self.data_dir.get_file(path, self.REDIRECT_FILE),
self.data_dir.get_file(path, self.PROXY_FILE),
)
self.logger.info("Updated %s", self.pages[path])
@@ -52,14 +54,26 @@ class Registry:
if path in self.pages and self.pages[path].token_hash != token_hash:
self.data_dir.set_file(path, self.TOKEN_FILE, token_hash, 0o600)
self.pages[path].token_hash = token_hash
self.logger.debug("Updated %s", self.pages[path])
self.logger.debug("Updated %s (token)", self.pages[path])
def set_redirect(self, path: str, redirect: str) -> None:
if path in self.pages and self.pages[path].redirect != redirect:
if path not in self.pages or self.pages[path].redirect != redirect:
self.data_dir.empty(path)
self.data_dir.set_file(path, self.REDIRECT_FILE, redirect)
if path not in self.pages:
self.pages[path] = Page(path)
self.pages[path].redirect = redirect
self.logger.debug("Updated %s", self.pages[path])
def set_proxy(self, path: str, proxy: str) -> None:
if path not in self.pages or self.pages[path].proxy != proxy:
self.data_dir.empty(path)
self.data_dir.set_file(path, self.PROXY_FILE, proxy)
if path not in self.pages:
self.pages[path] = Page(path)
self.pages[path].proxy = proxy
self.logger.debug("Updated %s", self.pages[path])
def remove(self, path: str) -> None:
if path in self.pages:
page = self.pages[path]