From 6d20d53406a2420ec586d863987000138191ee42 Mon Sep 17 00:00:00 2001
From: q <q@neveremailme.fuckoff>
Date: Sat, 15 Jun 2024 20:23:24 -0400
Subject: [PATCH] Adding blacklist version

---
 README.md                |  2 +-
 blacklist.txt            |  1 +
 custommedia-blacklist.py | 75 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 blacklist.txt
 create mode 100755 custommedia-blacklist.py

diff --git a/README.md b/README.md
index 843e13b..fc907d8 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # CustomMedia
 
-A fork of [Plan9's CustomMedia project](https://gitea.plan9.rocks/cat/CustomMedia) that adds a whitelist-based version for choosing specific homeservers to pull media from directly. Conserve your matrix homeserver's precious storage space and bandwidth with this one simple webserver.
+A fork of [Plan9's CustomMedia project](https://gitea.plan9.rocks/cat/CustomMedia) that adds whitelist and blacklist versions for choosing specific homeservers to pull media from directly. Conserve your matrix homeserver's precious storage space and bandwidth with this one simple webserver.
 
 # Requirements
 - Python 3
diff --git a/blacklist.txt b/blacklist.txt
new file mode 100644
index 0000000..de54ac6
--- /dev/null
+++ b/blacklist.txt
@@ -0,0 +1 @@
+example.com
diff --git a/custommedia-blacklist.py b/custommedia-blacklist.py
new file mode 100755
index 0000000..c720597
--- /dev/null
+++ b/custommedia-blacklist.py
@@ -0,0 +1,75 @@
+#!/usr/bin/python3
+import urllib.request
+import json
+
+headers = {
+    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36'
+}
+
+class MyServer:
+    def __init__(self, environ, start_response):
+        self.environ = environ
+        self.start_response = start_response
+        self.mapping_file = 'mapping.json'
+        self.mappings = self.load_mappings()
+        self.default_hs = "https://matrix-client.matrix.org"
+        self.blacklist = set(line.strip() for line in open('blacklist.txt'))
+
+    def __iter__(self):
+        hs = self.environ['PATH_INFO'].split('/')[5]
+
+        if not hs in self.blacklist:
+            hsu = self.mappings.get(hs)
+
+            if not hsu:
+                wellknown = "https://" + hs + "/.well-known/matrix/client"
+                req = urllib.request.Request(wellknown, headers=headers)
+                try:
+                    with urllib.request.urlopen(req, timeout=2) as cont:
+                        hsu = json.load(cont)
+                        hsu = hsu["m.homeserver"]["base_url"].rstrip('/')
+                        self.mappings[hs] = hsu
+                        self.save_mappings()
+                except Exception as e: # I don't care to fix this properly, not my problem
+                    hsu = self.default_hs
+        else:
+            hsu = self.default_hs
+
+        hsp = hsu + self.environ['PATH_INFO'] + '?' + self.environ['QUERY_STRING']
+        self.start_response('301 Moved Permanently', [('Location', hsp)])
+        return iter([])
+
+    def load_mappings(self):
+        try:
+            with open(self.mapping_file, 'r') as f:
+                return json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            return {}
+
+    def save_mappings(self):
+        with open(self.mapping_file, 'w') as f:
+            json.dump(self.mappings, f)
+
+if __name__ == "__main__":
+    from gunicorn.app.base import BaseApplication
+
+    class GunicornServer(BaseApplication):
+        def __init__(self, app, options=None):
+            self.options = options or {}
+            self.application = app
+            super().__init__()
+
+        def load_config(self):
+            for key, value in self.options.items():
+                self.cfg.set(key, value)
+
+        def load(self):
+            return self.application
+
+    options = {
+        'bind': 'localhost:9999',
+        'workers': 48,  # Adjust the number of workers based on your system's resources - ChatGPT
+    }
+
+    server = GunicornServer(MyServer, options)
+    server.run()