Update app_logic.py

app_logic.py

@@ -1,65 +1,151 @@
 import os
 import re
 import tempfile
-import shutil
-import git
+import git # Used by Repository
 from huggingface_hub import (
     create_repo,
     upload_folder,
     list_repo_files,
-    # delete_file, # Not used
-    Repository,
+    Repository, # Used in update_space_file
     whoami,
 )
 import logging
 from pathlib import Path
+from PIL import Image # For type hinting from gr.Image(type="pil")
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+# Attempt to import keylock_decode
+try:
+    from keylock_decode import decode_from_image_pil
+    KEYLOCK_DECODE_AVAILABLE = True
+except ImportError:
+    KEYLOCK_DECODE_AVAILABLE = False
+    decode_from_image_pil = None # Placeholder
+    logging.warning("keylock-decode library not found. Image decoding feature will be disabled.")
+
+
+# Configure logging (ensure this is suitable for your environment, might be set in app.py too)
+# If keylock-decode uses its own logger (e.g., logging.getLogger('keylock_decode')),
+# you might need to adjust its level separately if you want more/less verbose output from it.
+logging.basicConfig(
+    level=logging.INFO, 
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
 logger = logging.getLogger(__name__)
+if KEYLOCK_DECODE_AVAILABLE: # Example: Set keylock_decode's logger to DEBUG if you want its verbose output
+    # logging.getLogger('keylock_decode.decoder').setLevel(logging.DEBUG) # Check actual logger name in keylock-decode
+    pass
 
-# Function to parse markdown input
+
+# --- New Helper Function to Get API Token ---
+def _get_api_token(ui_token_from_textbox=None):
+    """
+    Retrieves the Hugging Face API token.
+    Prioritizes HF_TOKEN from environment variables (potentially set by KeyLock-Decode),
+    then falls back to UI input.
+    Returns (token, error_message)
+    """
+    env_token = os.getenv('HF_TOKEN')
+    if env_token:
+        logger.info("Using HF_TOKEN from environment.")
+        return env_token, None
+    
+    if ui_token_from_textbox:
+        logger.info("Using API token from UI textbox.")
+        return ui_token_from_textbox, None
+    
+    logger.warning("HF API token not found in environment or UI textbox.")
+    return None, "Error: Hugging Face API token not provided. Please enter it in the textbox or load it from an image."
+
+
+# --- New Function for KeyLock-Decode Integration ---
+def load_token_from_image_and_set_env(image_pil_object: Image.Image, password: str):
+    """
+    Decodes data from a PIL Image object using keylock-decode, sets environment variables,
+    and returns status messages.
+    """
+    if not KEYLOCK_DECODE_AVAILABLE:
+        return "Error: KeyLock-Decode library is not installed. This feature is disabled."
+
+    if image_pil_object is None:
+        return "Error: No image provided for decoding."
+    if not password:
+        return "Error: Password cannot be empty for image decoding."
+
+    status_messages_display = []
+    
+    try:
+        logger.info(f"Attempting to decode from image using KeyLock-Decode...")
+        # decode_from_image_pil expects a PIL Image object
+        decoded_data, status_msgs_from_lib = decode_from_image_pil(
+            image_pil_object,
+            password,
+            set_environment_variables=True # CRITICAL: This sets os.environ
+        )
+
+        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 was found and has been set in the current process environment.** It will be prioritized for operations.")
+            elif 'HF_TOKEN' in decoded_data: # Should be rare if set_env_vars=True
+                 status_messages_display.append(f"\nWarning: HF_TOKEN was in decoded data but os.getenv('HF_TOKEN') is not picking it up. This is unexpected if 'set_environment_variables=True' worked.")
+            else:
+                status_messages_display.append("\nNote: HF_TOKEN was not specifically found in the decoded data from the image.")
+        else:
+            # This case is often covered by status_msgs_from_lib, but good to have a fallback.
+            if not any("No payload data found" in msg for msg in status_msgs_from_lib):
+                 status_messages_display.append("No data was decoded, or the decoded data was empty.")
+
+    except ValueError as e: # Specific errors from keylock-decode (e.g., bad password, corrupted data)
+        logger.error(f"KeyLock-Decode ValueError: {e}")
+        status_messages_display.append(f"**Decoding Error:** {e}")
+    except Exception as e:
+        logger.exception("An unexpected error occurred during image decoding with KeyLock-Decode:")
+        status_messages_display.append(f"**An unexpected error occurred:** {str(e)}")
+    
+    return "\n".join(status_messages_display)
+
+
+# --- Original `parse_markdown` function (unchanged) ---
 def parse_markdown(markdown_input):
     """Parse markdown input to extract space details and file structure."""
     space_info = {"repo_name_md": "", "owner_md": "", "files": []}
     current_file = None
     file_content = []
-    in_file_content = False # Tracks if we are inside a ### File: block content
-    in_code_block = False # Tracks if we are inside a ``` code block ```
+    in_file_content = False 
+    in_code_block = False 
 
     lines = markdown_input.strip().split("\n")
     for line_idx, line_content_orig in enumerate(lines):
         line_content_stripped = line_content_orig.strip()
 
-        # Handle file content collection, especially for code blocks
         if in_file_content:
             if line_content_stripped.startswith("```"):
-                if in_code_block: # Closing ```
-                    file_content.append(line_content_orig) # Keep the closing backticks as part of content
+                if in_code_block: 
+                    file_content.append(line_content_orig) 
                     in_code_block = False
-                    # Don't immediately save here, let the next ### File or end of input handle it
-                    # This allows for text after a code block but before the next file.
-                else: # Opening ```
+                else: 
                     in_code_block = True
                     file_content.append(line_content_orig)
-            elif in_code_block: # Inside a code block
+            elif in_code_block: 
                 file_content.append(line_content_orig)
-            elif not in_code_block: # Plain text line within ### File: block but outside ```
-                # Check if this line is a new file marker, if so, current file ends.
+            elif not in_code_block: 
                 if line_content_stripped.startswith("### File:") or line_content_stripped.startswith("## File Structure") or line_content_stripped.startswith("# Space:"):
                     if current_file and file_content:
                         space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
-                    current_file = None # Reset
+                    current_file = None 
                     file_content = []
-                    in_file_content = False # Current file ended
-                    # Reprocess this line if it's a new file marker (will be handled by outer ifs)
-                else: # Regular content line
+                    in_file_content = False 
+                else: 
                      file_content.append(line_content_orig)
 
-
-        # Detect major structural elements
         if line_content_stripped.startswith("# Space:"):
-            if current_file and file_content: # Save previous file if any
+            if current_file and file_content: 
                 space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
             full_space_name_md = line_content_stripped.replace("# Space:", "").strip()
             if "/" in full_space_name_md:
@@ -72,58 +158,55 @@ def parse_markdown(markdown_input):
             in_code_block = False
 
         elif line_content_stripped.startswith("## File Structure"):
-            if current_file and file_content: # Save previous file if any
+            if current_file and file_content: 
                 space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
             current_file = None
             file_content = []
             in_file_content = False
             in_code_block = False
-            continue # Just a section header
+            continue 
 
         elif line_content_stripped.startswith("### File:"):
-            if current_file and file_content: # Save content of the previous file
+            if current_file and file_content: 
                 space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
             
             current_file = line_content_stripped.replace("### File:", "").strip()
-            file_content = [] # Reset for new file
-            in_file_content = True # Start collecting content lines for this file
-            in_code_block = False # Reset code block state for new file
-
-        # Note: 📄 and 📁 are ignored if ### File: is the primary mechanism as implemented.
-        # If they are meant to define empty files, that logic would need to be added.
-        # Current parser prioritizes ### File: sections for content.
-
-    # Append the last file's content if any
+            file_content = [] 
+            in_file_content = True 
+            in_code_block = False 
+            
     if current_file and file_content:
         space_info["files"].append({"path": current_file, "content": "\n".join(file_content)})
     
-    space_info["files"] = [f for f in space_info["files"] if f.get("path")] # Filter out empty path entries
+    space_info["files"] = [f for f in space_info["files"] if f.get("path")]
     return space_info
 
 
-def _determine_repo_id(api_token, space_name_ui, owner_ui):
-    """
-    Determines the final owner and constructs the repo_id.
-    space_name_ui should be just the name, not 'owner/name'.
-    owner_ui is the value from the UI's owner field.
-    Returns (repo_id, error_message)
-    """
+# --- Updated `_determine_repo_id` to use `_get_api_token` ---
+def _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui):
     if not space_name_ui:
         return None, "Error: Space Name cannot be empty."
-    if "/" in space_name_ui: # User should not put slash in space name field
+    if "/" in space_name_ui:
         return None, "Error: Space Name should not contain '/'. Please use the Owner field for the namespace."
 
     final_owner = owner_ui
     error_message = None
 
     if not final_owner:
-        if not api_token:
-            return None, "Error: API token is required to automatically determine owner when Owner field is empty."
+        # Get the actual token to use for whoami
+        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
+        if token_err:
+            return None, token_err 
+        # If resolved_api_token is None here but token_err is None, it means _get_api_token had an issue not returning error string
+        # which _get_api_token is designed not to do, it should always return error if no token.
+        if not resolved_api_token: 
+             return None, "Error: API token is required to automatically determine owner when Owner field is empty (internal check failed)."
+
         try:
-            user_info = whoami(token=api_token)
+            user_info = whoami(token=resolved_api_token)
             if user_info and 'name' in user_info:
                 final_owner = user_info['name']
-                logger.info(f"Determined owner: {final_owner} from API token.")
+                logger.info(f"Determined owner: {final_owner} using API token.")
             else:
                 logger.error(f"whoami(token=...) returned: {user_info} - 'name' field missing or user_info is None.")
                 error_message = "Error: Could not retrieve username from API token. Ensure token is valid and has 'Read profile' permissions. Or, specify Owner manually."
@@ -134,132 +217,118 @@ def _determine_repo_id(api_token, space_name_ui, owner_ui):
         if error_message:
             return None, error_message
 
-    if not final_owner: 
+    if not final_owner:
         return None, "Error: Owner could not be determined. Please provide an owner or ensure your API token is valid."
 
     return f"{final_owner}/{space_name_ui}", None
 
-# Function to create and populate a Space
-def create_space(api_token, space_name_ui, owner_ui, sdk_ui, markdown_input):
-    """Create a Hugging Face Space and populate it with files from markdown input."""
+
+# --- Updated Core Functions to use `_get_api_token` ---
+
+def create_space(ui_api_token_from_textbox, space_name_ui, owner_ui, sdk_ui, markdown_input):
+    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
-        if not api_token:
-            return "Error: Please provide a valid Hugging Face API token."
+        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
+        if token_err:
+            return token_err
 
-        repo_id, err = _determine_repo_id(api_token, space_name_ui, owner_ui)
+        # _determine_repo_id will also call _get_api_token internally if owner_ui is blank
+        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
         if err:
             return err
+        repo_id_for_error_logging = repo_id # Update for more accurate logging if determination succeeds
         
         space_info = parse_markdown(markdown_input)
         if not space_info["files"]:
-            return "Error: No files found in the markdown input. Ensure '### File: path/to/file.ext' markers are used correctly with content."
+            return "Error: No files found in the markdown input. Ensure '### File: path/to/file.ext' markers are used correctly."
 
-        # Create temporary directory
         with tempfile.TemporaryDirectory() as temp_dir:
-            repo_local_path = Path(temp_dir) / "repo_content_for_upload"
+            repo_local_path = Path(temp_dir) / "repo_upload_content"
             repo_local_path.mkdir(exist_ok=True)
 
-            # Write files to temporary directory
             for file_info in space_info["files"]:
-                if not file_info.get("path"):
-                    logger.warning(f"Skipping file with no path: {file_info}")
-                    continue
+                if not file_info.get("path"): continue
                 file_path_abs = repo_local_path / file_info["path"]
                 file_path_abs.parent.mkdir(parents=True, exist_ok=True)
-                with open(file_path_abs, "w", encoding="utf-8") as f:
-                    f.write(file_info["content"])
+                with open(file_path_abs, "w", encoding="utf-8") as f: f.write(file_info["content"])
                 logger.info(f"Wrote file: {file_path_abs}")
 
-            # Create repository on Hugging Face
             try:
                 create_repo(
-                    repo_id=repo_id,
-                    token=api_token,
-                    repo_type="space", # Correctly set
-                    space_sdk=sdk_ui,
-                    private=False,
+                    repo_id=repo_id, token=resolved_api_token, repo_type="space",
+                    space_sdk=sdk_ui, private=False,
                 )
                 logger.info(f"Created Space repo: {repo_id}")
             except Exception as e:
                 err_str = str(e).lower()
-                if "already exists" in err_str or "you already created this repo" in err_str or "exists" in err_str: # More robust check
-                    logger.info(f"Space {repo_id} already exists, proceeding to upload/update files.")
+                if "already exists" in err_str or "you already created this repo" in err_str or "exists" in err_str:
+                    logger.info(f"Space {repo_id} already exists. Proceeding to upload/update files.")
                 else:
                     return f"Error creating Space '{repo_id}': {str(e)}"
             
-            # Push to Hugging Face Space
             upload_folder(
-                repo_id=repo_id,
-                folder_path=str(repo_local_path), # upload_folder expects string path
-                path_in_repo=".",
-                token=api_token,
-                repo_type="space",  # ***** ADD THIS LINE *****
+                repo_id=repo_id, folder_path=str(repo_local_path), path_in_repo=".",
+                token=resolved_api_token, repo_type="space",
                 commit_message=f"Initial Space setup of {repo_id} via Builder",
-                # allow_patterns=["*.py", "*.md", "*.txt", "Dockerfile", ".gitattributes", "*.json", "*.yaml", "*.yml"], # Example: be more specific if needed
-                # ignore_patterns=["*.git/*", ".*", "__pycache__/*"], # Example
             )
             logger.info(f"Uploaded files to Space: {repo_id}")
             return f"Successfully created/updated Space: [{repo_id}](https://huggingface.co/spaces/{repo_id})"
 
     except Exception as e:
-        logger.exception(f"Error in create_space for {repo_id if 'repo_id' in locals() else 'unknown repo'}:") # Log full traceback
+        logger.exception(f"Error in create_space for {repo_id_for_error_logging}:")
         return f"Error during Space creation/update: {str(e)}"
 
-# Function to view Space files
-def view_space_files(api_token, space_name_ui, owner_ui):
-    """List files in a Hugging Face Space."""
-    repo_id_for_error = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
+
+def view_space_files(ui_api_token_from_textbox, space_name_ui, owner_ui):
+    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
-        if not api_token:
-            return "Error: Please provide a valid Hugging Face API token."
+        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
+        if token_err:
+            return token_err
         
-        repo_id, err = _determine_repo_id(api_token, space_name_ui, owner_ui)
+        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
         if err:
             return err
-        
-        files = list_repo_files(repo_id=repo_id, token=api_token, repo_type="space") # Correctly set
+        repo_id_for_error_logging = repo_id
+            
+        files = list_repo_files(repo_id=repo_id, token=resolved_api_token, repo_type="space")
         if files:
             return f"Files in `{repo_id}`:\n\n" + "\n".join([f"- `{f}`" for f in files])
-        else:
-            return f"No files found in the Space `{repo_id}`."
+        return f"No files found in the Space `{repo_id}`."
     except Exception as e:
-        logger.exception(f"Error in view_space_files for {repo_id_for_error}:")
-        return f"Error listing files for `{repo_id_for_error}`: {str(e)}"
+        logger.exception(f"Error in view_space_files for {repo_id_for_error_logging}:")
+        return f"Error listing files for `{repo_id_for_error_logging}`: {str(e)}"
+
 
-# Function to update a Space file
-def update_space_file(api_token, space_name_ui, owner_ui, file_path_in_repo, file_content, commit_message_ui):
-    """Update a file in a Hugging Face Space with a commit."""
-    repo_id_for_error = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
+def update_space_file(ui_api_token_from_textbox, space_name_ui, owner_ui, file_path_in_repo, file_content, commit_message_ui):
+    repo_id_for_error_logging = f"{owner_ui}/{space_name_ui}" if owner_ui else space_name_ui
     try:
-        if not api_token:
-            return "Error: Please provide a valid Hugging Face API token."
+        resolved_api_token, token_err = _get_api_token(ui_api_token_from_textbox)
+        if token_err:
+            return token_err
 
-        repo_id, err = _determine_repo_id(api_token, space_name_ui, owner_ui)
+        repo_id, err = _determine_repo_id(ui_api_token_from_textbox, space_name_ui, owner_ui)
         if err:
             return err
+        repo_id_for_error_logging = repo_id
         
-        if not file_path_in_repo:
-            return "Error: File Path cannot be empty."
-        if not commit_message_ui:
-            commit_message_ui = f"Update {file_path_in_repo} via Space Builder"
+        if not file_path_in_repo: return "Error: File Path cannot be empty."
+        commit_message_ui = commit_message_ui or f"Update {file_path_in_repo} via Space Builder"
 
         with tempfile.TemporaryDirectory() as temp_dir:
-            repo_local_clone_path = Path(temp_dir) / "cloned_space_repo"
+            repo_local_clone_path = Path(temp_dir) / "cloned_space_repo_for_update"
             
             cloned_repo = Repository(
-                local_dir=str(repo_local_clone_path), # Repository expects string path
-                clone_from=f"https://huggingface.co/spaces/{repo_id}", # Ensure this URL is correct
-                repo_type="space", # Correctly set
-                use_auth_token=api_token,
-                git_user="Space Builder Bot", 
-                git_email="[email protected]"
+                local_dir=str(repo_local_clone_path),
+                clone_from=f"https://huggingface.co/spaces/{repo_id}",
+                repo_type="space", use_auth_token=resolved_api_token,
+                git_user="Space Builder Bot", git_email="[email protected]"
             )
             logger.info(f"Cloned Space {repo_id} to {repo_local_clone_path}")
 
-            full_local_file_path = cloned_repo.local_dir / file_path_in_repo # Path object arithmetic
+            full_local_file_path = Path(cloned_repo.local_dir) / file_path_in_repo
             full_local_file_path.parent.mkdir(parents=True, exist_ok=True)
-            with open(full_local_file_path, "w", encoding="utf-8") as f:
-                f.write(file_content)
+            with open(full_local_file_path, "w", encoding="utf-8") as f: f.write(file_content)
             logger.info(f"Wrote updated file {file_path_in_repo} locally.")
 
             cloned_repo.push_to_hub(commit_message=commit_message_ui)
@@ -268,5 +337,5 @@ def update_space_file(api_token, space_name_ui, owner_ui, file_path_in_repo, fil
             return f"Successfully updated `{file_path_in_repo}` in Space [{repo_id}](https://huggingface.co/spaces/{repo_id})"
 
     except Exception as e:
-        logger.exception(f"Error in update_space_file for {repo_id_for_error}:")
-        return f"Error updating file for `{repo_id_for_error}`: {str(e)}"
+        logger.exception(f"Error in update_space_file for {repo_id_for_error_logging}:")
+        return f"Error updating file for `{repo_id_for_error_logging}`: {str(e)}"