split out constant, improve logic some, add tests for channel_hash and generate_channel_hash

meshtastic/util.py

@@ -40,6 +40,8 @@ whitelistVids = dict.fromkeys([0x239a, 0x303a])
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_KEY = base64.b64decode("1PG7OiApB1nwvP+rz05pAQ==".encode("utf-8"))
+
 def quoteBooleans(a_string: str) -> str:
     """Quote booleans
     given a string that contains ": true", replace with ": 'true'" (or false)
@@ -372,24 +374,21 @@ def channel_hash(data: bytes) -> int:
         result ^= char
     return result
 
-def generate_channel_hash(name, key) -> int:
+def generate_channel_hash(name: Union[str, bytes], key: Union[str, bytes]) -> int:
     """generate the channel number by hashing the channel name and psk (accepts str or bytes for both)"""
     # Handle key as str or bytes
-    if isinstance(key, bytes):
-        key = key.decode("utf-8")
-    if key == "AQ==":
-        key = "1PG7OiApB1nwvP+rz05pAQ=="
-    replaced_key = key.replace("-", "+").replace("_", "/")
-    key_bytes = base64.b64decode(replaced_key.encode("utf-8"))
+    if isinstance(key, str):
+        key = base64.b64decode(key.replace("-", "+").replace("_", "/").encode("utf-8"))
+
+    if len(key) == 1:
+        key = DEFAULT_KEY[:-1] + key
 
     # Handle name as str or bytes
-    if isinstance(name, bytes):
-        name_bytes = name
-    else:
-        name_bytes = name.encode("utf-8")
+    if isinstance(name, str):
+        name = name.encode("utf-8")
 
-    h_name = channel_hash(name_bytes)
-    h_key = channel_hash(key_bytes)
+    h_name = channel_hash(name)
+    h_key = channel_hash(key)
     result: int = h_name ^ h_key
     return result