Update app.py

app.py

@@ -1,12 +1,3 @@
-# Initialize session state variables if they don't exist
-if 'failed_attempts' not in streamlit.session_state:
-    streamlit.session_state.failed_attempts = 0
-if 'stored_data' not in streamlit.session_state:
-    streamlit.session_state.stored_data = {}
-if 'current_page' not in streamlit.session_state:
-    streamlit.session_state.current_page = "Home"
-if 'last_attempt_time' not in streamlit.session_state:
-    streamlit.session_state.last_attempt_time = 0
 
 # Function to hash passkey
 def hash_passkey(passkey):
@@ -37,16 +28,11 @@ def decrypt_data(encrypted_text, passkey, data_id):
             streamlit.session_state.failed_attempts += 1
             streamlit.session_state.last_attempt_time = time.time()
             return None
-    except Exception as e:
+    except Exception:
         streamlit.session_state.failed_attempts += 1
         streamlit.session_state.last_attempt_time = time.time()
         return None
 
-# Function to generate a unique ID for data
-def generate_data_id():
-    import uuid
-    return str(uuid.uuid4())
-
 # Function to reset failed attempts
 def reset_failed_attempts():
     streamlit.session_state.failed_attempts = 0
@@ -55,108 +41,125 @@ def reset_failed_attempts():
 def change_page(page):
     streamlit.session_state.current_page = page
 
-# Streamlit UI
-streamlit.title("🔒 Secure Data Encryption System")
-
-# Navigation
-menu = ["Home", "Store Data", "Retrieve Data", "Login"]
-choice = streamlit.sidebar.selectbox("Navigation", menu, index=menu.index(streamlit.session_state.current_page))
-
-# Update current page based on selection
-streamlit.session_state.current_page = choice
-
-# Check if too many failed attempts
-if streamlit.session_state.failed_attempts >= 3:
-    streamlit.session_state.current_page = "Login"
-    streamlit.warning("🔒 Too many failed attempts! Reauthorization required.")
-
-# Display current page
-if streamlit.session_state.current_page == "Home":
-    streamlit.subheader("🏠 Welcome to the Secure Data System")
-    streamlit.write("Use this app to **securely store and retrieve data** using unique passkeys.")
-
-    col1, col2 = streamlit.columns(2)
-    with col1:
-        if streamlit.button("Store New Data", use_container_width=True):
-            change_page("Store Data")
-    with col2:
-        if streamlit.button("Retrieve Data", use_container_width=True):
-            change_page("Retrieve Data")
-
-    streamlit.info(f"Currently storing {len(streamlit.session_state.stored_data)} encrypted data entries.")
-
-elif streamlit.session_state.current_page == "Store Data":
-    streamlit.subheader("📂 Store Data Securely")
-    user_data = streamlit.text_area("Enter Data:")
-    passkey = streamlit.text_input("Enter Passkey:", type="password")
-    confirm_passkey = streamlit.text_input("Confirm Passkey:", type="password")
-
-    if streamlit.button("Encrypt & Save"):
-        if user_data and passkey and confirm_passkey:
-            if passkey != confirm_passkey:
-                streamlit.error("⚠️ Passkeys do not match!")
+def main():
+    # Initialize session state variables
+    if 'failed_attempts' not in streamlit.session_state:
+        streamlit.session_state.failed_attempts = 0
+    if 'stored_data' not in streamlit.session_state:
+        streamlit.session_state.stored_data = {}
+    if 'current_page' not in streamlit.session_state:
+        streamlit.session_state.current_page = "Home"
+    if 'last_attempt_time' not in streamlit.session_state:
+        streamlit.session_state.last_attempt_time = 0
+
+    streamlit.title("🔒 Secure Data Encryption System")
+
+    # Navigation
+    menu = ["Home", "Store Data", "Retrieve Data", "Login"]
+    choice = streamlit.sidebar.selectbox("Navigation", menu, index=menu.index(streamlit.session_state.current_page))
+    streamlit.session_state.current_page = choice
+
+    # Block if too many failed attempts
+    if streamlit.session_state.failed_attempts >= 3:
+        streamlit.session_state.current_page = "Login"
+        streamlit.warning("🔒 Too many failed attempts! Reauthorization required.")
+
+    # HOME PAGE
+    if streamlit.session_state.current_page == "Home":
+        streamlit.subheader("🏠 Welcome to the Secure Data System")
+        streamlit.write("Use this app to **securely store and retrieve data** using unique passkeys.")
+
+        col1, col2 = streamlit.columns(2)
+        with col1:
+            if streamlit.button("Store New Data", use_container_width=True):
+                change_page("Store Data")
+        with col2:
+            if streamlit.button("Retrieve Data", use_container_width=True):
+                change_page("Retrieve Data")
+
+        streamlit.info(f"Currently storing {len(streamlit.session_state.stored_data)} encrypted data entries.")
+
+    # STORE DATA PAGE
+    elif streamlit.session_state.current_page == "Store Data":
+        streamlit.subheader("📂 Store Data Securely")
+        user_data = streamlit.text_area("Enter Data:")
+        passkey = streamlit.text_input("Enter Passkey:", type="password")
+        confirm_passkey = streamlit.text_input("Confirm Passkey:", type="password")
+
+        if streamlit.button("Encrypt & Save"):
+            if user_data and passkey and confirm_passkey:
+                if passkey != confirm_passkey:
+                    streamlit.error("⚠️ Passkeys do not match!")
+                else:
+                    data_id = str(uuid.uuid4())
+                    hashed_passkey = hash_passkey(passkey)
+                    encrypted_text = encrypt_data(user_data, passkey)
+
+                    streamlit.session_state.stored_data[data_id] = {
+                        "encrypted_text": encrypted_text,
+                        "passkey": hashed_passkey
+                    }
+
+                    streamlit.success("✅ Data stored securely!")
+                    streamlit.code(data_id, language="text")
+                    streamlit.info("⚠️ Save this Data ID! You'll need it to retrieve your data.")
             else:
-                data_id = generate_data_id()
-                hashed_passkey = hash_passkey(passkey)
-                encrypted_text = encrypt_data(user_data, passkey)
-                streamlit.session_state.stored_data[data_id] = {
-                    "encrypted_text": encrypted_text,
-                    "passkey": hashed_passkey
-                }
-                streamlit.success("✅ Data stored securely!")
-                streamlit.code(data_id, language="text")
-                streamlit.info("⚠️ Save this Data ID! You'll need it to retrieve your data.")
-        else:
-            streamlit.error("⚠️ All fields are required!")
-
-elif streamlit.session_state.current_page == "Retrieve Data":
-    streamlit.subheader("🔍 Retrieve Your Data")
-
-    attempts_remaining = 3 - streamlit.session_state.failed_attempts
-    streamlit.info(f"Attempts remaining: {attempts_remaining}")
+                streamlit.error("⚠️ All fields are required!")
+
+    # RETRIEVE DATA PAGE
+    elif streamlit.session_state.current_page == "Retrieve Data":
+        streamlit.subheader("🔍 Retrieve Your Data")
+        attempts_remaining = 3 - streamlit.session_state.failed_attempts
+        streamlit.info(f"Attempts remaining: {attempts_remaining}")
+
+        data_id = streamlit.text_input("Enter Data ID:")
+        passkey = streamlit.text_input("Enter Passkey:", type="password")
+
+        if streamlit.button("Decrypt"):
+            if data_id and passkey:
+                if data_id in streamlit.session_state.stored_data:
+                    encrypted_text = streamlit.session_state.stored_data[data_id]["encrypted_text"]
+                    decrypted_text = decrypt_data(encrypted_text, passkey, data_id)
+
+                    if decrypted_text:
+                        streamlit.success("✅ Decryption successful!")
+                        streamlit.markdown("### Your Decrypted Data:")
+                        streamlit.code(decrypted_text, language="text")
+                    else:
+                        streamlit.error(f"❌ Incorrect passkey! Attempts remaining: {3 - streamlit.session_state.failed_attempts}")
+                else:
+                    streamlit.error("❌ Data ID not found!")
 
-    data_id = streamlit.text_input("Enter Data ID:")
-    passkey = streamlit.text_input("Enter Passkey:", type="password")
+                if streamlit.session_state.failed_attempts >= 3:
+                    streamlit.warning("🔒 Too many failed attempts! Redirecting to Login Page.")
+                    streamlit.session_state.current_page = "Login"
+                    streamlit.rerun()
+            else:
+                streamlit.error("⚠️ Both fields are required!")
 
-    if streamlit.button("Decrypt"):
-        if data_id and passkey:
-            if data_id in streamlit.session_state.stored_data:
-                encrypted_text = streamlit.session_state.stored_data[data_id]["encrypted_text"]
-                decrypted_text = decrypt_data(encrypted_text, passkey, data_id)
+    # LOGIN PAGE
+    elif streamlit.session_state.current_page == "Login":
+        streamlit.subheader("🔑 Reauthorization Required")
 
-                if decrypted_text:
-                    streamlit.success("✅ Decryption successful!")
-                    streamlit.markdown("### Your Decrypted Data:")
-                    streamlit.code(decrypted_text, language="text")
+        if time.time() - streamlit.session_state.last_attempt_time < 10 and streamlit.session_state.failed_attempts >= 3:
+            remaining_time = int(10 - (time.time() - streamlit.session_state.last_attempt_time))
+            streamlit.warning(f"🕒 Please wait {remaining_time} seconds before trying again.")
+        else:
+            login_pass = streamlit.text_input("Enter Master Password:", type="password")
+
+            if streamlit.button("Login"):
+                if login_pass == "admin123":
+                    reset_failed_attempts()
+                    streamlit.success("✅ Reauthorized successfully!")
+                    streamlit.session_state.current_page = "Home"
+                    streamlit.rerun()
                 else:
-                    streamlit.error(f"❌ Incorrect passkey! Attempts remaining: {3 - streamlit.session_state.failed_attempts}")
-            else:
-                streamlit.error("❌ Data ID not found!")
+                    streamlit.error("❌ Incorrect password!")
 
-            if streamlit.session_state.failed_attempts >= 3:
-                streamlit.warning("🔒 Too many failed attempts! Redirecting to Login Page.")
-                streamlit.session_state.current_page = "Login"
-                streamlit.rerun()
-        else:
-            streamlit.error("⚠️ Both fields are required!")
-
-elif streamlit.session_state.current_page == "Login":
-    streamlit.subheader("🔑 Reauthorization Required")
-
-    if time.time() - streamlit.session_state.last_attempt_time < 10 and streamlit.session_state.failed_attempts >= 3:
-        remaining_time = int(10 - (time.time() - streamlit.session_state.last_attempt_time))
-        streamlit.warning(f"🕒 Please wait {remaining_time} seconds before trying again.")
-    else:
-        login_pass = streamlit.text_input("Enter Master Password:", type="password")
-
-        if streamlit.button("Login"):
-            if login_pass == "admin123":
-                reset_failed_attempts()
-                streamlit.success("✅ Reauthorized successfully!")
-                streamlit.session_state.current_page = "Home"
-                streamlit.rerun()
-            else:
-                streamlit.error("❌ Incorrect password!")
+    # Footer
+    streamlit.markdown("---")
+    streamlit.markdown("🔐 Secure Data Encryption System | Educational Project")
 
-streamlit.markdown("---")
-streamlit.markdown("🔐 Secure Data Encryption System | Educational Project")
+# Run app
+if __name__ == "__main__":
+    main()