1 files changed, 30 insertions, 19 deletions
diff --git a/samloader/auth.py b/samloader/auth.py
index 5b6a1d5..3758bf8 100644
--- a/samloader/auth.py
+++ b/samloader/auth.py
@@ -1,38 +1,49 @@
 # SPDX-License-Identifier: GPL-3.0+
 # Copyright (C) 2020 nlscc
 
-# FUS authentication functions (decrypting nonce, calculating auth token)
+""" FUS authentication functions (decrypting nonce, calculating auth token) """
 
-from Cryptodome.Cipher import AES
 import base64
-import requests
+from Cryptodome.Cipher import AES
 
+# Constant key input values.
 KEY_1 = "hqzdurufm2c8mf6bsjezu1qgveouv7c7"
 KEY_2 = "w13r4cvf4hctaujv"
 
-unpad = lambda d: d[:-d[-1]]
-pad = lambda d: d + bytes([16 - (len(d) % 16)]) * (16 - (len(d) % 16))
+# PKCS#7 padding functions.
+pkcs_unpad = lambda d: d[:-d[-1]]
+pkcs_pad = lambda d: d + bytes([16 - (len(d) % 16)]) * (16 - (len(d) % 16))
 
-def aes_encrypt(inp, key):
-    cipher = AES.new(key, AES.MODE_CBC, key[:16])
-    return cipher.encrypt(pad(inp))
+def aes_encrypt(inp: bytes, key: bytes) -> bytes:
+    """ Perform an AES-CBC encryption. Encrypts /inp/ with key /key/. """
+    enc_iv = key[:16] # IV is first 16 bytes of key
+    cipher = AES.new(key, AES.MODE_CBC, enc_iv)
+    return cipher.encrypt(pkcs_pad(inp))
 
-def aes_decrypt(inp, key):
-    cipher = AES.new(key, AES.MODE_CBC, key[:16])
-    return unpad(cipher.decrypt(inp))
+def aes_decrypt(inp: bytes, key: bytes) -> bytes:
+    """ Perform an AES-CBC decryption. Decrypts /inp/ with key /key/. """
+    enc_iv = key[:16]
+    cipher = AES.new(key, AES.MODE_CBC, enc_iv)
+    return pkcs_unpad(cipher.decrypt(inp))
 
-def getfkey(inp):
+def derive_key(nonce: str) -> bytes:
+    """ Calculate the AES key from the FUS input nonce. """
     key = ""
+    # First 16 bytes are offsets into KEY_1
     for i in range(16):
-        key += KEY_1[inp[i]]
+        key += KEY_1[ord(nonce[i]) % 16]
+    # Last 16 bytes are static
     key += KEY_2
     return key.encode()
 
-def getauth(nonce):
-    keydata = [ord(c) % 16 for c in nonce]
-    fkey = getfkey(keydata)
-    return base64.b64encode(aes_encrypt(nonce.encode(), fkey)).decode()
+def getauth(nonce: str) -> str:
+    """ Calculate the response token from a given nonce. """
+    nkey = derive_key(nonce)
+    auth_data = aes_encrypt(nonce.encode(), nkey)
+    return base64.b64encode(auth_data).decode()
 
-def decryptnonce(inp):
-    nonce = aes_decrypt(base64.b64decode(inp), KEY_1.encode()).decode()
+def decryptnonce(inp: str) -> str:
+    """ Decrypt the nonce returned by the server. """
+    inp_data = base64.b64decode(inp)
+    nonce = aes_decrypt(inp_data, KEY_1.encode()).decode()
     return nonce