Use AI to generate system commands

bash-ai.py

#!/usr/bin/env python3
import os
import sys
import requests
import json
import argparse

def generate_bash_command(prompt):
    api_key = os.getenv('BASH_AI_KEY')
    api_url = os.getenv('BASH_URL_KEY', "http://localhost:11434/v1/chat/completions")

    if not api_key:
        print("Error: BASH_AI_KEY environment variable not set")
        return None

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    data = {
        "model": "gpt-3.5-turbo",
        "messages": [
            {
                "role": "system",
                "content": "You are an expert Linux/Bash system administrator. Generate only working bash commands or scripts without explanations. Do not include markdown code blocks."
            },
            {
                "role": "user",
                "content": f"Generate a bash command or script that {prompt}. Output only the command/script without any explanations."
            }
        ],
        "max_tokens": 200,
        "temperature": 0.3
    }

    try:
        response = requests.post(api_url, headers=headers, json=data)

        if response.status_code == 200:
            result = response.json()
            command = result['choices'][0]['message']['content'].strip()

            # Clean up the response to ensure it's just bash code
            lines = command.split('\n')
            clean_command = []

            for line in lines:
                if not line.startswith('```') and not line.startswith('bash'):
                    clean_command.append(line)

            return '\n'.join(clean_command).strip()
        else:
            print(f"Error: {response.status_code} - {response.text}")
            return None

    except Exception as e:
        print(f"Error making request: {e}")
        return None

def main():
    parser = argparse.ArgumentParser(description="Generate bash commands from natural language")
    parser.add_argument("prompt", nargs="*", help="Prompt for the bash command (or use -f to read from file)")
    parser.add_argument("-f", "--file", help="Read prompt from file")
    parser.add_argument("-o", "--output", help="Save output to executable script file")
    parser.add_argument("-e", "--execute", action="store_true", help="Execute the generated command immediately")

    args = parser.parse_args()

    # Get prompt from either command line or file
    if args.file:
        try:
            with open(args.file, 'r') as f:
                prompt = f.read().strip()
        except FileNotFoundError:
            print(f"Error: File '{args.file}' not found")
            return
        except Exception as e:
            print(f"Error reading file: {e}")
            return
    else:
        if not args.prompt:
            print("Usage: bash-ai \"description of what you want to do\"")
            print("Example: bash-ai \"rename all files skill.md to rules.md recursively\"")
            return
        prompt = " ".join(args.prompt)

    print(f"Generating bash command for: {prompt}")
    print("-" * 50)

    command = generate_bash_command(prompt)

    if command:
        print("Generated bash command:")
        print(command)
        print("-" * 50)

        # Save to file if requested
        if args.output:
            try:
                with open(args.output, 'w') as f:
                    f.write(command)
                # Make it executable
                os.chmod(args.output, 0o755)
                print(f"Command saved to {args.output} and made executable")
            except Exception as e:
                print(f"Error saving file: {e}")

        # Execute immediately if requested
        if args.execute:
            print("Executing...")
            os.system(command)
        elif not args.output:
            # Ask user if they want to execute it (only if not saving to file)
            confirm = input("Do you want to execute this command? (y/N): ")
            if confirm.lower() == 'y':
                print("Executing...")
                os.system(command)
    else:
        print("Failed to generate command")

if __name__ == "__main__":
    main()