Update app_logic.py

app_logic.py

@@ -11,202 +11,90 @@ from huggingface_hub import (
     list_repo_files,
     Repository,
     whoami,
-    hf_hub_download,
+    hf_hub_download, # New import
 )
 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_data_from_image_pil,
-        save_decoded_data_locally_encrypted,
-        load_decoded_data_locally_encrypted
-    )
+    from keylock_decode import decode_from_image_pil
     KEYLOCK_DECODE_AVAILABLE = True
 except ImportError:
     KEYLOCK_DECODE_AVAILABLE = False
-    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.")
-
+    decode_from_image_pil = None
+    logging.warning("keylock-decode library not found. KeyLock Wallet image feature will be disabled.")
 
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.INFO, 
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
-# --- 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 ---
+# --- Helper Function to Get API Token (Unchanged) ---
 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:
-        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
-
+    if env_token: return env_token, None
+    if ui_token_from_textbox: return ui_token_from_textbox, None
     return None, "Error: Hugging Face API token not provided."
 
-
-# --- 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
+# --- `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."
     status_messages_display = []
-
-    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
+    # ... (rest of the function, ensure debug logic is as intended or removed)
     try:
         logger.info(f"Attempting to decode from KeyLock Wallet image...")
-        decoded_data, status_msgs_from_lib = decode_data_from_image_pil(image_pil_object, password)
+        decoded_data, status_msgs_from_lib = decode_from_image_pil(image_pil_object, password, set_environment_variables=True)
         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}")
-
-            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 = {}
-
-
+            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)}")
     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):
-    # Processes markdown by stripping comments if a specific marker is present
+    """Process a commented markdown string by stripping '# ' from each line if '# # Space:' 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:"):
@@ -231,39 +119,41 @@ def parse_markdown(markdown_input):
     space_info["files"] = [f for f in space_info["files"] if f.get("path")]
     return space_info
 
-# --- Helper Function: Determine Repository ID ---
+
+
+# --- `_determine_repo_id` (Unchanged) ---
 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."
-        except Exception as e: error_message = f"Error retrieving username: {str(e)}."
+            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."
         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
 
-# --- Function: Fetch File Content from Hub ---
+
+# --- New Function to 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 None, token_err # Return error as second element for consistency
 
         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."
 
@@ -273,42 +163,45 @@ 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}")
-        return content, None
-
+        logger.info(f"Successfully downloaded and read file: {file_path_in_repo} from {repo_id}")
+        return content, None # Return content and no error
+        
     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}:")
+        logger.exception(f"Error fetching file content for {file_path_in_repo} from {repo_id_for_error_logging}:")
         return None, f"Error fetching file content: {str(e)}"
 
-# --- Function: List Space Files ---
+# --- Function to list files (reused, but now distinct from fetching content) ---
 def list_space_files_for_browsing(ui_api_token_from_textbox, space_name_ui, owner_ui):
-    # Lists files in a Hugging Face Space
+    """Lists files in a Hugging Face Space, returns list or error."""
     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 files, None
+            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
     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 Function: Create/Update Space ---
+
+# --- Core Functions: `create_space`, `update_space_file` (Unchanged from previous correct versions) ---
 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)
@@ -338,9 +231,7 @@ 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)