#!/usr/bin/env python3
import argparse
import json
import re
import sys
from typing import Dict, List, Tuple

BLOCKED_RULES: List[Tuple[str, str]] = [
    (r'rm\s+-rf\s+(/|/\*|~|~/\*)', 'Recursive forced deletion targeting root or home filesystem'),
    (r'mkfs\s+/dev/', 'Filesystem format operation on block device'),
    (r'dd\s+.*if=/dev/(zero|urandom).*of=/dev/', 'Disk overwrite operation on block device'),
    (r':\(\)\s*{\s*:\|:\s*&\s*};\s*:', 'Fork bomb command'),
    (r'chmod\s+-R\s+(777|000)\s+(/|/\*)', 'Mass permission modification on root filesystem'),
    (r'>\s*/dev/(sd|hd|nvme)', 'Direct write to block device'),
    (r'(curl|wget).*\|\s*(bash|sh|zsh)', 'Remote code execution via pipe'),
    (r'echo.*base64\s+-d.*\|\s*(bash|sh|zsh)', 'Encoded payload execution'),
    (r'(shutdown|reboot|halt|poweroff)', 'System power control command'),
    (r'mv\s+/\s+/dev/null', 'Root filesystem redirection to null device'),
    (r'--no-preserve-root', 'Root deletion without preservation flag'),
    (r'history\s+-c.*rm.*\.bash_history', 'Audit trail destruction'),
    (r'printf\s+.*\\x72\\x6d\\x20\\x2d\\x72\\x66\\x20\\x2f', 'Hex-encoded root deletion'),
    (r'CMD\s*=\s*["\']rm["\'].*ARGS\s*=\s*["\']-rf /["\']', 'Variable-based root deletion evasion'),
    (r'alias\s+.*=\s*["\']rm -rf /["\']', 'Alias-based root deletion evasion'),
    (r'echo\s+["\']rm["\'].*xargs.*-rf\s+/', 'Pipe-based root deletion evasion'),
]

CONFIRM_RULES: List[Tuple[str, str]] = [
    (r'rm\s+-r?f?\s+[^/]+', 'Recursive file deletion'),
    (r'sudo\s+', 'Privilege escalation command'),
    (r'kill\s+-9\s+\d+', 'Force process termination'),
    (r'killall\s+\w+', 'Mass process termination by name'),
    (r'(chmod|chown|chattr)\s+', 'File permission/attribute modification'),
    (r'systemctl\s+(stop|disable|restart)', 'System service management'),
    (r'(apt|brew|pip)\s+(remove|uninstall)', 'Package removal'),
    (r'pip\s+install.*--user', 'Global Python package installation'),
    (r'npm\s+install\s+-g', 'Global NPM package installation'),
    (r'git\s+push\s+--force', 'Forced git push'),
    (r'git\s+reset\s+--hard', 'Destructive git reset'),
    (r'(DROP|DELETE|TRUNCATE)\s+TABLE', 'Destructive SQL operation'),
    (r'(iptables|ufw)\s+', 'Firewall configuration change'),
    (r'crontab\s+-r', 'Crontab removal'),
    (r'docker\s+(rm|rmi)\s+', 'Docker container/image removal'),
    (r'(mount|umount)\s+', 'Filesystem mount operation'),
    (r'export\s+\w*(KEY|TOKEN|PASSWORD|SECRET)', 'Sensitive environment variable export'),
    (r'(curl|wget)\s+.*-X\s+POST', 'HTTP POST request to external URL'),
    (r'eval\s+\$\(', 'Dynamic code evaluation'),
]

ALLOWED_PREFIXES = [
    'ls', 'cat', 'head', 'tail', 'less', 'more', 'echo', 'printf', 'pwd', 'whoami',
    'hostname', 'grep', 'awk', 'sed', 'find', 'cd', 'pushd', 'popd', 'cp', 'mkdir',
    'touch', 'python3', 'node', 'npm run', 'npm test', 'npm start', 'pip install',
    'git status', 'git log', 'git diff', 'git branch', 'git add', 'git commit',
    'git pull', 'env', 'printenv', 'wc', 'sort', 'uniq', 'cut', 'tr', 'date', 'cal',
    'uptime', 'tree'
]

def classify_command(command: str) -> Dict:
    command = command.strip()
    
    for pattern, reason in BLOCKED_RULES:
        if re.search(pattern, command, re.IGNORECASE):
            return {
                'verdict': 'BLOCKED',
                'risk_level': 'critical',
                'reason': reason,
                'matched_rules': [pattern],
                'suggestion': get_safe_suggestion(pattern, command)
            }
    
    for pattern, reason in CONFIRM_RULES:
        if re.search(pattern, command, re.IGNORECASE):
            return {
                'verdict': 'CONFIRM',
                'risk_level': 'medium',
                'reason': reason,
                'matched_rules': [pattern],
                'suggestion': 'Verify this is intended before execution.'
            }
    
    for prefix in ALLOWED_PREFIXES:
        if command.startswith(prefix):
            return {
                'verdict': 'ALLOWED',
                'risk_level': 'low',
                'reason': 'Command matches allowed prefix list',
                'matched_rules': [],
                'suggestion': ''
            }
    
    return {
        'verdict': 'CONFIRM',
        'risk_level': 'unknown',
        'reason': 'Command not in known safe list',
        'matched_rules': [],
        'suggestion': 'Verify this command is safe before execution.'
    }

def get_safe_suggestion(pattern: str, command: str) -> str:
    if 'rm -rf /' in pattern:
        return "Use 'rm -rf ./<specific_dir>' to target a specific directory instead."
    elif 'mkfs' in pattern:
        return "Filesystem formatting is a destructive operation. Use with extreme caution."
    elif 'dd' in pattern:
        return "Disk overwriting can cause permanent data loss. Double-check device paths."
    elif 'fork bomb' in pattern:
        return "Fork bombs crash systems by exhausting process resources."
    elif 'chmod -R' in pattern:
        return "Use chmod on specific files/directories instead of recursive root changes."
    elif '| bash' in pattern:
        return "Download and inspect scripts before executing them locally."
    else:
        return "This command is dangerous and cannot be executed."

def main():
    parser = argparse.ArgumentParser(description='Safe Command Guard - Validate system commands before execution')
    parser.add_argument('--command', help='Command string to validate')
    parser.add_argument('--output', choices=['json', 'text'], default='json', help='Output format')
    args = parser.parse_args()
    
    command = args.command
    if not command:
        command = sys.stdin.read().strip()
    
    if not command:
        result = {
            'verdict': 'CONFIRM',
            'risk_level': 'unknown',
            'reason': 'Empty command provided',
            'matched_rules': [],
            'suggestion': 'Provide a valid command to check.'
        }
    else:
        result = classify_command(command)
    
    if args.output == 'json':
        print(json.dumps(result, indent=2, ensure_ascii=False))
    else:
        verdict = result['verdict']
        icon = '🔴' if verdict == 'BLOCKED' else '🟡' if verdict == 'CONFIRM' else '🟢'
        print(f"{icon} {verdict}: {result['reason']}")
        if result['suggestion']:
            print(f"💡 Suggestion: {result['suggestion']}")

if __name__ == '__main__':
    main()