Update app.py

app.py

@@ -3,13 +3,13 @@ import google.generativeai as genai
 import re
 import datetime
 import os
-import subprocess
+import ipaddress
 
 # Secret key and Google Gemini API configuration
 API_KEY = st.secrets["GOOGLE_API_KEY"]
 
 # Page configuration
-st.set_page_config(page_title="☁️ Gemini2 Terraform Gen Pro", page_icon="☁️", layout="wide")
+st.set_page_config(page_title="⚙️ Gemini2 Ansible Gen Pro", page_icon="⚙️", layout="wide")
 
 # --- Helper Functions ---
 def send_message_to_model(message, model_name, temperature, top_p, top_k, max_tokens):
@@ -34,214 +34,120 @@ def send_message_to_model(message, model_name, temperature, top_p, top_k, max_to
         st.error(f"❌ Error communicating with the AI: {e}")
         return None
 
-def generate_terraform_code(prompt_base, detail_level, provider, security_level, model_name, temperature, top_p, top_k, max_tokens, prompt_detail, encoding, add_header, log_level, infrastructure_type, custom_requirements, specific_details):
-    """Generates a Terraform code based on user settings."""
+def generate_ansible_playbook(prompt_base, detail_level, os_type, security_level, model_name, temperature, top_p, top_k, max_tokens, prompt_detail, encoding, add_header, log_level, ansible_user, ansible_ssh_pass, ansible_become, custom_requirements, specific_details, target_hosts):
+    """Generates an Ansible playbook based on user settings."""
     prompt = f"""
-    You are an expert DevOps engineer specializing in Terraform. Your task is to generate a complete, secure, and efficient Terraform configuration based on the following description:
+    You are an expert Ansible automation engineer. Your task is to generate a complete, secure, and efficient Ansible playbook based on the following description:
 
-    **Goal:** Create the most complete, detailed, efficient, and secure Terraform configuration possible, considering all variables, edge cases, and potential scenarios.
+    **Goal:** Create the most complete, detailed, efficient, and secure Ansible playbook possible, considering all variables, edge cases, and potential scenarios.
 
-    **Configuration Description:** {prompt_base}
+    **Playbook Description:** {prompt_base}
 
     **Detail Level:** {detail_level}
     **Security Level:** {security_level}
     **Prompt Detail Level**:{prompt_detail}
-    **Infrastructure Type:** {infrastructure_type}
-    **Provider:** {provider}
+    **Operating System:** {os_type}
+    **Target Hosts:** {target_hosts if target_hosts else "None"}
+    **Ansible User:** {ansible_user if ansible_user else "root"}
+    **Ansible SSH Password:** {ansible_ssh_pass if ansible_ssh_pass else "None"}
+     **Ansible Become:** {ansible_become}
     **Custom Requirements:** {custom_requirements if custom_requirements else "None"}
     **Specific Details:** {specific_details if specific_details else "None"}
 
     **Response Format:**
-    - Respond in Markdown format, including a HCL code block with its original formatting, without line breaks.
-    - The HCL code block must be delimited by ```terraform and ```.
-    - Do not include comments, explanations, or any other text outside the code block.
+    - Respond in Markdown format, including a YAML code block with its original formatting, without line breaks.
+    - The YAML code block must be delimited by ```yaml and ```.
+    - Do not include comments, explanations, or any other text outside the code block, unless the prompt requires.
     - The 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, always prioritizing security and efficiency.
-    - Use advanced Terraform resources such as modules, loops, variables, functions and provisioners when necessary.
-    - Unless the user specifies otherwise, use the most current and secure versions of the resources, using terraform and following best practices.
-    - If the description asks to create a file, the terraform code 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 terraform code should read the file directly from the file system and not use a screen input for this.
+    - Use advanced Ansible resources such as modules, roles, variables, functions, conditionals, loops, and handlers, when necessary.
+    - Unless the user specifies otherwise, use the most current and secure versions of the modules, using ansible and following best practices.
+    - If the description asks to create a file, the ansible playbook 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 ansible playbook should read the file directly from the file system and not use a screen input for this.
+     - If possible, use Jinja2 templates when required and do not create hard code in your code
     - Use incremental reasoning to add improvements, expansions, and considerations to your code.
     - Use the history of the conversations so that the response is incremental.
 
     **Log Level**:{log_level}
 
     **Important:**
-    - Generate only one configuration at a time.
+    - Generate only one playbook 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 configuration.
-    - Use contextual information (such as the provider and infrastructure type) to generate the code.
-    - If possible, use incremental reasoning to add improvements, expansions, and considerations to your code.
-    - If the prompt asks to manage a service or process, consider different init systems (systemd, init, etc.).
-    - If the user asks to create something using modules, use official modules or public repositories, and if it is not possible, create a module with all parameters that it can have.
-    - If the description includes networking, create a secure and complete network infrastructure with all the necessary parameters, avoiding defaults.
-    - If the description includes storage, create secure and complete storage configurations, using all the available options from the provider.
-     - If the description includes compute resources, create the most secure and complete setup using the provider options.
-
-Always create cloud infrastructure with the following considerations:
-
-Security:
-
-Implement strict security groups/firewalls allowing only necessary ports.
-Configure IAM roles and policies with the principle of least privilege.
-Enable logging and monitoring for all resources.
-Networking:
-
-Set up a Virtual Private Cloud (VPC) or equivalent, with public and private subnets.
-Use NAT gateways for internet access in private subnets.
-Configure proper routing tables and internet gateways.
-Automation:
-
-Use Infrastructure as Code (IaC) tools like Terraform, CloudFormation, or ARM templates.
-Include automation for updates, scaling, and backups.
-Resilience:
-
-Ensure resources are highly available, using multiple availability zones where possible.
-Configure automatic backups and disaster recovery plans.
-Resource Optimization:
-
-Choose appropriate instance types and sizes based on the workload.
-Enable resource tagging for cost management and organization.
-Specific Instructions for Cloud Providers:
-
-For AWS:
-
-Use AWS services like EC2, S3, RDS, and CloudWatch.
-Always select the latest AMI for operating systems.
-Configure CloudTrail for auditing.
-Use Auto Scaling Groups for EC2 instances to ensure high availability.
-For Google Cloud:
-
-Use Google Compute Engine for VMs, Cloud Storage, and Cloud Monitoring.
-Set up VPC networks with subnets in multiple regions for high availability.
-Use Google Cloud IAM for access management.
-For Azure:
-
-Utilize Azure Virtual Machines, Azure Storage, and Azure Monitor.
-Set up Azure Virtual Networks with subnets and NSGs.
-Use Azure Role-Based Access Control (RBAC) for permissions.
-Never forget to:
-
-Configure proper logging and monitoring for observability.
-Include alerts and notifications for critical resource changes.
-Use version control for all IaC files.
-Ensure compliance with the cloud provider’s best practices and guidelines.
-
-When creating cloud infrastructure, always adhere to state-of-the-art practices across different cloud providers (AWS, Google Cloud, Azure). Your goal is to ensure the architecture is secure, scalable, resilient, and cost-optimized while maintaining high automation and compliance standards. Follow these detailed instructions:
-
-General Best Practices (Applicable to All Providers):
-Security:
-
-Implement multi-layered security:
-Use security groups/firewalls to restrict inbound and outbound traffic to the minimum necessary.
-Enforce encryption for data at rest (e.g., EBS volumes, databases) and in transit (e.g., HTTPS, TLS).
-Use secrets management services for sensitive data (e.g., AWS Secrets Manager, Google Secret Manager, Azure Key Vault).
-IAM Policies:
-Define and apply the principle of least privilege (PoLP) to all roles and users.
-Use separate accounts/projects/subscriptions for different environments (e.g., development, testing, production).
-Monitoring and Logging:
-Enable centralized logging (e.g., CloudWatch Logs, Stackdriver Logging, Azure Monitor).
-Set up alerts for unusual activities (e.g., unauthorized access attempts).
-Networking:
-
-Create VPCs or equivalent for isolated networking environments.
-Design subnets (public/private) based on the application’s accessibility needs.
-Use NAT Gateways for private subnet instances to access the internet securely.
-Apply Network Access Control Lists (ACLs) to provide an additional layer of security.
-Implement VPC Peering, VPN, or Direct Connect for secure on-premises connectivity.
-Automation and IaC:
-
-Use Terraform for declaring all cloud resources, ensuring infrastructure is version-controlled and repeatable.
-Apply modular design in Terraform for reusability and maintainability.
-Use remote state storage (e.g., S3 with state locking via DynamoDB, Google Cloud Storage with versioning) for state management.
-Integrate CI/CD pipelines for automatic deployment and validation of infrastructure code.
-Resilience and High Availability:
-
-Design for multi-region and multi-availability zone deployments for critical applications.
-Implement load balancers (e.g., ALB, Cloud Load Balancer, Azure Load Balancer) to distribute traffic across multiple instances.
-Use Auto Scaling Groups or equivalent to handle variable traffic loads automatically.
-Set up disaster recovery strategies, including backups and failover mechanisms.
-Cost Optimization:
-
-Tag all resources for cost tracking and management.
-Use Reserved Instances or Savings Plans for predictable workloads.
-Enable idle resource management (e.g., stopping unused instances, rightsizing underutilized resources).
-Compliance and Governance:
-
-Implement policies and guardrails using services like AWS Config, Google Organization Policies, and Azure Policy.
-Use audit logs to ensure traceability of actions.
-Regularly conduct compliance checks and audits.
-AWS-Specific Best Practices:
-Compute:
-Use EC2 Auto Scaling Groups for dynamic scaling.
-Leverage Lambda for serverless applications to reduce costs and increase scalability.
-Storage:
-Use S3 with lifecycle policies for automatic transition of data between storage classes.
-Enable S3 Bucket Versioning and MFA Delete for data protection.
-Database:
-Use RDS Multi-AZ for high availability.
-Leverage Aurora for managed relational databases with built-in replication.
-Networking:
-Configure VPC Flow Logs for network traffic analysis.
-Use CloudFront for CDN and enhanced security with AWS Shield.
-Google Cloud-Specific Best Practices:
-Compute:
-Use Managed Instance Groups for high availability and scaling.
-Leverage Cloud Functions or Cloud Run for serverless deployments.
-Storage:
-Enable Bucket Versioning and use Nearline/Coldline storage classes for cost efficiency.
-Set up Data Loss Prevention (DLP) to safeguard sensitive data.
-Networking:
-Use Cloud NAT for private instance internet access.
-Implement VPC Service Controls for enhanced data security.
-Monitoring:
-Set up Stackdriver Monitoring with custom metrics and dashboards.
-Azure-Specific Best Practices:
-Compute:
-Use Virtual Machine Scale Sets for scalable VMs.
-Leverage Azure Functions for event-driven, serverless applications.
-Storage:
-Use Azure Blob Storage with Soft Delete enabled for accidental recovery.
-Implement Storage Account Firewalls and VNET Service Endpoints.
-Networking:
-Configure NSGs (Network Security Groups) with detailed inbound/outbound rules.
-Use Azure Bastion for secure RDP/SSH connectivity.
-Monitoring:
-Enable Azure Monitor with Log Analytics for resource insights.
-Set up Application Insights for performance monitoring.
+     - Use the `delegate_to`, `local_action` and `register` options when necessary.
+     - Use the `with_items` and `with_dict` options when necessary.
+     - Do not use `sudo` unless absolutely required, prefer `become`
+     -  If possible, use `ansible.builtin` modules instead of other modules.
 
     """
     response = send_message_to_model(prompt, model_name, temperature, top_p, top_k, max_tokens)
     return response
 
-def parse_and_save_tf(ai_code, short_title, encoding, add_header):
-    """Parses the markdown and saves the terraform code as .tf."""
-    match = re.search(r'```terraform\s*(.*?)\s*```', ai_code, re.DOTALL | re.IGNORECASE)
+def parse_and_save_yml(ai_code, short_title, encoding, add_header):
+    """Parses the markdown and saves the ansible playbook as .yml."""
+    match = re.search(r'```yaml\s*(.*?)\s*```', ai_code, re.DOTALL | re.IGNORECASE)
     if match:
-        tf_code = match.group(1).strip()
+        yml_code = match.group(1).strip()
     else:
-        tf_code = ai_code.strip()
+        yml_code = ai_code.strip()
 
-    file_name = f"terraform_{short_title}.tf"
+    file_name = f"ansible_{short_title}.yml"
 
     if add_header:
         current_date = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         header = f"""#===============================================================================
-# Terraform Configuration Generated by Google Gemini 2 Terraform Gen Pro
+# Ansible Playbook Generated by Google Gemini 2 Ansible Gen Pro
 # Date: {current_date}
 # Author: Elias Andrade AKA Chaos4455
 #===============================================================================\n"""
-        tf_code = header + tf_code
+        yml_code = header + yml_code
 
     with open(file_name, "w", encoding=encoding) as f:
-        f.write(tf_code)
+        f.write(yml_code)
         
-    return file_name, tf_code
+    return file_name, yml_code
+
+def validate_ip_input(ip_input):
+    """Validates and parses IP address input."""
+    try:
+        if "," in ip_input:
+            ips = [ip.strip() for ip in ip_input.split(",")]
+            for ip in ips:
+                 if "/" in ip:
+                    ipaddress.ip_network(ip, strict=False) #verify if is a network
+                 elif "-" in ip:
+                      parts = ip.split("-")
+                      if len(parts) != 2:
+                          return False, "Invalid Range Format"
+                      
+                      ipaddress.ip_address(parts[0])
+                      ipaddress.ip_address(parts[1])
+                 else:
+                     ipaddress.ip_address(ip)  #verify if is a valid address
+            return True, ips
+        elif "/" in ip_input:
+              ipaddress.ip_network(ip_input, strict=False) #verify if is a network
+              return True, [ip_input]
+        elif "-" in ip_input:
+              parts = ip_input.split("-")
+              if len(parts) != 2:
+                  return False, "Invalid Range Format"
+              
+              ipaddress.ip_address(parts[0])
+              ipaddress.ip_address(parts[1])
+              return True,[ip_input]
+        else:
+            ipaddress.ip_address(ip_input)
+            return True, [ip_input]
+    except ValueError as e:
+        return False, str(e)
+    except Exception as e:
+        return False, str(e)
 
 def main():
-    st.title("☁️ Gemini2 Terraform Gen Pro by [Elias Andrade](https://github.com/chaos4455)")
-    st.markdown("Generate advanced Terraform configurations with ease! 🚀")
+    st.title("⚙️ Gemini2 Ansible Gen Pro by [Elias Andrade](https://github.com/chaos4455)")
+    st.markdown("Generate advanced Ansible playbooks with ease! 🚀")
     st.markdown("---")
 
     # Layout in columns (sidebar and main area)
@@ -258,53 +164,57 @@ def main():
             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", "Create a VPC", "Create an EC2 instance", "Create a Kubernetes cluster", "Create a Docker container", "Create a storage account", "Create a virtual machine"], index=0, help="Choose a predefined prompt.")
+            prompt_presets = st.selectbox("🎯 Predefined Prompts", ["None", "Install packages", "Configure services", "Manage users", "File management", "System updates", "Deploy applications"], 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("☁️ Terraform Settings"):
-            provider = st.selectbox(
-                "☁️ Provider",
-                [
-                 "aws", "azurerm", "google", "kubernetes", "docker", "vmware"
-                ],
-                index=0, help="Choose the target provider."
-            )
-            infrastructure_type = st.selectbox(
-                "🏗️ Infrastructure Type",
+        with st.expander("⚙️ Ansible Settings"):
+            os_type = st.selectbox(
+                "💻 Operating System",
                 [
-                    "Networking", "Compute", "Storage", "Database", "Containers", "Virtual Machines", "Functions", "Other"
+                    "Ubuntu", "CentOS", "RedHat", "Debian", "Windows", "Other"
                 ],
-                index=0, help="Choose the infrastructure type."
+                index=0, help="Choose the target operating system."
             )
-            encoding = st.selectbox("🔤 Encoding", ["utf-8", "ansi"], index=0, help="Choose the encoding of the .tf file.")
-            add_header = st.checkbox("📜 Add Header", value=True, help="Add a header with information in the .tf file.")
+            
+            ansible_user = st.text_input("👤 Ansible User", value="root", help="User for Ansible SSH connection")
+            ansible_ssh_pass = st.text_input("🔑 Ansible SSH Password", type="password", help="Password for Ansible SSH connection (Use with caution)")
+            ansible_become = st.checkbox("🔥 Use Become", value=True, help="Enable privilege escalation")
+
+            encoding = st.selectbox("🔤 Encoding", ["utf-8", "ansi"], index=0, help="Choose the encoding of the .yml file.")
+            add_header = st.checkbox("📜 Add Header", value=True, help="Add a header with information in the .yml file.")
             log_level = st.selectbox("🗂️ Logging Level", ["Detailed", "Default", "Minimum"], index=1, help="Defines the detail level of logs.")
             security_level = st.radio("Security Level", ["High", "Medium", "Low"], index=1)
-            custom_requirements = st.text_input("Custom Requirements", placeholder="Ex: Use specific modules, parameters", help="Add specific requirements for the generated configuration.")
-            specific_details = st.text_input("Specific Details", placeholder="Ex: Specific configurations, edge cases", help="Add specific details for configuration generation.")
+            custom_requirements = st.text_input("Custom Requirements", placeholder="Ex: Use specific modules, parameters", help="Add specific requirements for the generated playbook.")
+            specific_details = st.text_input("Specific Details", placeholder="Ex: Specific configurations, edge cases", help="Add specific details for playbook generation.")
             
             detail_level = st.selectbox("Detail Level", ["More detailed", "Default", "More concise"], index=1)
 
 
     with col2:
         # User's base prompt
-        prompt_base = st.text_input("Describe the Terraform Configuration:", placeholder="Ex: Create a VPC with public and private subnets in AWS", key="prompt_base")
+        prompt_base = st.text_input("Describe the Ansible Playbook:", placeholder="Ex: Install apache2 and configure a default page for all the servers", key="prompt_base")
+        target_hosts = st.text_input("Target Hosts (IP(s), Range or CIDR, separated by comma):", placeholder="192.168.0.1,192.168.0.2,192.168.0.1/24 or 192.168.0.1-255",key="target_hosts")
 
         if prompt_presets != "None":
             if prompt_base:
                prompt_base = f"{prompt_presets} , {prompt_base}"
             else:
                prompt_base = prompt_presets;
-        if st.button("✨ Generate Terraform Configuration"):
+        if st.button("✨ Generate Ansible Playbook"):
             if not prompt_base:
-                st.error("⚠️ Please enter a configuration description.")
+                st.error("⚠️ Please enter a playbook description.")
                 return
+            
+            valid_ip, parsed_ips = validate_ip_input(target_hosts)
+            if not valid_ip:
+               st.error(f"⚠️ Invalid IP Address or Range: {parsed_ips}")
+               return
 
-            with st.spinner("⏳ Generating configuration..."):
-                ai_code = generate_terraform_code(
+            with st.spinner("⏳ Generating playbook..."):
+                ai_code = generate_ansible_playbook(
                    prompt_base, 
                    detail_level, 
-                   provider, 
+                   os_type,
                    security_level,
                    model_name,
                    temperature,
@@ -315,26 +225,29 @@ def main():
                    encoding,
                    add_header,
                    log_level,
-                   infrastructure_type,
+                   ansible_user,
+                   ansible_ssh_pass,
+                   ansible_become,
                    custom_requirements,
-                   specific_details
+                   specific_details,
+                   parsed_ips if parsed_ips else None
                    )
 
                 if ai_code:
-                    st.markdown("### ✅ Generated Configuration:")
-                    st.code(ai_code, language="terraform")
+                    st.markdown("### ✅ Generated Playbook:")
+                    st.code(ai_code, language="yaml")
 
                     short_title = prompt_base[:30].strip().replace(" ", "_").lower()
-                    file_name_tf, tf_code = parse_and_save_tf(ai_code, short_title, encoding, add_header)
+                    file_name_yml, yml_code = parse_and_save_yml(ai_code, short_title, encoding, add_header)
 
                     st.download_button(
-                        label="⬇️ Download Configuration (.tf)",
-                        data=tf_code,
-                        file_name=file_name_tf,
-                        mime="application/x-tf",
+                        label="⬇️ Download Playbook (.yml)",
+                        data=yml_code,
+                        file_name=file_name_yml,
+                        mime="application/x-yaml",
                     )
                 else:
-                    st.error("❌ Error generating the configuration. Check the connection with the AI and try again.")
+                    st.error("❌ Error generating the playbook. Check the connection with the AI and try again.")
 
 if __name__ == "__main__":
     main()