Update app_logic.py

app_logic.py

@@ -11,90 +11,202 @@ from huggingface_hub import (
     list_repo_files,
     Repository,
     whoami,
-    hf_hub_download, # New import
+    hf_hub_download,
 )
 import logging
 from pathlib import Path
-from PIL import Image 
+from PIL import Image
 
+# --- Import functions from keylock_decode ---
 try:
-    from keylock_decode import decode_from_image_pil
+    from keylock_decode import (
+        decode_data_from_image_pil,
+        save_decoded_data_locally_encrypted,
+        load_decoded_data_locally_encrypted
+    )
     KEYLOCK_DECODE_AVAILABLE = True
 except ImportError:
     KEYLOCK_DECODE_AVAILABLE = False
-    decode_from_image_pil = None
-    logging.warning("keylock-decode library not found. KeyLock Wallet image feature will be disabled.")
+    decode_data_from_image_pil = None
+    save_decoded_data_locally_encrypted = None
+    load_decoded_data_locally_encrypted = None
+    logging.warning("keylock-decode library not available or missing functions. KeyLock Wallet features will be disabled.")
+
 
 logging.basicConfig(
-    level=logging.INFO, 
+    level=logging.INFO,
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
-# --- Helper Function to Get API Token (Unchanged) ---
+# --- Global variable for in-memory keys and local storage path ---
+_loaded_local_keys = {}
+LOCAL_KEYS_FILE = "/data/.space_keys.enc" # Persistent storage in Hugging Face Spaces
+
+# --- Helper Function: Load Local Keys into Memory ---
+def _load_keys_into_memory(password: str) -> list[str]:
+    # Loads and decrypts keys from the local file into memory
+    global _loaded_local_keys
+    status_messages = []
+
+    if not KEYLOCK_DECODE_AVAILABLE or not load_decoded_data_locally_encrypted:
+        status_messages.append("KeyLock-Decode library not available for local loading.")
+        _loaded_local_keys = {}
+        return status_messages
+
+    if not password:
+        status_messages.append("Error: Password required to load local keys.")
+        _loaded_local_keys = {}
+        return status_messages
+
+    if not Path(LOCAL_KEYS_FILE).exists():
+         status_messages.append(f"Info: Local key file not found at `{LOCAL_KEYS_FILE}`. No keys loaded.")
+         _loaded_local_keys = {}
+         return status_messages
+
+    try:
+        _loaded_local_keys = load_decoded_data_locally_encrypted(password, LOCAL_KEYS_FILE)
+
+        if _loaded_local_keys:
+            status_messages.append(f"Successfully loaded {len(_loaded_local_keys)} keys from local storage.")
+            masked_keys = {k: ('********' if any(k_word in k.upper() for k_word in ['TOKEN', 'KEY', 'SECRET', 'PASS']) else v) for k, v in _loaded_local_keys.items()}
+            status_messages.append(f"Loaded keys: {masked_keys}")
+            logger.info(f"Keys loaded into memory from {LOCAL_KEYS_FILE}")
+        else:
+             status_messages.append(f"Info: Local key file `{LOCAL_KEYS_FILE}` was empty or contained no valid data after decryption.")
+             _loaded_local_keys = {}
+
+    except FileNotFoundError:
+        status_messages.append(f"Info: Local key file not found at {LOCAL_KEYS_FILE}. No keys loaded.")
+        _loaded_local_keys = {}
+    except ValueError as e:
+        status_messages.append(f"Error loading local keys: {e}. Password may be incorrect.")
+        _loaded_local_keys = {}
+    except IOError as e:
+         status_messages.append(f"IO Error loading local keys: {e}")
+         _loaded_local_keys = {}
+    except Exception as e:
+        logger.exception(f"Unexpected error loading local keys from {LOCAL_KEYS_FILE}:")
+        status_messages.append(f"Unexpected error loading local keys: {str(e)}")
+        _loaded_local_keys = {}
+
+    return status_messages
+
+
+# --- Helper Function: Get API Token ---
 def _get_api_token(ui_token_from_textbox=None):
+    # Retrieves token from in-memory keys, env vars, or textbox
+    in_memory_token = _loaded_local_keys.get('HF_TOKEN')
+    if in_memory_token:
+        logger.info("Using HF_TOKEN from in-memory loaded keys.")
+        return in_memory_token, None
+
     env_token = os.getenv('HF_TOKEN')
-    if env_token: return env_token, None
-    if ui_token_from_textbox: return ui_token_from_textbox, None
+    if env_token:
+        logger.info("Using HF_TOKEN from environment variable.")
+        return env_token, None
+
+    if ui_token_from_textbox:
+        logger.info("Using HF_TOKEN from UI textbox.")
+        return ui_token_from_textbox, None
+
     return None, "Error: Hugging Face API token not provided."
 
-# --- `load_token_from_image_and_set_env` (Unchanged - Terminology and debug logic as before) ---
-def load_token_from_image_and_set_env(image_pil_object: Image.Image, password: str):
-    if not KEYLOCK_DECODE_AVAILABLE: return "Error: KeyLock-Decode library is not installed."
-    if image_pil_object is None: return "Error: No KeyLock Wallet image provided."
-    if not password: return "Error: Password cannot be empty."
+
+# --- Main Function: Process KeyLock Image and Store Locally ---
+def process_keylock_image_and_store_locally(image_pil_object: Image.Image, password: str):
+    # Decodes image, saves data encrypted locally, and loads into memory
     status_messages_display = []
-    # ... (rest of the function, ensure debug logic is as intended or removed)
+
+    if not KEYLOCK_DECODE_AVAILABLE or not decode_data_from_image_pil or \
+       not save_decoded_data_locally_encrypted or not load_decoded_data_locally_encrypted:
+        status_messages_display.append("Error: KeyLock-Decode library not available.")
+        global _loaded_local_keys
+        _loaded_local_keys = {}
+        return "\n".join(status_messages_display)
+
+    if image_pil_object is None:
+        status_messages_display.append("Error: No KeyLock Wallet image provided.")
+        return "\n".join(status_messages_display)
+
+    if not password:
+        status_messages_display.append("Error: Password cannot be empty.")
+        return "\n".join(status_messages_display)
+
+    decoded_data = None
     try:
         logger.info(f"Attempting to decode from KeyLock Wallet image...")
-        decoded_data, status_msgs_from_lib = decode_from_image_pil(image_pil_object, password, set_environment_variables=True)
+        decoded_data, status_msgs_from_lib = decode_data_from_image_pil(image_pil_object, password)
         status_messages_display.extend(status_msgs_from_lib)
+
         if decoded_data:
             status_messages_display.append("\n**Decoded Data Summary (sensitive values masked):**")
             for key, value in decoded_data.items():
                 display_value = '********' if any(k_word in key.upper() for k_word in ['TOKEN', 'KEY', 'SECRET', 'PASS']) else value
                 status_messages_display.append(f"- {key}: {display_value}")
-            if os.getenv('HF_TOKEN'): status_messages_display.append(f"\n**SUCCESS: HF_TOKEN set from KeyLock Wallet image.**")
-            # ... (other status messages)
-    except ValueError as e: status_messages_display.append(f"**Decoding Error:** {e}")
-    except Exception as e: status_messages_display.append(f"**Unexpected decoding error:** {str(e)}")
+
+            try:
+                save_decoded_data_locally_encrypted(decoded_data, password, LOCAL_KEYS_FILE)
+                status_messages_display.append(f"\n**SUCCESS: Decoded data saved to encrypted local file:** `{LOCAL_KEYS_FILE}`")
+
+                load_status = _load_keys_into_memory(password)
+                status_messages_display.extend(load_status)
+                if not _loaded_local_keys:
+                    status_messages_display.append("Warning: No keys were loaded into memory after saving.")
+
+
+            except Exception as e:
+                 status_messages_display.append(f"Error saving or loading data locally: {str(e)}")
+                 logger.error(f"Error during local save/load after decoding: {e}")
+                 global _loaded_local_keys
+                 _loaded_local_keys = {}
+
+        elif not status_msgs_from_lib:
+             status_messages_display.append("Info: Decoding process completed, but no data was extracted.")
+
+
+    except ValueError as e:
+        status_messages_display.append(f"**Decoding Error:** {e}. Please check password and image.")
+        global _loaded_local_keys
+        _loaded_local_keys = {}
+    except Exception as e:
+        status_messages_display.append(f"**Unexpected error during processing:** {str(e)}")
+        logger.exception("Unexpected error during keylock image processing:")
+        global _loaded_local_keys
+        _loaded_local_keys = {}
+
+
     return "\n".join(status_messages_display)
 
+# --- Function: Load Keys from Local File ---
+def load_keys_from_local_file(password: str) -> str:
+    # Triggers loading keys from the encrypted local file into memory
+    status_messages = _load_keys_into_memory(password)
+    if not status_messages:
+         if _loaded_local_keys:
+             return f"Keys successfully loaded from `{LOCAL_KEYS_FILE}`."
+         else:
+             return f"Attempted to load keys from `{LOCAL_KEYS_FILE}`, but no keys were loaded. Check password."
 
+    return "\n".join(status_messages)
 
-'''def process_commented_markdown(commented_input):
-    """Process a commented markdown string by stripping '# ' from each line if '# # Space:' is present."""
-    lines = commented_input.strip().split("\n")
-    print(type(lines))
-    # Check for '# # Space:' or variations (e.g., '# Space:') in any line
-    if any( "# # Space:" in line.strip() for line in lines):
-        print("YES")
-        cleaned_lines = [line.lstrip("# ") for line in lines]
-        return cleaned_lines
-    return lines'''
 
+# --- Markdown Processing Functions ---
 def process_commented_markdown(commented_input):
-    """Process a commented markdown string by stripping '# ' from each line if '# # Space:' is present."""
+    # Processes markdown by stripping comments if a specific marker is present
     lines = commented_input.strip().split("\n")
-    print(type(lines)) # Original debug print
-    # Check for '# # Space:' or variations (e.g., '# Space:') in any line
     if any( "# # Space:" in line.strip() for line in lines):
-        print("YES") # Original debug print
-        
         cleaned_lines = [line[2:] if line.startswith("# ") else line for line in lines]
-        
         return cleaned_lines
     return lines
 
 def parse_markdown(markdown_input):
+    # Parses markdown input to extract space info and file content
     space_info = {"repo_name_md": "", "owner_md": "", "files": []}
     current_file_path = None; current_file_content_lines = []
     in_file_definition = False; in_code_block = False
-    print(markdown_input)
     lines = process_commented_markdown(markdown_input)
-    print(lines)
-    #lines = markdown_input.strip().split("\n")
-    
+
     for line_content_orig in lines:
         line_content_stripped = line_content_orig.strip()
         if line_content_stripped.startswith("### File:"):
@@ -119,41 +231,39 @@ def parse_markdown(markdown_input):
     space_info["files"] = [f for f in space_info["files"] if f.get("path")]
     return space_info
 
-
-
-# --- `_determine_repo_id` (Unchanged) ---
+# --- Helper Function: Determine Repository ID ---
 def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
+    # Determines the full repo ID (owner/space_name)
     if not space_name_ui: return None, "Error: Space Name cannot be empty."
     if "/" in space_name_ui: return None, "Error: Space Name should not contain '/'. Use Owner field."
     final_owner = owner_ui; error_message = None
     if not final_owner:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
-        if token_err: return None, token_err 
+        if token_err: return None, token_err
         if not resolved_api_token: return None, "Error: API token required for auto owner determination."
         try:
             user_info = whoami(token=resolved_api_token)
             if user_info and 'name' in user_info: final_owner = user_info['name']
-            else: error_message = "Error: Could not retrieve username. Check token/permissions or specify Owner."
-        except Exception as e: error_message = f"Error retrieving username: {str(e)}. Specify Owner."
+            else: error_message = "Error: Could not retrieve username. Check token/permissions."
+        except Exception as e: error_message = f"Error retrieving username: {str(e)}."
         if error_message: return None, error_message
     if not final_owner: return None, "Error: Owner could not be determined."
     return f"{final_owner}/{space_name_ui}", None
 
-
-# --- New Function to Fetch File Content from Hub ---
+# --- Function: Fetch File Content from Hub ---
 def get_space_file_content(ui_api_token_from_textbox, space_name_ui, owner_ui, file_path_in_repo):
-    """Fetches content of a specific file from a Hugging Face Space."""
+    # Fetches content of a specific file from a Hugging Face Space
     repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
         if token_err:
-            return None, token_err # Return error as second element for consistency
+            return None, token_err
 
         repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
         if err:
             return None, err
         repo_id_for_error_logging = repo_id
-        
+
         if not file_path_in_repo:
             return None, "Error: File path cannot be empty."
 
@@ -163,45 +273,42 @@ def get_space_file_content(ui_api_token_from_textbox, space_name_ui, owner_ui, f
             filename=file_path_in_repo,
             repo_type="space",
             token=resolved_api_token,
-            # revision="main", # Optional: specify a branch/commit
-            # cache_dir=... # Optional: manage cache
         )
-        
+
         content = Path(downloaded_file_path).read_text(encoding="utf-8")
-        logger.info(f"Successfully downloaded and read file: {file_path_in_repo} from {repo_id}")
-        return content, None # Return content and no error
-        
+        logger.info(f"Successfully downloaded and read file: {file_path_in_repo}")
+        return content, None
+
     except Exception as e:
-        # Catch specific huggingface_hub.utils.HFValidationError for not found etc.
         if "404" in str(e) or "not found" in str(e).lower():
              logger.warning(f"File not found {file_path_in_repo} in {repo_id_for_error_logging}: {e}")
              return None, f"Error: File '{file_path_in_repo}' not found in Space '{repo_id_for_error_logging}'."
-        logger.exception(f"Error fetching file content for {file_path_in_repo} from {repo_id_for_error_logging}:")
+        logger.exception(f"Error fetching file content for {file_path_in_repo}:")
         return None, f"Error fetching file content: {str(e)}"
 
-# --- Function to list files (reused, but now distinct from fetching content) ---
+# --- Function: List Space Files ---
 def list_space_files_for_browsing(ui_api_token_from_textbox, space_name_ui, owner_ui):
-    """Lists files in a Hugging Face Space, returns list or error."""
+    # Lists files in a Hugging Face Space
     repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
         if token_err: return None, token_err
-        
+
         repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
         if err: return None, err
         repo_id_for_error_logging = repo_id
-            
+
         files = list_repo_files(repo_id=repo_id, token=resolved_api_token, repo_type="space")
         if not files:
-            return [], f"No files found in Space `{repo_id}`." # Return empty list and info message
-        return files, None # Return list of files and no error
+            return [], f"No files found in Space `{repo_id}`."
+        return files, None
     except Exception as e:
         logger.exception(f"Error listing files for {repo_id_for_error_logging}:")
         return None, f"Error listing files for `{repo_id_for_error_logging}`: {str(e)}"
 
-
-# --- Core Functions: `create_space`, `update_space_file` (Unchanged from previous correct versions) ---
+# --- Core Function: Create/Update Space ---
 def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
+    # Creates or updates a Hugging Face Space with files from markdown
     repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
@@ -231,7 +338,9 @@ def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, mar
         logger.exception(f"Error in create_space for {repo_id_for_error_logging}:")
         return f"Error during Space creation/update: {str(e)}"
 
+# --- Core Function: Update Space File ---
 def update_space_file(ui_api_token_from_textbox, space_name_ui, owner_ui, file_path_in_repo, file_content, commit_message_ui):
+    # Updates a specific file in a Hugging Face Space
     repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
         resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)