tests: decryoted AUTH utility

Author: faretek1

.gitignore

@@ -7,7 +7,7 @@ setup.py
 scratchattach/test.py
 scratchattach.code-workspace
 **/.DS_Store
-setup.py
-setup.py
 .env
 tests/manual_tests/**
+tests/util/localauth.toml
+.venv

tests/test_auth.py

@@ -1,8 +1,11 @@
 import sys
 import os
+import util
+
 
 def test_import():
-  sys.path.insert(0, ".")
+    sys.path.insert(0, ".")
 
-  assert "FERNET_KEY" in os.environ
-  assert len(os.environ["FERNET_KEY"]) == 32
+    assert "FERNET_KEY" in os.environ
+    assert len(os.environ["FERNET_KEY"]) == 32
+    assert util.AUTH["test"]["note"]

tests/util/__init__.py

@@ -0,0 +1,3 @@
+# utility methods for testing
+# includes special handlers for authentication etc.
+from .keyhandler import AUTH

tests/util/__main__.py

@@ -0,0 +1,50 @@
+# allows you to generate encrypted data
+
+import argparse
+
+import secrets
+
+try:
+    from .keyhandler import FERNET
+except ImportError:
+    from keyhandler import FERNET
+
+
+def gen_keystr():
+    return secrets.token_urlsafe(32)
+
+
+def main():
+    class Args(argparse.Namespace):
+        command: str
+        content: str
+
+    parser = argparse.ArgumentParser()
+
+    if command := parser.add_subparsers(dest="command"):
+        if encrypt := command.add_parser("e", help="Encrypt content"):
+            encrypt.add_argument("content", nargs="?")
+        if decrypt := command.add_parser("d", help="Decrypt content"):
+            decrypt.add_argument("content", nargs="?")
+        if keygen := command.add_parser("keygen", help="Generate a key. You could set this to $FERNET_KEY if you want"):
+            ...
+
+    args = parser.parse_args(namespace=Args())
+
+    match args.command:
+        case "e":
+            if not args.content:
+                args.content = input("content: ")
+            print(FERNET.encrypt(args.content.encode()).decode())
+
+        case "d":
+            if not args.content:
+                args.content = input("content: ")
+            print(FERNET.decrypt(args.content.encode()).decode())
+
+        case "keygen":
+            print(gen_keystr())
+
+
+if __name__ == "__main__":
+    main()

tests/util/auth.toml

@@ -0,0 +1,7 @@
+# Encrypted authentication data using the FERNET_KEY
+# If you want to set local auth data, make a file 'localauth.toml' in this same directory.
+# any values - that be entire categories - in localauth will overwrite those in auth.toml
+
+[test]
+# this key is just used to test if auth is working
+note = "gAAAAABozoVANH3_fARiCOFvv-O6pHd8DU2kb5vzs7xIxzE3JoWnsFByK5LdoLnnfh0pXpmSpo0KS1j4uwl3XvF-Qm3Sm63Fwy_H224EAdqF_d2hWdVxJkn7HmAJTT445AyrtdFYgrg9"

tests/util/keyhandler.py

@@ -0,0 +1,63 @@
+import warnings
+import os
+import tomllib
+from pathlib import Path
+from typing import TypeVar
+
+from cryptography.fernet import Fernet
+from base64 import urlsafe_b64encode
+
+
+def str_2_key(gen: str) -> bytes:
+    if (length := len(gen)) < 32:
+        warnings.warn(f"Short length {length}")
+        gen = gen.zfill(32)
+    elif length > 32:
+        warnings.warn(f"Long length {length}")
+        gen = gen[:32]
+
+    return urlsafe_b64encode(gen.encode())
+
+
+_fernet_key_raw = os.getenv("FERNET_KEY")
+FERNET_KEY = str_2_key(_fernet_key_raw)
+FERNET = Fernet(FERNET_KEY)
+
+T = TypeVar("T")
+
+
+def _decrypt_val(v: T) -> T:
+    if isinstance(v, str):
+        return FERNET.decrypt(v).decode()
+    if isinstance(v, list):
+        return _decrypt_list(v)
+    if isinstance(v, dict):
+        return _decrypt_dict(v)
+
+    return v
+
+
+def _decrypt_list(data: list) -> list:
+    ret = []
+    for v in data:
+        ret.append(_decrypt_val(v))
+    return ret
+
+
+def _decrypt_dict(data: dict) -> dict:
+    ret = {}
+    for k, v in data.items():
+        ret[k] = _decrypt_val(v)
+
+    return ret
+
+
+__fp__ = Path(__file__).parent
+_auth = _decrypt_dict(tomllib.load(
+    (__fp__ / "auth.toml").open("rb")
+))
+_local_auth = tomllib.load(
+    (__fp__ / "localauth.toml").open("rb")
+)
+
+AUTH = _auth | _local_auth