app.py

import streamlit as st
import google.generativeai as genai
import re
import datetime
import os
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 Ansible Gen Pro", page_icon="⚙️", layout="wide")

# --- Helper Functions ---
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": top_p,
            "top_k": top_k,
            "response_mime_type": "text/plain",
            "max_output_tokens": max_tokens,
        }
        MODEL = genai.GenerativeModel(
            model_name=model_name,
            generation_config=GENERATION_CONFIG,
        )
        
        response = MODEL.start_chat(history=[]).send_message(message)
        return response.text
    except Exception as e:
        st.error(f"❌ Error communicating with the AI: {e}")
        return None

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_description):
    """Generates an Ansible playbook based on user settings."""
    prompt = f"""
    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 Ansible playbook possible, considering all variables, edge cases, and potential scenarios.

    **Playbook Description:** {prompt_base}

    **Detail Level:** {detail_level}
    **Security Level:** {security_level}
    **Prompt Detail Level**:{prompt_detail}
    **Operating System:** {os_type}
    **Target Hosts:** {target_hosts_description if target_hosts_description 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 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.
     - The `hosts` section of the playbook, must include the list of the target hosts in the description, if the target hosts are not specified you must create your own hosts list, please avoid using "all" if the target hosts are specified.
    - Explore different approaches, techniques, and advanced practices, always prioritizing security and efficiency.
    - 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 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 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_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:
        yml_code = match.group(1).strip()
    else:
        yml_code = ai_code.strip()

    file_name = f"ansible_{short_title}.yml"

    if add_header:
        current_date = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        header = f"""#===============================================================================
# Ansible Playbook Generated by Google Gemini 2 Ansible Gen Pro
# Date: {current_date}
# Author: Elias Andrade AKA Chaos4455
#===============================================================================\n"""
        yml_code = header + yml_code

    with open(file_name, "w", encoding=encoding) as f:
        f.write(yml_code)
        
    return file_name, yml_code


def main():
    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)
    col1, col2 = st.columns([1, 3])
    
    with col1:
        st.header("⚙️ Settings")
        
        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", "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("⚙️ Ansible Settings"):
            os_type = st.selectbox(
                "💻 Operating System",
                [
                    "Ubuntu", "CentOS", "RedHat", "Debian", "Windows", "Other"
                ],
                index=0, help="Choose the target operating system."
            )
            
            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 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 Ansible Playbook:", placeholder="Ex: Install apache2 and configure a default page for all the servers", key="prompt_base")
        target_hosts_description = st.text_input("Target Hosts Description (text format, IP(s), Range or CIDR):", placeholder="Ex: all the servers in 192.168.0.0/24 network",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 Ansible Playbook"):
            if not prompt_base:
                st.error("⚠️ Please enter a playbook description.")
                return

            with st.spinner("⏳ Generating playbook..."):
                ai_code = 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_description # Pass only the description
                   )

                if ai_code:
                    st.markdown("### ✅ Generated Playbook:")
                    st.code(ai_code, language="yaml")

                    short_title = prompt_base[:30].strip().replace(" ", "_").lower()
                    file_name_yml, yml_code = parse_and_save_yml(ai_code, short_title, encoding, add_header)

                    st.download_button(
                        label="⬇️ Download Playbook (.yml)",
                        data=yml_code,
                        file_name=file_name_yml,
                        mime="application/x-yaml",
                    )
                else:
                    st.error("❌ Error generating the playbook. Check the connection with the AI and try again.")

if __name__ == "__main__":
    main()