From c39d5777c5c38430e90798c9b58a34a5e591b5a8 Mon Sep 17 00:00:00 2001
From: 132ikl <132@ikl.sh>
Date: Sat, 6 Nov 2021 23:03:31 -0400
Subject: [PATCH] Add Google Safe Browsing check

---
 liteshort/config.template.yml | 18 +++++++++++++-----
 liteshort/main.py             | 25 +++++++++++++++++++++++--
 liteshort/util.py             | 28 ++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 7 deletions(-)

diff --git a/liteshort/config.template.yml b/liteshort/config.template.yml
index 29a057d..10b3f0e 100644
--- a/liteshort/config.template.yml
+++ b/liteshort/config.template.yml
@@ -1,6 +1,6 @@
 # String: Username to make admin API requests
 # Default: 'admin'
-admin_username: 'admin'
+admin_username: "admin"
 
 # String: Plaintext password to make admin API requests
 # Safe to remove if admin_hashed_password is set
@@ -24,7 +24,7 @@ secret_key: CHANGE_ME
 
 # String: Filename of the URL database without extension
 # Default: 'urls'
-database_name: 'urls'
+database_name: "urls"
 
 # Integer: Length of random short URLs by default
 # Default: 4
@@ -32,7 +32,7 @@ random_length: 4
 
 # String: Allowed URL characters
 # Default: ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_
-allowed_chars: 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_'
+allowed_chars: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
 
 # Amount of time in seconds to spend generating random short URLs until timeout
 # Default: 5
@@ -40,7 +40,7 @@ random_gen_timeout: 5
 
 # String: Name shown on tab while on site and on page header
 # Default: 'liteshort'
-site_name: 'liteshort'
+site_name: "liteshort"
 
 # String: Domain where the shortlinks will be served from. Useful if using the web interface on a subdomain.
 # If not set, it is automatically taken from the URL the shorten request is sent to.
@@ -58,7 +58,7 @@ subdomain:
 # Short URLs cannot be created with this string if set
 # Unset to disable
 # Default: l
-latest: 'l'
+latest: "l"
 
 # Boolean: Show link to project repository on GitHub at bottom right corner of page
 # Default: true
@@ -75,3 +75,11 @@ selflinks: false
 #     - subdomain.blocklisted.net
 # Default: []
 blocklist: []
+
+# String: API key to use Google Safe Browsing to verify links. Leave Unset to not use Safe Browsing.
+# Default: unset
+#safe_browse_key:
+
+# String: URL to replace malicious links (as determined by Safe Browsing) with
+# Default: unset
+#malicious_replace:
diff --git a/liteshort/main.py b/liteshort/main.py
index 2259b98..045d8ed 100644
--- a/liteshort/main.py
+++ b/liteshort/main.py
@@ -12,6 +12,7 @@ from bcrypt import checkpw
 from flask import current_app, g, redirect, render_template, request, url_for
 
 from .config import load_config
+from .util import check_url
 
 logging.basicConfig(level=logging.INFO)
 LOGGER = logging.getLogger(__name__)
@@ -222,6 +223,17 @@ def query_db(query, args=(), one=False, row_factory=sqlite3.Row):
     return (rv[0] if rv else None) if one else rv
 
 
+def safe_check(config, long):
+    if key := config.get("safe_browse_key"):
+        if check_url(key, long):
+            if replace := config.get("malicious_replace"):
+                return replace
+            return None
+        else:
+            return long
+    return long
+
+
 @app.teardown_appcontext
 def close_db(error):
     if hasattr(g, "sqlite_db"):
@@ -327,10 +339,19 @@ def main_post():
                 get_baseUrl() + long_exists,
                 "Error: Failed to return pre-existing random shortlink",
             )
+
+        long = safe_check(current_app.config, request.form["long"])
+        if not long:
+            return response(
+                request,
+                None,
+                "Error: Refusing to create short link for malicious site",
+            )
+
         get_db().cursor().execute(
-            "INSERT INTO urls (long,short) VALUES (?,?)", (request.form["long"], short)
+            "INSERT INTO urls (long,short) VALUES (?,?)", (long, short)
         )
-        set_latest(request.form["long"])
+        set_latest(long)
         get_db().commit()
         return response(request, get_baseUrl() + short, "Error: Failed to generate")
     else:
diff --git a/liteshort/util.py b/liteshort/util.py
index 29a383e..eb13a96 100644
--- a/liteshort/util.py
+++ b/liteshort/util.py
@@ -1,6 +1,34 @@
 from getpass import getpass
+from json import dumps
+from sys import argv
 
 import bcrypt
+import requests
+
+
+def check_url(key, url):
+    out = requests.post(
+        f"https://safebrowsing.googleapis.com/v4/threatMatches:find?key={key}",
+        data=dumps(
+            {
+                "client": {"clientId": "liteshort"},
+                "threatInfo": {
+                    "threatTypes": [
+                        "MALWARE",
+                        "SOCIAL_ENGINEERING",
+                        "UNWANTED_SOFTWARE",
+                        "POTENTIALLY_HARMFUL_APPLICATION",
+                    ],
+                    "platformTypes": ["ANY_PLATFORM"],
+                    "threatEntryTypes": ["URL"],
+                    "threatEntries": [
+                        {"url": url},
+                    ],
+                },
+            }
+        ),
+    )
+    return bool(out.json())
 
 
 def hash_passwd():