From 5a5d5bb7dd4f87002a0b7fe990333776499ebcf4 Mon Sep 17 00:00:00 2001
From: Hykilpikonna <me@hydev.org>
Date: Mon, 15 Aug 2022 17:11:23 -0400
Subject: [PATCH] [+] safebrowsing filter

---
 api.py | 39 ++++++++++++++++++++++++++++-----------
 1 file changed, 28 insertions(+), 11 deletions(-)

diff --git a/api.py b/api.py
index ccf36c8..ed27b72 100644
--- a/api.py
+++ b/api.py
@@ -1,8 +1,12 @@
 import json
+import os
+import re
 from pathlib import Path
 
+import requests
 import uvicorn
 from fastapi import FastAPI, Body
+from pysafebrowsing import SafeBrowsing
 from starlette.responses import RedirectResponse, HTMLResponse, FileResponse, PlainTextResponse
 
 app = FastAPI()
@@ -19,6 +23,10 @@ last_id_path = data_path / 'last_id.txt'
 chars = 'abcdefghijkmnopqrstuvwxyz'
 base = len(chars)
 
+# URL checks
+re_url = re.compile(r"""^https?://(?:\S+(?::\S*)?@)?(?:(?!(?:10|127)(?:\.\d{1,3}){3})(?!(?:169\.254|192\.168)(?:\.\d{1,3}){2})(?!172\.(?:1[6-9]|2\d|3[0-1])(?:\.\d{1,3}){2})(?:[1-9]\d?|1\d\d|2[01]\d|22[0-3])(?:\.(?:1?\d{1,2}|2[0-4]\d|25[0-5])){2}(?:\.(?:[1-9]\d?|1\d\d|2[0-4]\d|25[0-4]))|(?:(?:[a-z\u00a1-\uffff0-9]-*)*[a-z\u00a1-\uffff0-9]+)(?:\.(?:[a-z\u00a1-\uffff0-9]-*)*[a-z\u00a1-\uffff0-9]+)*(?:\.(?:[a-z\u00a1-\uffff]{2,}))\.?)(?::\d{2,5})?(?:[/?#]\S*)?$""")
+safe_browsing = SafeBrowsing(os.environ['GOOGLE_API_KEY'])
+
 
 def store():
     data_path.mkdir(parents=True, exist_ok=True)
@@ -68,21 +76,30 @@ def get():
 
 @app.put('/')
 def put(name: str | None = None, body: str = Body()):
-    global last_id
+    try:
+        global last_id
 
-    # Generate name
-    while not name:
-        last_id += 1
-        name = encode(last_id)
-        if name in links:
-            name = None
+        # Check valid html
+        assert re_url.match(body)
+        sb = safe_browsing.lookup_url(body)
+        print(sb)
+        assert not sb['malicious'], f'Link is malicious ({",".join(sb["threats"]).lower()})'
 
-    # Put name
-    links[name] = body
-    store()
+        # Generate name
+        while not name:
+            last_id += 1
+            name = encode(last_id)
+            if name in links:
+                name = None
 
-    return PlainTextResponse(f'/{name}')
+        # Put name
+        links[name] = body
+        store()
 
+        return PlainTextResponse(f'/{name}')
+
+    except AssertionError as e:
+        return PlainTextResponse(f'Error: {e}', status_code=400)
 
 if __name__ == '__main__':
     load()