Update app.py

app.py

@@ -19,6 +19,53 @@ seed = hashlib.sha256(seed_text.encode()).digest()
 SERVER_KEY = Keypair.from_seed(seed)
 
 
+
+
+CHAIN_LENGTH = 256
+N = 32  # SHA256 output bytes
+
+def verify_wots(signature, message, public_key):
+    """
+    signature: list[str] (base58 encoded, length 32)
+    message: bytes
+    public_key: list[str] (base58 encoded, length 32)
+    """
+
+    # Basic sanity checks
+    if len(signature) != N or len(public_key) != N:
+        print("Invalid lengths")
+        return False
+
+    # Hash the message
+    msg_hash = hashlib.sha256(message).digest()
+
+    for i in range(N):
+        try:
+            sig_i = base58.b58decode(signature[i])
+            pk_i = base58.b58decode(public_key[i])
+        except Exception as e:
+            print(f"Base58 decode error at index {i}: {e}")
+            return False
+
+        check = sig_i
+
+        # Walk forward in hash chain
+        steps = CHAIN_LENGTH - msg_hash[i]
+
+        for _ in range(steps):
+            check = hashlib.sha256(check).digest()
+
+        # Compare with public key element
+        if check != pk_i:
+            print(f"Mismatch at index {i}")
+            return False
+
+    return True
+
+
+
+
+
 @app.post("/relay-txn")
 async def relay_txn(payload: dict):
     # 1. Verification Math
@@ -56,6 +103,9 @@ async def relay_txn(payload: dict):
 
 
 
+
+
+
 @app.post("/sign-message")
 async def sign_message(req: dict):
     #message = bytes(payload["message"])
@@ -83,20 +133,12 @@ async def sign_message(req: dict):
     signature = payload["signature"]        # list[str]
     public_key = payload["public_key"]      # list[str]
 
-    msg_hash = hashlib.sha256(message).digest()
-
-    for i in range(len(signature)):
-        sig_i = base58.b58decode(signature[i])
-        pk_i = base58.b58decode(public_key[i])
+    is_valid = verify_wots(signature, message, public_key)
 
-        check = sig_i
-        for _ in range(CHAIN_LENGTH - msg_hash[i]):
-            check = hashlib.sha256(check).digest()
-
-        if check != pk_i:
-            return {"status":False}
+    if not is_valid:
+        return {"success": False, "error": "Invalid WOTS signature"}
 
-    return {"status":True}
+    return {"success": True}