[+] safebrowsing filter

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,8 +76,15 @@ def get():
 
 @app.put('/')
 def put(name: str | None = None, body: str = Body()):
+    try:
         global last_id
 
+        # 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()})'
+
         # Generate name
         while not name:
             last_id += 1
@@ -83,6 +98,8 @@ def put(name: str | None = None, body: str = Body()):
 
         return PlainTextResponse(f'/{name}')
 
+    except AssertionError as e:
+        return PlainTextResponse(f'Error: {e}', status_code=400)
 
 if __name__ == '__main__':
     load()