From 3d3c95c782549035d0e2710ae65ffdfcdc262136 Mon Sep 17 00:00:00 2001
From: Jess Koester <jess@neocodenetworks.com>
Date: Wed, 20 Mar 2024 17:36:34 -0700
Subject: [PATCH] Add base64 filters and convert thumbor filter to simple
 signing

---
 sssg/__init__.py |  4 ++--
 sssg/filters.py  | 53 ++++++++++++++++++++++++++++++++----------------
 2 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/sssg/__init__.py b/sssg/__init__.py
index 57df0f1..492f787 100644
--- a/sssg/__init__.py
+++ b/sssg/__init__.py
@@ -10,7 +10,7 @@
 import frontmatter
 import urllib.request
 from typing import Generator, Optional
-from .filters import thumbor_signed
+from .filters import add_custom_filters
 
 
 class BuildError(Exception):
@@ -48,7 +48,7 @@ def __init__(self, source_dir: str):
 
         self.jinja = Environment(loader=FileSystemLoader(template_paths))
         self.jinja.filters["markdown"] = lambda text: Markup(self.md.convert(text))
-        self.jinja.filters["thumbor_signed"] = lambda text: thumbor_signed(text)
+        add_custom_filters(self.jinja.filters)
 
     def read_metadata(self, content: str) -> tuple[str, dict[str, str]]:
         '''Attempt to read metadata from a file.'''
diff --git a/sssg/filters.py b/sssg/filters.py
index 2c1135a..626215b 100644
--- a/sssg/filters.py
+++ b/sssg/filters.py
@@ -4,22 +4,41 @@
 import os
 
 
-def thumbor_signed(url: str):
+def b64encode(text: str) -> str:
+    return base64.b64encode(text.encode("utf-8")).decode("utf-8")
+
+
+def b64decode(text: str) -> str:
+    return base64.b64decode(text).decode("utf-8")
+
+
+def urlsafe_b64encode(text: str) -> str:
+    return base64.urlsafe_b64encode(text.encode("utf-8")).decode("utf-8")
+
+
+def urlsafe_b64decode(text: str) -> str:
+    return base64.urlsafe_b64decode(text).decode("utf-8")
+
+
+def sign_sha1(value: str) -> str:
     use_key: str | None = None
-    if "SSSG_THUMBOR_KEY" in os.environ:
-        use_key = os.environ["SSSG_THUMBOR_KEY"]
+    if "SSSG_SIGN_KEY" in os.environ:
+        use_key = os.environ["SSSG_SIGN_KEY"]
     if not use_key:
-        print("SSSG_THUMBOR_KEY not found, falling back to unsafe")
-        return "unsafe/" + url
-
-    return (
-        base64.urlsafe_b64encode(
-            hmac.new(
-                use_key.encode("utf-8"),
-                url.encode("utf-8"),
-                hashlib.sha1,
-            ).digest()
-        ).decode("utf-8")
-        + "/"
-        + url
-    )
+        raise Exception("SSSG_SIGN_KEY not found")
+
+    return base64.urlsafe_b64encode(
+        hmac.new(
+            use_key.encode("utf-8"),
+            value.encode("utf-8"),
+            hashlib.sha1,
+        ).digest()
+    ).decode("utf-8")
+
+
+def add_custom_filters(filters: dict[str, any]):
+    filters["b64decode"] = b64decode
+    filters["b64encode"] = b64encode
+    filters["sign_sha1"] = sign_sha1
+    filters["urlsafe_b64decode"] = urlsafe_b64decode
+    filters["urlsafe_b64encode"] = urlsafe_b64encode