From 37379e4255e13a1b94ef691971e2eee80439e976 Mon Sep 17 00:00:00 2001
From: Nathan Woodburn <github@nathan.woodburn.au>
Date: Fri, 1 Mar 2024 17:26:32 +1100
Subject: [PATCH] feat: Add more replace options

---
 .env.example | 3 ++-
 README.md    | 2 +-
 main.py      | 4 ++++
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/.env.example b/.env.example
index 7430307..0491eef 100644
--- a/.env.example
+++ b/.env.example
@@ -1,3 +1,4 @@
 PROXY=http://nathanwoodburn:5000/
 TLD=woodburn
-RESTRICTED=["admin"]
\ No newline at end of file
+RESTRICTED=["admin"]
+REPLACE=["https://nathan.woodburn.au/"]
\ No newline at end of file
diff --git a/README.md b/README.md
index 764d037..6987a79 100644
--- a/README.md
+++ b/README.md
@@ -4,4 +4,4 @@ Environment variables:
 - `PROXY` - The proxy to use for restricted TLDs. Example: `http://localhost:3128`
 - `TLD` - The TLD to restrict. Example: `g`
 - `RESTRICTED` - List of restricted paths. Example: `["path1", "path2"]` will require the user to have a .g domain to access `path1/*`, `path2/*`
-
+- `REPLACE` - List of urls to replace with the host to fix redirects and links
\ No newline at end of file
diff --git a/main.py b/main.py
index 4c12cf6..23dc15f 100644
--- a/main.py
+++ b/main.py
@@ -14,6 +14,8 @@ RESTRICTED = os.getenv('RESTRICTED')
 RESTRICTED = json.loads(RESTRICTED)
 RESTRICTED = [f'{i.lower()}/' for i in RESTRICTED]
 TLD = os.getenv('TLD')
+REPLACEMENT = os.getenv('REPLACE')
+REPLACEMENT = json.loads(REPLACEMENT)   
 
 # Load cookies
 cookies = []
@@ -118,6 +120,8 @@ def catch_all(path):
     if 'text/html' in res.headers['Content-Type']:
         content = res.content.decode('utf-8')
         content = content.replace(URL, request.host_url)
+        for i in REPLACEMENT:
+            content = content.replace(i, request.host_url)
         response = make_response(content, res.status_code, headers)
         return response