# docker-compose.security.yml - Log collection services
version: '3.8'
services:
  loki:
    image: grafana/loki:3.0
    ports:
      - "3100:3100"
    volumes:
      - ./loki/config.yaml:/etc/loki/local-config.yaml
      - loki-data:/loki
    command: -config.file=/etc/loki/local-config.yaml

  promtail:
    image: grafana/promtail:3.0
    volumes:
      - /var/log:/var/log
      - /var/lib/docker/containers:/var/lib/docker/containers:ro
      - ./promtail/config.yaml:/etc/promtail/config.yaml
    command: -config.file=/etc/promtail/config.yaml

  # AI log analysis agent
  security-agent:
    build: ./security-agent
    environment:
      - LLM_ENDPOINT=http://ollama:11434
      - MODEL_NAME=mistral
      - LOGI_ENDPOINT=http://loki:3100
    volumes:
      - ./security-agent/rules:/rules
    depends_on:
      - loki
      - ollama

#!/bin/bash
# security-agent/scripts/config_snapshot.sh
# Collect key VPS config files for AI auditing

SNAPSHOT_DIR="/var/security/snapshots/$(date +%Y%m%d_%H%M%S)"
mkdir -p "$SNAPSHOT_DIR"

echo "[*] Collecting system configuration..."
# Users and permissions
cp /etc/passwd "$SNAPSHOT_DIR/"
cp /etc/shadow "$SNAPSHOT_DIR/" 2>/dev/null || true
cp /etc/sudoers "$SNAPSHOT_DIR/"
getent group sudo wheel docker admin > "$SNAPSHOT_DIR/groups.txt" 2>/dev/null

# SSH configuration
ssh-keygen -lf /etc/ssh/sshd_config 2>/dev/null
cp /etc/ssh/sshd_config "$SNAPSHOT_DIR/"
ls -la /root/.ssh/ > "$SNAPSHOT_DIR/root_ssh_keys.txt" 2>/dev/null

# Firewall rules
iptables-save > "$SNAPSHOT_DIR/iptables_rules.txt"
ufw status verbose > "$SNAPSHOT_DIR/ufw_status.txt" 2>/dev/null

# Docker security config
docker ps --format '{{.Names}}: {{.Image}} (ports: {{.Ports}})' \
  > "$SNAPSHOT_DIR/docker_containers.txt"

# Check exposed ports
ss -tlnp > "$SNAPSHOT_DIR/listening_ports.txt"

# Cron jobs
crontab -l > "$SNAPSHOT_DIR/crontab.txt" 2>/dev/null || true

# Sensitive env var check
printenv > "$SNAPSHOT_DIR/environment_vars.txt" 2>/dev/null

echo "[+] Config snapshot complete: $SNAPSHOT_DIR"

# security-agent/collector/process_monitor.py
"""Real-time process and network monitoring for anomaly detection"""
import psutil
import socket
from datetime import datetime

class ProcessMonitor:
    def __init__(self):
        self.baseline_procs = {}  # Baseline process snapshot
        self.alert_thresholds = {
            'cpu_percent': 90,
            'memory_percent': 85,
            'open_files': 500,
        }

    def collect_snapshot(self):
        """Collect current process state"""
        snapshot = []
        for proc in psutil.process_iter(['pid', 'name', 'cmdline', 
                                          'cpu_percent', 'memory_percent',
                                          'create_time']):
            info = proc.info
            if info['create_time'] == 0:
                continue
            snapshot.append({
                'pid': info['pid'],
                'name': info['name'],
                'cmdline': ' '.join(info['cmdline'] or [])[:200],
                'cpu': info['cpu_percent'] or 0,
                'mem': info['memory_percent'] or 0,
            })
        return snapshot

    def detect_anomalies(self, current, baseline):
        """Compare against baseline, detect anomalous processes"""
        anomalies = []
        
        # Unknown processes
        known_pids = {p['pid'] for p in baseline}
        for proc in current:
            if proc['pid'] not in known_pids:
                anomalies.append({
                    'type': 'unknown_process',
                    'pid': proc['pid'],
                    'name': proc['name'],
                    'cmdline': proc['cmdline'],
                    'severity': 'high',
                })
        
        # Resource anomalies
        for proc in current:
            if proc['cpu'] > self.alert_thresholds['cpu_percent']:
                anomalies.append({
                    'type': 'high_cpu',
                    'pid': proc['pid'],
                    'name': proc['name'],
                    'cpu': proc['cpu'],
                    'severity': 'medium',
                })
        
        return anomalies

# security-agent/auditor/compliance_checker.py
"""Use LLM for VPS configuration compliance auditing"""
import json
from pathlib import Path

class ComplianceChecker:
    def __init__(self, llm_client):
        self.llm = llm_client
        self.checks = self._load_check_rules()

    def _load_check_rules(self):
        """Load security check rules"""
        return {
            'ssh_hardening': {
                'name': 'SSH Security Hardening',
                'priority': 'critical',
                'items': [
                    {'check': 'PermitRootLogin', 'expected': 'no', 'desc': 'Disable root remote login'},
                    {'check': 'PasswordAuthentication', 'expected': 'no', 'desc': 'Disable password auth'},
                    {'check': 'Port', 'expected': 'non_standard', 'desc': 'Use non-standard port'},
                    {'check': 'MaxAuthTries', 'expected': '<4', 'desc': 'Max auth attempts'},
                    {'check': 'X11Forwarding', 'expected': 'no', 'desc': 'Disable X11 forwarding'},
                ]
            },
            'docker_security': {
                'name': 'Docker Security Configuration',
                'priority': 'high',
                'items': [
                    {'check': 'rootless_mode', 'expected': True, 'desc': 'Use Rootless Docker'},
                    {'check': 'no_latest_tag', 'expected': True, 'desc': "Don't use :latest tag"},
                    {'check': 'healthcheck', 'expected': True, 'desc': 'Container health check'},
                    {'check': 'readonly_rootfs', 'expected': True, 'desc': 'Read-only root filesystem'},
                ]
            },
            'network_security': {
                'name': 'Network Security Configuration',
                'priority': 'high',
                'items': [
                    {'check': 'firewall_enabled', 'expected': True, 'desc': 'Firewall enabled'},
                    {'check': 'unnecessary_ports', 'expected': 0, 'desc': 'Close unnecessary ports'},
                    {'check': 'fail2ban', 'expected': True, 'desc': 'Fail2ban running'},
                ]
            },
        }

    async def audit_ssh_config(self, config_content: str) -> dict:
        """Audit SSH configuration file"""
        prompt = f"""You are a professional Linux security auditor. Analyze the following SSH configuration and identify all security issues and compliance items.

Configuration content:

# security-agent/detector/anomaly_detector.py
"""Use LLM for log anomaly detection and threat analysis"""
import re
from datetime import datetime, timedelta

class AnomalyDetector:
    def __init__(self, llm_client):
        self.llm = llm_client
        self.threat_patterns = self._load_threat_signatures()

    def _load_threat_signatures(self):
        """Load known attack pattern signatures"""
        return {
            'brute_force': {
                'pattern': r'Failed password for .* from (\d+\.\d+\.\d+\.\d+)',
                'threshold': 5,
                'window_minutes': 10,
                'severity': 'high',
            },
            'port_scan': {
                'pattern': r'Connection closed by .* port \d+.*preauth',
                'threshold': 20,
                'window_minutes': 5,
                'severity': 'medium',
            },
            'sql_injection': {
                'pattern': r"(?i)(union\s+select|or\s+1=1|drop\s+table|;\s*delete)",
                'threshold': 1,
                'window_minutes': 60,
                'severity': 'critical',
            },
            'directory_traversal': {
                'pattern': r'(\.\./){2,}',
                'threshold': 3,
                'window_minutes': 5,
                'severity': 'high',
            },
        }

    async def analyze_auth_logs(self, log_lines: list[str]) -> dict:
        """Analyze auth logs, detect brute force attacks"""
        ip_counts = {}
        for line in log_lines:
            match = re.search(r'from\s+(\d+\.\d+\.\d+\.\d+)', line)
            if match:
                ip = match.group(1)
                ip_counts[ip] = ip_counts.get(ip, 0) + 1

        threats = []
        for ip, count in ip_counts.items():
            if count >= 5:
                threats.append({
                    'type': 'potential_brute_force',
                    'source_ip': ip,
                    'attempt_count': count,
                    'severity': 'high' if count > 20 else 'medium',
                    'recommendation': f'Recommend blocking IP {ip} or enabling fail2ban',
                })

        # Deep analysis with LLM
        if threats:
            llm_prompt = f"""Analyze the following security events and determine if they are real attacks or false positives:

Event list:
{json.dumps(threats, indent=2)}

Please provide:
1. Attack type classification
2. Risk level assessment
3. Recommended response measures
4. Whether to report to threat intelligence platforms

Return structured analysis results."""
            analysis = await self.llm.generate(llm_prompt, model='mistral')
            threats[0]['llm_analysis'] = analysis

        return {'threats': threats, 'total_events': len(log_lines)}

    async def analyze_web_logs(self, access_log_path: str) -> dict:
        """Analyze web access logs for web attacks"""
        with open(access_log_path) as f:
            lines = f.readlines()[-1000:]  # Last 1000 lines

        suspicious_requests = []
        for line in lines:
            for pattern_name, config in self.threat_patterns.items():
                if re.search(config['pattern'], line):
                    suspicious_requests.append({
                        'pattern': pattern_name,
                        'line': line.strip()[:200],
                    })

        if suspicious_requests:
            llm_prompt = f"""Detected the following suspicious requests. Analyze if they are real attacks:

{suspicious_requests[:20]}

Analyze:
1. Attack type and intent
2. Target system vulnerabilities
3. Urgent response steps
4. Long-term protection recommendations"""
            analysis = await self.llm.generate(llm_prompt, model='mistral')
            return {'suspicious': suspicious_requests, 'analysis': analysis}

        return {'suspicious': [], 'status': 'clean'}

# security-agent/alerts/notifier.py
"""Multi-channel security alert notifications"""
import smtplib
import requests
from email.mime.text import MIMEText

class AlertNotifier:
    def __init__(self, config: dict):
        self.config = config
        self.channels = []

    def register_channel(self, channel_type: str, **kwargs):
        """Register alert channels"""
        if channel_type == 'email':
            self.channels.append(('email', kwargs))
        elif channel_type == 'telegram':
            self.channels.append(('telegram', kwargs))
        elif channel_type == 'webhook':
            self.channels.append(('webhook', kwargs))

    async def send_alert(self, severity: str, title: str, message: str,
                         details: dict = None):
        """Send alerts to all registered channels"""
        alert = {
            'severity': severity,
            'title': title,
            'message': message,
            'timestamp': datetime.now().isoformat(),
            'host': self.config.get('hostname', 'unknown'),
            'details': details,
        }

        for channel, kwargs in self.channels:
            try:
                if channel == 'email':
                    self._send_email(alert, kwargs)
                elif channel == 'telegram':
                    await self._send_telegram(alert, kwargs)
                elif channel == 'webhook':
                    await self._send_webhook(alert, kwargs)
            except Exception as e:
                print(f"[!] Alert failed ({channel}): {e}")

    def _send_email(self, alert: dict, config: dict):
        msg = MIMEText(f"""
VPS Security Alert
{'=' * 40}
Severity: {alert['severity'].upper()}
Title: {alert['title']}
Time: {alert['timestamp']}
Host: {alert['host']}

Details:
{alert['message']}
        """, 'plain', 'utf-8')
        msg['Subject'] = f"[{alert['severity'].upper()}] {alert['title']}"
        msg['From'] = config['from']
        msg['To'] = config['to']

        with smtplib.SMTP(config['smtp_server'], config['smtp_port']) as server:
            server.starttls()
            server.login(config['username'], config['password'])
            server.send_message(msg)

    async def _send_telegram(self, alert: dict, config: dict):
        severity_emoji = {
            'critical': '🔴',
            'high': '🟠',
            'medium': '🟡',
            'low': '🔵',
        }
        emoji = severity_emoji.get(alert['severity'], '⚪')
        text = f"{emoji} *[{alert['severity'].upper()}]*\n"
        text += f"{alert['title']}\n\n"
        text += f"{alert['message']}\n\n"
        text += f"🖥 {alert['host']} | 🕐 {alert['timestamp']}"

        url = f"https://api.telegram.org/bot{config['bot_token']}/sendMessage"
        await requests.post(url, json={
            'chat_id': config['chat_id'],
            'text': text,
            'parse_mode': 'Markdown',
        })

# security-agent/remediation/auto_fixer.py
"""Automated security remediation based on LLM analysis"""
import subprocess
import shlex

class AutoFixer:
    def __init__(self, dry_run: bool = True):
        self.dry_run = dry_run
        self.fix_history = []

    def apply_ssh_fix(self, issue: dict) -> dict:
        """Automatically fix SSH security issues"""
        fix_actions = []

        if issue.get('rule') == 'PermitRootLogin':
            fix_actions.append({
                'action': 'edit_sshd_config',
                'command': 'sed -i "s/^PermitRootLogin.*/PermitRootLogin no/" /etc/ssh/sshd_config',
                'requires_restart': True,
            })

        if issue.get('rule') == 'PasswordAuthentication':
            fix_actions.append({
                'action': 'edit_sshd_config',
                'command': 'sed -i "s/^PasswordAuthentication.*/PasswordAuthentication no/" /etc/ssh/sshd_config',
                'requires_restart': False,
            })

        if issue.get('rule') == 'MaxAuthTries':
            fix_actions.append({
                'action': 'edit_sshd_config',
                'command': 'grep -q "^MaxAuthTries" /etc/ssh/sshd_config && \\\n  sed -i "s/^MaxAuthTries.*/MaxAuthTries 3/" /etc/ssh/sshd_config || \\\n  echo "MaxAuthTries 3" >> /etc/ssh/sshd_config',
                'requires_restart': False,
            })

        for action in fix_actions:
            if self.dry_run:
                print(f"[DRY RUN] {action['command']}")
            else:
                result = subprocess.run(action['command'], shell=True,
                                       capture_output=True, text=True)
                if result.returncode == 0:
                    print(f"[✓] Fix applied: {action['action']}")
                else:
                    print(f"[✗] Fix failed: {result.stderr}")

        return {'actions_applied': len(fix_actions), 'dry_run': self.dry_run}

    def apply_docker_fix(self, issue: dict) -> dict:
        """Automatically fix Docker security issues"""
        fixes = []

        if issue.get('issue_type') == 'root_socket_mount':
            fixes.append({
                'action': 'remove_docker_socket_mount',
                'description': 'Remove Docker socket mount from container',
                'manual_step': 'Edit docker-compose.yml, remove /var/run/docker.sock from volumes',
            })

        if issue.get('issue_type') == 'missing_resource_limits':
            fixes.append({
                'action': 'add_resource_limits',
                'template': '''deploy:
  resources:
    limits:
      cpus: '1.0'
      memory: 512M
    reservations:
      cpus: '0.25'
      memory: 128M''',
                'description': 'Add resource limits to container',
            })

        return {'fixes': fixes}

    def generate_fix_report(self, issues: list, fixes: dict) -> str:
        """Generate remediation report"""
        report = f"""# VPS Security Fix Report
Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
Mode: {'Dry Run Preview' if self.dry_run else 'Live Execution'}

## Summary
- Issues Found: {len(issues)}
- Auto-Fixed: {sum(len(v) for v in fixes.values())}
- Manual Action Needed: {len(issues) - sum(len(v) for v in fixes.values())}

## Detailed Fix Records
"""
        for issue, fix_list in zip(issues, fixes.values()):
            report += f"\n### {issue.get('rule', 'Unknown')}\n"
            for fix in fix_list:
                report += f"- {fix.get('description', fix.get('action'))}\n"

        return report

# Update system
sudo apt update && sudo apt upgrade -y

# Install dependencies
sudo apt install -y python3 python3-pip docker.io \
  fail2ban ufw jq

# Install Ollama (local LLM runtime)
curl -fsSL https://ollama.com/install.sh | sh

# Pull a small model suitable for security auditing
ollama pull mistral:7b

# Clone the security agent project
git clone https://github.com/selfvps/ai-security-agent.git
cd ai-security-agent

# Create configuration file
cat > config.yaml << 'EOF'
llm:
  endpoint: http://localhost:11434
  model: mistral:7b

audit:
  schedule: "0 */6 * * *"  # Audit every 6 hours
  ssh_config: /etc/ssh/sshd_config
  docker_compose_dir: /opt/services

alerts:
  email:
    enabled: true
    smtp_server: smtp.gmail.com
    smtp_port: 587
    from: alert@yourdomain.com
    to: admin@yourdomain.com
  telegram:
    enabled: true
    bot_token: YOUR_BOT_TOKEN
    chat_id: YOUR_CHAT_ID

remediation:
  dry_run: true  # Start with dry-run mode!
  allowed_fixes:
    - ssh_hardening
    - firewall_rules
    - docker_security
EOF

# Add to crontab
crontab -e

# Full security audit daily at 2 AM
0 2 * * * /opt/ai-security-agent/run_audit.sh --report
# Check log anomalies every 6 hours
0 */6 * * * /opt/ai-security-agent/check_anomalies.sh
# Generate weekly report every Sunday
0 9 * * 0 /opt/ai-security-agent/generate_weekly_report.sh

# /etc/fail2ban/jail.local
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 3
backend = auto

[sshd-ai]
enabled = true
port = ssh
filter = sshd
logpath = /var/log/auth.log
maxretry = 3
bantime = 86400

# AI-enhanced custom filter rules
[nginx-ai-bot]
enabled = true
port = http,https
filter = nginx-ai-bot
logpath = /var/log/nginx/access.log
maxretry = 10
bantime = 43200

# Query current security score
$ echo "How is my VPS security?" | ./security-agent/chat

📊 Your VPS Current Security Score: 78/100
Overall Status: 🟡 Needs Attention

Main Risks:
1. SSH allows password authentication (Medium risk)
2. 2 Docker containers lack resource limits (Low risk)
3. UFW firewall rules partially missing (Low risk)

Priority Recommendations:
→ Disable SSH password auth, switch to key-based login
→ Add CPU/memory limits to all containers

# Query recent threat events
$ echo "Any recent security threats?" | ./security-agent/chat

🔒 Security Events (Last 24 Hours):

[🔴 Critical] SSH brute force detected from 185.220.xx.xx
  - Attempts: 342
  - Source IP blacklisted
  - Recommendation: Verify fail2ban rules are active

[🟡 Warning] SQL injection attempt in Nginx access logs
  - Source: 45.142.xx.xx
  - Attack vector: /api/search?q=1' OR 1=1--
  - Automatically blocked

[🔵 Info] System certificate expires in 15 days
  - Certificate: *.yourdomain.com
  - Recommendation: Run certbot renew proactively