From eea62d86852d4c97a798edf2f7faabfdedea0cdc Mon Sep 17 00:00:00 2001
From: "Maxence G. de Montauzan" <maxence@gdemontauzan.fr>
Date: Sun, 12 Nov 2023 00:58:52 +0100
Subject: [PATCH] Use a configurable TOTP protection

---
 conf.ini         |  4 ++++
 notifier.py      | 55 ++++++++++++++++++++++++++++++++++++++----------
 requirements.txt |  1 +
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/conf.ini b/conf.ini
index c0867e8..354a048 100644
--- a/conf.ini
+++ b/conf.ini
@@ -5,6 +5,10 @@ smtp_server = mailpit
 sender_email = sender@host.eu
 receiver_email = receiver@anotherhost.eu
 
+[api.totp]
+enabled = true
+key = mysuperkey
+
 [projects]
 projects = [
 	"borgbackup/borg",
diff --git a/notifier.py b/notifier.py
index f4154dc..bd7453d 100644
--- a/notifier.py
+++ b/notifier.py
@@ -7,9 +7,14 @@ import sys
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
 
+import pyotp
 import requests
-from fastapi import FastAPI
+from fastapi import FastAPI, status
+from fastapi.responses import JSONResponse
 
+#
+# CONF PART
+#
 SCRIPT_DIR = os.path.dirname(__file__)
 CONF_FILE = os.path.join(SCRIPT_DIR, "conf.ini")
 TEMPLATE_FILE = os.path.join(SCRIPT_DIR, "template.html")
@@ -31,10 +36,36 @@ class EnvInjection(configparser.Interpolation):
         return env_value if env_value else file_value
 
 
+def load_conf(conf_file=CONF_FILE) -> configparser.ConfigParser:
+    parser = configparser.ConfigParser(
+        default_section="config", interpolation=EnvInjection()
+    )
+    parser.read(conf_file)
+    return parser
+
+
+def load_totp():
+    parser = load_conf()
+
+    if parser.getboolean("api.totp", "enabled", fallback=True):
+        totp = pyotp.TOTP(parser.get("api.totp", "key", fallback=pyotp.random_base32()))
+        # TODO Catch totp.now() error
+        print(
+            f"""TOTP enabled.
+    Information: {totp.provisioning_uri()}
+    TOTP check: {totp.now()}"""
+        )
+        return totp
+    else:
+        print("WARNING: Api is open without security")
+        return type("totp", (object,), {"verify": lambda str: True})
+
+
 #
 # API PARTS
 #
 app = FastAPI()
+TOTP = load_totp()
 
 
 @app.get("/subscriptions")
@@ -45,13 +76,23 @@ def get_projects():
 
 
 @app.put("/subscriptions")
-def put_projects(project: str, author: str | None = None):
+def put_projects(
+    project: str, author: str | None = None, credentials: str | None = None
+):
+    if not TOTP.verify(credentials):
+        return JSONResponse(
+            status_code=status.HTTP_401_UNAUTHORIZED, content="Unauthorized"
+        )
     # TODO Check if project really exist?
     parser = load_conf()
     projects = json.loads(parser.get("projects", "projects"))
 
     if author:
         project = f"{author}/{project}"
+
+    if project in projects:
+        return project
+
     projects.append(project)
 
     # TODO Watch a converter for list: https://stackoverflow.com/a/53274707/1346391
@@ -60,15 +101,7 @@ def put_projects(project: str, author: str | None = None):
     with open("conf.ini", "w", encoding="utf-8") as configfile:
         parser.write(configfile)
 
-    return project
-
-
-def load_conf(conf_file=CONF_FILE) -> configparser.ConfigParser:
-    parser = configparser.ConfigParser(
-        default_section="config", interpolation=EnvInjection()
-    )
-    parser.read(conf_file)
-    return parser
+    return JSONResponse(status_code=status.HTTP_201_CREATED, content=project)
 
 
 #
diff --git a/requirements.txt b/requirements.txt
index fc62a3b..227cc3e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 black
 fastapi
 pre-commit
+pyotp
 requests
 uvicorn[standard]