From 72fbeef4feb582a33c34b292a9598f15efb07f74 Mon Sep 17 00:00:00 2001
From: Filippo Romani <filippo@romani.cc>
Date: Fri, 29 Nov 2024 21:00:34 +0100
Subject: [PATCH] feat: error handling + various fixes

---
 .gitignore         |  2 ++
 README.md          | 17 ++++++++++++++---
 main.py            | 11 -----------
 test.py            | 12 ++++++++++++
 zerokno/zerokno.py | 21 +++++++++++++++------
 5 files changed, 43 insertions(+), 20 deletions(-)
 delete mode 100644 main.py
 create mode 100644 test.py

diff --git a/.gitignore b/.gitignore
index 0d3786d..0ada2ed 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,5 @@
 *.gz
 zerokno/zerokno.egg-info/PKG-INFO
 *.txt
+zerokno/__pycache__/zerokno.cpython-312.pyc
+test
diff --git a/README.md b/README.md
index 6874b8d..ba4ed79 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,12 @@
 # ZeroKno
 
 ZeroKno is a simple, easy to use and lightweight zero-knowledge password storage method for Python.
+The keys are stored in plain text, while the passwords are stored as base58-encoded SHA-256 hashes.
+
+This makes it impossible for the application to know the password, while still being able to verify it. 
+If you want to use this in a production environment, please make sure to use a secure storage method for the keys (e.g. a secure database).
+
+Note that the passwords or values are not encrypted, but hashed. This means that the password cannot be retrieved from the hash, but the hash can be used to verify the password. Basically, you can't recover the password if you lose it.
 
 ## Installation
 
@@ -13,10 +19,10 @@ pip install zerokno
 ```python
 from zerokno import ZeroKno
 
-# Create a new ZeroKno instance
-zk = ZeroKno()
+# Create a new ZeroKno instance -- app_secret is a secret key for the application and storage is a directory to store the password hashes
+zk = ZeroKno(app_secret, storage) 
 
-# Add a new password
+# Add a new password -- please note that the password is stored as a hash, while the user id is stored as plain text
 zk.store("password", "userid")
 
 # Check password match
@@ -28,3 +34,8 @@ zk.validate("password", "userid1")
 # Error: Userid not found
 
 ```
+
+## Errors 
+
+- `Key not found` - This error is raised when the user id is not found in the storage
+- Any other error is raised when the storage file is not found or the storage file is corrupted or in any other way not accessible
\ No newline at end of file
diff --git a/main.py b/main.py
deleted file mode 100644
index 1ffd0ed..0000000
--- a/main.py
+++ /dev/null
@@ -1,11 +0,0 @@
-import zerokno.zerokno as zerokno
-
-z = zerokno.Zerokno()
-
-pwd1 = input("Enter password: ")
-
-z.store(pwd1, "test")
-
-pwd2 = input("Enter password: ")
-
-print(z.validate(pwd2, "test"))
\ No newline at end of file
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..12b4f5d
--- /dev/null
+++ b/test.py
@@ -0,0 +1,12 @@
+import zerokno.zerokno as zerokno # local import
+
+z = zerokno.ZeroKno(storage="./")
+
+pwd1 = input("Enter password: ")
+
+z.store(pwd1, "test")
+
+pwd2 = input("Enter password: ")
+
+print(z.validate(pwd2, "test"))
+# if this prints True, then the code is working as expected
\ No newline at end of file
diff --git a/zerokno/zerokno.py b/zerokno/zerokno.py
index 453e822..631646a 100644
--- a/zerokno/zerokno.py
+++ b/zerokno/zerokno.py
@@ -3,18 +3,27 @@
 
 
 class ZeroKno:
-    def __init__(self, app_secret: str = None, storage: str = None):
+    def __init__(self, app_secret: str = "", storage: str = None):
         self.path = storage
         self.secret = app_secret
 
     @staticmethod
     def algo(password: str, secret: str) -> str:
-        return b58encode(sha256(password.encode()).digest())
+        try: return b58encode(sha256(password.encode()).digest()).decode()
+        except Exception as e: raise Exception(e)
     
     def store(self, password: str, uid: str) -> bool:
-        with open(self.path+uid, "w+") as f:
-            f.write(self.algo(password, self.secret))
+        try:
+            with open(self.path+uid, "w+") as f:
+                f.write(self.algo(password, self.secret))
+                return True
+        except Exception as e:
+            raise Exception(e)
+
 
     def validate(self, password: str, uid: str) -> bool:
-        with open(self.path+uid, "r") as f:
-            return self.algo(password, self.secret) == f.read()
\ No newline at end of file
+        try:
+            with open(self.path+uid, "r") as f:
+                return self.algo(password, self.secret) == f.read()
+        except FileNotFoundError:
+            raise Exception("Key not found")