From df0a637d2b5da4a2fc6dae0c7f07bcd7f50e4828 Mon Sep 17 00:00:00 2001
From: Kevin Morris <kevr@0cost.org>
Date: Thu, 28 Jan 2021 16:52:56 -0800
Subject: [PATCH] add aurweb.captcha, a CAPTCHA utility module

This CAPTCHA workflow is the same workflow used by our current
PHP implementation of account registration.

Signed-off-by: Kevin Morris <kevr@0cost.org>
---
 aurweb/captcha.py    | 54 +++++++++++++++++++++++++++++++++++++++
 aurweb/templates.py  |  6 ++++-
 test/test_captcha.py | 60 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 119 insertions(+), 1 deletion(-)
 create mode 100644 aurweb/captcha.py
 create mode 100644 test/test_captcha.py

diff --git a/aurweb/captcha.py b/aurweb/captcha.py
new file mode 100644
index 00000000..5475d85f
--- /dev/null
+++ b/aurweb/captcha.py
@@ -0,0 +1,54 @@
+""" This module consists of aurweb's CAPTCHA utility functions and filters. """
+import hashlib
+
+import jinja2
+
+from aurweb.db import query
+from aurweb.models.user import User
+
+
+def get_captcha_salts():
+    """ Produce salts based on the current user count. """
+    count = query(User).count()
+    salts = []
+    for i in range(0, 6):
+        salts.append(f"aurweb-{count - i}")
+    return salts
+
+
+def get_captcha_token(salt):
+    """ Produce a token for the CAPTCHA salt. """
+    return hashlib.md5(salt.encode()).hexdigest()[:3]
+
+
+def get_captcha_challenge(salt):
+    """ Get a CAPTCHA challenge string (shell command) for a salt. """
+    token = get_captcha_token(salt)
+    return f"LC_ALL=C pacman -V|sed -r 's#[0-9]+#{token}#g'|md5sum|cut -c1-6"
+
+
+def get_captcha_answer(token):
+    """ Compute the answer via md5 of the real template text, return the
+    first six digits of the hexadecimal hash. """
+    text = r"""
+ .--.                  Pacman v%s.%s.%s - libalpm v%s.%s.%s
+/ _.-' .-.  .-.  .-.   Copyright (C) %s-%s Pacman Development Team
+\  '-. '-'  '-'  '-'   Copyright (C) %s-%s Judd Vinet
+ '--'
+                       This program may be freely redistributed under
+                       the terms of the GNU General Public License.
+""" % tuple([token] * 10)
+    return hashlib.md5((text + "\n").encode()).hexdigest()[:6]
+
+
+@jinja2.contextfilter
+def captcha_salt_filter(context):
+    """ Returns the most recent CAPTCHA salt in the list of salts. """
+    salts = get_captcha_salts()
+    return salts[0]
+
+
+@jinja2.contextfilter
+def captcha_cmdline_filter(context, salt):
+    """ Returns a CAPTCHA challenge for a given salt. """
+    return get_captcha_challenge(salt)
diff --git a/aurweb/templates.py b/aurweb/templates.py
index 4ea74a62..d548e92b 100644
--- a/aurweb/templates.py
+++ b/aurweb/templates.py
@@ -12,7 +12,7 @@ from fastapi.responses import HTMLResponse
 
 import aurweb.config
 
-from aurweb import l10n, time
+from aurweb import captcha, l10n, time
 
 # Prepare jinja2 objects.
 loader = jinja2.FileSystemLoader(os.path.join(
@@ -23,6 +23,10 @@ env = jinja2.Environment(loader=loader, autoescape=True,
 # Add tr translation filter.
 env.filters["tr"] = l10n.tr
 
+# Add captcha filters.
+env.filters["captcha_salt"] = captcha.captcha_salt_filter
+env.filters["captcha_cmdline"] = captcha.captcha_cmdline_filter
+
 
 def make_context(request: Request, title: str, next: str = None):
     """ Create a context for a jinja2 TemplateResponse. """
diff --git a/test/test_captcha.py b/test/test_captcha.py
new file mode 100644
index 00000000..ec19dee9
--- /dev/null
+++ b/test/test_captcha.py
@@ -0,0 +1,60 @@
+import hashlib
+
+from aurweb import captcha
+
+
+def test_captcha_salts():
+    """ Make sure we can get some captcha salts. """
+    salts = captcha.get_captcha_salts()
+    assert len(salts) == 6
+
+
+def test_captcha_token():
+    """ Make sure getting a captcha salt's token matches up against
+    the first three digits of the md5 hash of the salt. """
+    salts = captcha.get_captcha_salts()
+    salt = salts[0]
+
+    token1 = captcha.get_captcha_token(salt)
+    token2 = hashlib.md5(salt.encode()).hexdigest()[:3]
+
+    assert token1 == token2
+
+
+def test_captcha_challenge_answer():
+    """ Make sure that executing the captcha challenge via shell
+    produces the correct result by comparing it against a straight
+    up token conversion. """
+    salts = captcha.get_captcha_salts()
+    salt = salts[0]
+
+    challenge = captcha.get_captcha_challenge(salt)
+
+    token = captcha.get_captcha_token(salt)
+    challenge2 = f"LC_ALL=C pacman -V|sed -r 's#[0-9]+#{token}#g'|md5sum|cut -c1-6"
+
+    assert challenge == challenge2
+
+
+def test_captcha_salt_filter():
+    """ Make sure captcha_salt_filter returns the first salt from
+    get_captcha_salts().
+
+    Example usage:
+        <input type="hidden" name="captcha_salt" value="{{ captcha_salt }}">
+    """
+    salt = captcha.captcha_salt_filter(None)
+    assert salt == captcha.get_captcha_salts()[0]
+
+
+def test_captcha_cmdline_filter():
+    """ Make sure that the captcha_cmdline filter gives us the
+    same challenge that get_captcha_challenge does.
+
+    Example usage:
+        <code>{{ captcha_salt | captcha_cmdline }}</code>
+    """
+    salt = captcha.captcha_salt_filter(None)
+    display1 = captcha.captcha_cmdline_filter(None, salt)
+    display2 = captcha.get_captcha_challenge(salt)
+    assert display1 == display2