Update app.py

app.py

@@ -1,23 +1,33 @@
 import streamlit as st
 import google.generativeai as genai
-import json
 import re
 import datetime
+import os
+import subprocess
 
 # Secret key and Google Gemini API configuration
 API_KEY = st.secrets["GOOGLE_API_KEY"]
 
 # Page configuration
-st.set_page_config(page_title="🤖 Gemini2 PowerShell Command Gen", page_icon="🤖", layout="wide")
+st.set_page_config(page_title="🐧 Gemini2 Linux Command Gen Pro", page_icon="🐧", layout="wide")
 
-def send_message_to_model(message, model_name, temperature, max_tokens):
+# --- Helper Functions ---
+def check_command_availability(command):
+    """Checks if a command is available in the system's PATH."""
+    try:
+        subprocess.run(["which", command], check=True, capture_output=True)
+        return True
+    except subprocess.CalledProcessError:
+        return False
+
+def send_message_to_model(message, model_name, temperature, top_p, top_k, max_tokens):
     """Sends a message to the AI model and returns the response."""
     try:
         # AI model configurations
         GENERATION_CONFIG = {
             "temperature": temperature,
-            "top_p": 0.8,
-            "top_k": 40,
+            "top_p": top_p,
+            "top_k": top_k,
             "response_mime_type": "text/plain",
             "max_output_tokens": max_tokens,
         }
@@ -32,12 +42,12 @@ def send_message_to_model(message, model_name, temperature, max_tokens):
         st.error(f"❌ Error communicating with the AI: {e}")
         return None
 
-def generate_powershell_command(prompt_base, detail_level, script_type, security_level, model_name, temperature, max_tokens, prompt_detail, encoding, add_header, add_error_handling, log_level):
-    """Generates a PowerShell command based on user settings."""
+def generate_linux_command(prompt_base, detail_level, script_type, security_level, model_name, temperature, top_p, top_k, max_tokens, prompt_detail, encoding, add_header, add_error_handling, log_level, operating_system, command_execution_method, custom_command_requirements, command_specific_details):
+    """Generates a Linux command based on user settings."""
     prompt = f"""
-    You are a Powershell expert. Your task is to generate a single Powershell command based on the following description:
+    You are an expert Linux system administrator. Your task is to generate a single, secure and efficient Linux command or shell script based on the following description:
 
-    **Goal:** Create the most complete, detailed, and efficient Powershell command possible, considering all variables and scenarios.
+    **Goal:** Create the most complete, detailed, efficient, and secure Linux command possible, considering all variables, edge cases, and potential scenarios.
 
     **Command Description:** {prompt_base}
 
@@ -45,66 +55,78 @@ def generate_powershell_command(prompt_base, detail_level, script_type, security
     **Script Type:** {script_type}
     **Security Level:** {security_level}
     **Prompt Detail Level**:{prompt_detail}
+    **Operating System:** {operating_system}
+    **Command Execution Method:** {command_execution_method}
+    **Custom Command Requirements:** {custom_command_requirements if custom_command_requirements else "None"}
+    **Command Specific Details**:{command_specific_details if command_specific_details else "None"}
 
     **Response Format:**
-    - Respond in Markdown format, including a Powershell code block with its original formatting, without line breaks.
-    - The Powershell code block must be delimited by ```powershell and ```.
+    - Respond in Markdown format, including a bash code block with its original formatting, without line breaks.
+    - The bash code block must be delimited by ```bash and ```.
     - Do not include comments, explanations, or any other text outside the code block.
-    - The Powershell code must maintain its full vertical formatting, respecting indentation and line breaks.
+    - The bash code must maintain its full vertical formatting, respecting indentation and line breaks.
     - The code must be realistic, using real-world examples, data, and situations.
-    - Explore different approaches, techniques, and advanced practices.
-    - If generating a command with a chain of commands, do not use semicolons at the end or beginning.
-    - If necessary, use the pipe "|" to chain commands.
+    - Explore different approaches, techniques, and advanced practices, always prioritizing security and efficiency.
+    - If generating a command with a chain of commands, use "&&" to chain commands when appropriate and consider error handling within the chain.
     - Do not use any special formatting in the result, only the code.
     - If the description asks to create a file, the command should create the file directly in the file system and not use a screen output for this.
     - If the description asks to read a file, the command should read the file directly from the file system and not use a screen input for this.
-    - Use advanced PowerShell resources, such as pipelines, variables, functions, and script blocks, when necessary.
-     - The default operating system is Windows Server 2016, and the default Powershell version is 7, unless the user specifies otherwise.
-     - Make sure that the generated command is secure and follows the best Powershell practices.
+    - Use advanced Linux resources, such as pipelines, variables, functions, conditional statements and script blocks, when necessary.
+    - Unless the user specifies otherwise, use the most current and secure versions of the command, using bash and following best practices.
 
     **Log Level**:{log_level}
     **Error Handling**:{add_error_handling}
 
-     **Important:**
-    - Generate only one command at a time.
-    - Create the longest, most complete, and detailed code possible.
-    - Consider all the details of the request, expanding the response and improving the command.
-     - Use contextual information (such as PowerShell version and operating system) to generate the command.
-     - If possible, use incremental reasoning to add improvements, expansions, and considerations to your code.
-     - Use the history of the conversations so that the response is incremental.
+    **Important:**
+    - Generate only one command or script at a time.
+    - Create the longest, most complete, and detailed code possible to cover a wide range of possibilities and scenarios.
+    - Consider all the details of the request, expanding the response and improving the command or script.
+    - Use contextual information (such as the operating system) to generate the command or script.
+    - If possible, use incremental reasoning to add improvements, expansions, and considerations to your code.
+    - Use the history of the conversations so that the response is incremental.
+    - If the prompt asks to manage a service or process, consider different init systems (systemd, init, etc.).
+    - If a command is required but is not installed on the system, provide instructions to install such command.
 
     """
-    response = send_message_to_model(prompt, model_name, temperature, max_tokens)
+    response = send_message_to_model(prompt, model_name, temperature, top_p, top_k, max_tokens)
     return response
 
-def parse_and_save_ps1(ai_code, short_title, encoding, add_header):
-    """Parses the markdown and saves the Powershell code as .ps1."""
-    match = re.search(r'```powershell\s*(.*?)\s*```', ai_code, re.DOTALL | re.IGNORECASE)
+def parse_and_save_sh(ai_code, short_title, encoding, add_header, executable_permission):
+    """Parses the markdown and saves the bash code as .sh."""
+    match = re.search(r'```bash\s*(.*?)\s*```', ai_code, re.DOTALL | re.IGNORECASE)
     if match:
-        ps1_code = match.group(1).strip()
+        sh_code = match.group(1).strip()
     else:
-        ps1_code = ai_code.strip()
+        sh_code = ai_code.strip()
 
-    file_name = f"command_{short_title}.ps1"
+    file_name = f"command_{short_title}.sh"
 
     if add_header:
         current_date = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        header = f"""
+        header = f"""#!/bin/bash
 #===============================================================================
-# Script Generated by Google Gemini 2 PowerShell Command Gen
+# Script Generated by Google Gemini 2 Linux Command Gen Pro
 # Date: {current_date}
 # Author: Elias Andrade AKA Chaos4455
-#===============================================================================
-    """
-        ps1_code = header + ps1_code
+#===============================================================================\n"""
+        sh_code = header + sh_code
+        
 
     with open(file_name, "w", encoding=encoding) as f:
-        f.write(ps1_code)
-    return file_name, ps1_code
+        f.write(sh_code)
+
+    # Make the script executable if necessary
+    if executable_permission:
+        try:
+            os.chmod(file_name, 0o755)
+        except Exception as e:
+            st.warning(f"Warning: Could not set executable permission for {file_name}. Reason: {e}")
+            
+    return file_name, sh_code
 
 def main():
-    st.title("🤖 Gemini2 PowerShell Command Gen by [Elias Andrade](https://github.com/chaos4455)")
-    st.markdown("Create PowerShell commands easily and quickly! 🚀")
+    st.title("🐧 Gemini2 Linux Command Gen Pro by [Elias Andrade](https://github.com/chaos4455)")
+    st.markdown("Generate advanced Linux commands and scripts with ease! 🚀")
     st.markdown("---")
 
     # Layout in columns (sidebar and main area)
@@ -116,77 +138,90 @@ def main():
         with st.expander("✨ AI Settings"):
             model_name = st.selectbox("🤖 AI Model", ["gemini-2.0-flash-exp", "gemini-1.5-flash"], index=0, help="Choose the AI model.")
             temperature = st.slider("🌡️ Temperature", min_value=0.1, max_value=1.0, value=0.7, step=0.1, help="Adjust the AI's creativity.")
+            top_p = st.slider("Top P", min_value=0.1, max_value=1.0, value=0.8, step=0.1, help="Adjust the AI's sampling.")
+            top_k = st.slider("Top K", min_value=1, max_value=100, value=40, step=1, help="Adjust the AI's number of candidate tokens.")
             max_tokens = st.number_input("📏 Max Tokens", min_value=128, max_value=8192, value=8192, step=128, help="Adjust the maximum size of the response.")
             
         with st.expander("📝 Prompt Settings"):
-            prompt_presets = st.selectbox("🎯 Predefined Prompts", ["None", "List files", "Manage processes", "Manage services"], index=0, help="Choose a predefined prompt.")
+            prompt_presets = st.selectbox("🎯 Predefined Prompts", ["None", "List files", "Manage processes", "Manage services", "User Management", "System Information", "Network Configuration"], index=0, help="Choose a predefined prompt.")
             prompt_detail = st.selectbox("🧐 Prompt Detail", ["More descriptive", "Default", "Concise"], index=1, help="Defines the level of detail of the prompt")
 
-        with st.expander("🛠️ PowerShell Settings"):
-            powershell_version = st.selectbox("🎛️ PowerShell Version", ["7", "5.1"], index=0, help="Choose the PowerShell version.")
+        with st.expander("🐧 Linux Settings"):
             operating_system = st.selectbox(
                 "💻 Operating System",
                 [
-                    "Windows Server 2022", "Windows Server 2019", "Windows Server 2016",
-                    "Windows Server 2012 R2", "Windows Server 2012", "Windows Server 2008 R2",
-                    "Windows Server 2008", "Windows 11", "Windows 10", "Windows 8.1", "Windows 8", "Windows 7", "Other"
+                    "Ubuntu Server 22.04", "Ubuntu Server 20.04", "Ubuntu Server 18.04", "Ubuntu Server 16.04", "Ubuntu Server 14.04",
+                    "CentOS 8", "CentOS 7",
+                    "RedHat Enterprise Linux 9", "RedHat Enterprise Linux 8",
+                    "Debian 11", "Debian 10", "Debian 9", "Debian 8", "Debian 7", "Other"
                 ],
-                index=2, help="Choose the target operating system."
+                index=0, help="Choose the target operating system."
             )
-            encoding = st.selectbox("🔤 Encoding", ["utf-8", "ansi"], index=0, help="Choose the encoding of the .ps1 file.")
-            add_header = st.checkbox("📜 Add Header", value=True, help="Add a header with information in the .ps1 file.")
+            encoding = st.selectbox("🔤 Encoding", ["utf-8", "ansi"], index=0, help="Choose the encoding of the .sh file.")
+            add_header = st.checkbox("📜 Add Header", value=True, help="Add a header with information in the .sh file.")
             add_error_handling = st.checkbox("🛡️ Error Handling", value=True, help="Add standard error handling to the script.")
             log_level = st.selectbox("🗂️ Logging Level", ["Detailed", "Default", "Minimum"], index=1, help="Defines the detail level of logs.")
             
             detail_level = st.selectbox("Detail Level", ["More detailed", "Default", "More concise"], index=1)
             script_type = st.selectbox("Script Type", ["More automatic", "More interactive"], index=0)
             security_level = st.radio("Security Level", ["High", "Medium", "Low"], index=1)
+            command_execution_method = st.selectbox("Execution Method", ["Single Command", "Shell Script"], index=0, help="Defines if a single command or a shell script will be generated")
+            executable_permission = st.checkbox("Enable Execute Permission", value=False, help="Set the generated .sh file as executable.")
+            custom_command_requirements = st.text_input("Custom Command Requirements", placeholder="Ex: Use specific tools, parameters", help="Add specific requirements for the generated command.")
+            command_specific_details = st.text_input("Command Specific Details", placeholder="Ex: Specific configurations, edge cases", help="Add specific details for command generation.")
+
 
     with col2:
          # User's base prompt
-        prompt_base = st.text_input("Describe the PowerShell Command:", placeholder="Ex: List all running processes", key="prompt_base")
+        prompt_base = st.text_input("Describe the Linux Command/Script:", placeholder="Ex: List all files in /var/log recursively, including size and permissions", key="prompt_base")
 
         if prompt_presets != "None":
             if prompt_base:
                prompt_base = f"{prompt_presets} , {prompt_base}"
             else:
                prompt_base = prompt_presets;
-        if st.button("✨ Generate PowerShell Command"):
+        if st.button("✨ Generate Linux Command/Script"):
             if not prompt_base:
                 st.error("⚠️ Please enter a command description.")
                 return
 
-            with st.spinner("⏳ Generating command..."):
-                ai_code = generate_powershell_command(
+            with st.spinner("⏳ Generating command/script..."):
+                ai_code = generate_linux_command(
                    prompt_base, 
                    detail_level, 
                    script_type, 
                    security_level,
                    model_name,
                    temperature,
+                   top_p,
+                   top_k,
                    max_tokens,
                    prompt_detail,
                    encoding,
                    add_header,
                    add_error_handling,
-                   log_level
+                   log_level,
+                   operating_system,
+                   command_execution_method,
+                   custom_command_requirements,
+                   command_specific_details
                    )
 
                 if ai_code:
-                    st.markdown("### ✅ Generated Command:")
-                    st.code(ai_code, language="powershell")
+                    st.markdown("### ✅ Generated Command/Script:")
+                    st.code(ai_code, language="bash")
 
                     short_title = prompt_base[:30].strip().replace(" ", "_").lower()
-                    file_name_ps1, ps1_code = parse_and_save_ps1(ai_code, short_title, encoding, add_header)
+                    file_name_sh, sh_code = parse_and_save_sh(ai_code, short_title, encoding, add_header, executable_permission)
 
                     st.download_button(
-                        label="⬇️ Download Command (.ps1)",
-                        data=ps1_code,
-                        file_name=file_name_ps1,
-                        mime="application/powershell",
+                        label="⬇️ Download Command/Script (.sh)",
+                        data=sh_code,
+                        file_name=file_name_sh,
+                        mime="application/x-sh",
                     )
                 else:
-                    st.error("❌ Error generating the command. Check the connection with the AI and try again.")
+                    st.error("❌ Error generating the command/script. Check the connection with the AI and try again.")
 
 if __name__ == "__main__":
     main()