From 5153f5d1128c9106b8e4ba7d2984787acd86877e Mon Sep 17 00:00:00 2001
From: n9k <originalcocaineninja@protonmail.com>
Date: Sat, 5 Mar 2022 09:36:57 +0000
Subject: [PATCH] Add config option for old tripcode algorithm

---
 anonstream/__init__.py         |  2 ++
 anonstream/helpers/tripcode.py | 16 ++++++++++++----
 config.toml                    |  1 +
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/anonstream/__init__.py b/anonstream/__init__.py
index 30bf3f5..5cd2a54 100644
--- a/anonstream/__init__.py
+++ b/anonstream/__init__.py
@@ -25,6 +25,7 @@ def create_app(config_file):
         'lstrip_blocks': True,
     })
     app.config.update({
+        'SECRET_KEY_STRING': config['secret_key'],
         'SECRET_KEY': config['secret_key'].encode(),
         'AUTH_USERNAME': config['auth']['username'],
         'AUTH_PWHASH': auth_pwhash,
@@ -56,6 +57,7 @@ def create_app(config_file):
         'CHAT_NAME_MAX_LENGTH': config['chat']['max_name_length'],
         'CHAT_NAME_MIN_CONTRAST': config['chat']['min_name_contrast'],
         'CHAT_BACKGROUND_COLOUR': color_to_colour(config['chat']['background_color']),
+        'CHAT_LEGACY_TRIPCODE_ALGORITHM': config['chat']['legacy_tripcode_algorithm'],
         'FLOOD_DURATION': config['flood']['duration'],
         'FLOOD_THRESHOLD': config['flood']['threshold'],
         'CAPTCHA_LIFETIME': config['captcha']['lifetime'],
diff --git a/anonstream/helpers/tripcode.py b/anonstream/helpers/tripcode.py
index 0b0489a..576fa67 100644
--- a/anonstream/helpers/tripcode.py
+++ b/anonstream/helpers/tripcode.py
@@ -11,19 +11,27 @@ CONFIG = current_app.config
 def _generate_tripcode_digest_legacy(password):
     hexdigest, _ = werkzeug.security._hash_internal(
         'pbkdf2:sha256:150000',
-        CONFIG['SECRET_KEY'],
+        CONFIG['SECRET_KEY_STRING'],
         password,
     )
     digest = bytes.fromhex(hexdigest)
     return base64.b64encode(digest)[:8].decode()
 
-def generate_tripcode_digest(password):
+def _generate_tripcode_digest(password):
     parts = CONFIG['SECRET_KEY'] + b'tripcode\0' + password.encode()
     digest = hashlib.sha256(parts).digest()
     return base64.b64encode(digest)[:8].decode()
 
-def generate_tripcode(password, generate_digest=generate_tripcode_digest):
-    digest = generate_digest(password)
+def generate_tripcode_digest(password):
+    algorithm = (
+        _generate_tripcode_digest_legacy
+        if CONFIG['CHAT_LEGACY_TRIPCODE_ALGORITHM'] else
+        _generate_tripcode_digest
+    )
+    return algorithm(password)
+
+def generate_tripcode(password):
+    digest = generate_tripcode_digest(password)
     background_colour = generate_colour(
         seed='tripcode-background\0' + digest,
         bg=CONFIG['CHAT_BACKGROUND_COLOUR'],
diff --git a/config.toml b/config.toml
index 6632e90..c0043f1 100644
--- a/config.toml
+++ b/config.toml
@@ -45,6 +45,7 @@ max_comment_length = 512
 max_name_length = 24
 min_name_contrast = 3.0
 background_color = "#232327"
+legacy_tripcode_algorithm = false
 
 [flood]
 duration = 20.0