From 2d819e670a532ded21f0b74726061b895d5663cf Mon Sep 17 00:00:00 2001
From: Mauro Schiavinato <maurolschiavinato@gmail.com>
Date: Sun, 28 Jul 2024 16:30:48 -0300
Subject: [PATCH] add jinja2 time extension

---
 changedetectionio/jinja_extensions.py         | 67 +++++++++++++++++++
 changedetectionio/safe_jinja.py               |  2 +-
 changedetectionio/tests/test_jinja2.py        | 26 +++++++
 .../tests/unit/test_restock_logic.py          |  2 +-
 requirements.txt                              |  2 +-
 5 files changed, 96 insertions(+), 3 deletions(-)
 create mode 100644 changedetectionio/jinja_extensions.py

diff --git a/changedetectionio/jinja_extensions.py b/changedetectionio/jinja_extensions.py
new file mode 100644
index 00000000..67dd2744
--- /dev/null
+++ b/changedetectionio/jinja_extensions.py
@@ -0,0 +1,67 @@
+"""Jinja2 extensions."""
+import arrow
+from secrets import choice
+
+from jinja2 import nodes
+from jinja2.ext import Extension
+
+class TimeExtension(Extension):
+    """Jinja2 Extension for dates and times."""
+
+    tags = {'now'}
+
+    def __init__(self, environment):
+        """Jinja2 Extension constructor."""
+        super().__init__(environment)
+
+        environment.extend(datetime_format='%Y-%m-%d')
+
+    def _datetime(self, timezone, operator, offset, datetime_format):
+        d = arrow.now(timezone)
+
+        # parse shift params from offset and include operator
+        shift_params = {}
+        for param in offset.split(','):
+            interval, value = param.split('=')
+            shift_params[interval.strip()] = float(operator + value.strip())
+        d = d.shift(**shift_params)
+
+        if datetime_format is None:
+            datetime_format = self.environment.datetime_format
+        return d.strftime(datetime_format)
+
+    def _now(self, timezone, datetime_format):
+        if datetime_format is None:
+            datetime_format = self.environment.datetime_format
+        return arrow.now(timezone).strftime(datetime_format)
+
+    def parse(self, parser):
+        """Parse datetime template and add datetime value."""
+        lineno = next(parser.stream).lineno
+
+        node = parser.parse_expression()
+
+        if parser.stream.skip_if('comma'):
+            datetime_format = parser.parse_expression()
+        else:
+            datetime_format = nodes.Const(None)
+
+        if isinstance(node, nodes.Add):
+            call_method = self.call_method(
+                '_datetime',
+                [node.left, nodes.Const('+'), node.right, datetime_format],
+                lineno=lineno,
+            )
+        elif isinstance(node, nodes.Sub):
+            call_method = self.call_method(
+                '_datetime',
+                [node.left, nodes.Const('-'), node.right, datetime_format],
+                lineno=lineno,
+            )
+        else:
+            call_method = self.call_method(
+                '_now',
+                [node, datetime_format],
+                lineno=lineno,
+            )
+        return nodes.Output([call_method], lineno=lineno)
\ No newline at end of file
diff --git a/changedetectionio/safe_jinja.py b/changedetectionio/safe_jinja.py
index 8a6e1d38..5c34c5ff 100644
--- a/changedetectionio/safe_jinja.py
+++ b/changedetectionio/safe_jinja.py
@@ -12,7 +12,7 @@ JINJA2_MAX_RETURN_PAYLOAD_SIZE = 1024 * int(os.getenv("JINJA2_MAX_RETURN_PAYLOAD
 
 
 def render(template_str, **args: t.Any) -> str:
-    jinja2_env = jinja2.sandbox.ImmutableSandboxedEnvironment(extensions=['jinja2_time.TimeExtension'])
+    jinja2_env = jinja2.sandbox.ImmutableSandboxedEnvironment(extensions=['changedetectionio.jinja_extensions.TimeExtension'])
     output = jinja2_env.from_string(template_str).render(args)
     return output[:JINJA2_MAX_RETURN_PAYLOAD_SIZE]
 
diff --git a/changedetectionio/tests/test_jinja2.py b/changedetectionio/tests/test_jinja2.py
index fba9f227..f6b30266 100644
--- a/changedetectionio/tests/test_jinja2.py
+++ b/changedetectionio/tests/test_jinja2.py
@@ -3,6 +3,7 @@
 import time
 from flask import url_for
 from .util import live_server_setup, wait_for_all_checks
+from ..safe_jinja import render
 
 
 def test_setup(client, live_server, measure_memory_usage):
@@ -56,3 +57,28 @@ def test_jinja2_security_url_query(client, live_server, measure_memory_usage):
     assert b'is invalid and cannot be used' in res.data
     # Some of the spewed output from the subclasses
     assert b'dict_values' not in res.data
+
+def test_add_time(environment):
+    """Verify that added time offset can be parsed."""
+
+    finalRender = render("{% now 'utc' + 'hours=2,seconds=30' %}")
+
+    assert finalRender == "Thu, 10 Dec 2015 01:33:31"
+
+
+def test_substract_time(environment):
+    """Verify that substracted time offset can be parsed."""
+
+    finalRender = render("{% now 'utc' - 'minutes=11' %}")
+
+    assert finalRender == "Wed, 09 Dec 2015 23:22:01"
+
+
+def test_offset_with_format(environment):
+    """Verify that offset works together with datetime format."""
+
+    finalRender = render(
+        "{% now 'utc' - 'days=2,minutes=33,seconds=1', '%d %b %Y %H:%M:%S' %}"
+    )
+
+    assert finalRender == "07 Dec 2015 23:00:00"
\ No newline at end of file
diff --git a/changedetectionio/tests/unit/test_restock_logic.py b/changedetectionio/tests/unit/test_restock_logic.py
index 46fff2c7..db8a63ee 100644
--- a/changedetectionio/tests/unit/test_restock_logic.py
+++ b/changedetectionio/tests/unit/test_restock_logic.py
@@ -6,7 +6,7 @@
 import unittest
 import os
 
-from changedetectionio.processors import restock_diff
+import changedetectionio.processors.restock_diff.processor as restock_diff
 
 # mostly
 class TestDiffBuilder(unittest.TestCase):
diff --git a/requirements.txt b/requirements.txt
index 2e085cf6..884f1c80 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -63,7 +63,7 @@ werkzeug~=3.0
 
 # Templating, so far just in the URLs but in the future can be for the notifications also
 jinja2~=3.1
-jinja2-time
+arrow
 openpyxl
 # https://peps.python.org/pep-0508/#environment-markers
 # https://github.com/dgtlmoon/changedetection.io/pull/1009