Cpu killer

Scriptas skirtas automatiškai killinti userių procesus kurie pastoviai ėda >90% cpu time. Naudoju Free Shells projektui jau eilę metų, čia jau antra scripto iteracija, pirmoji buvo parašyta ant bash.. Scriptas taip pat "decloakina" hidden procesus, t.y pakeistais pavadinimais, pervadina procesą su prefix'u "_too_much_cpu_time". Pervadintas nepasileis automatiškai, jeigu uždėtas ant cron ar kokio kito auto paleisties mechanizmo t.y systemd-user.

#!/usr/bin/env python3
"""
CPU Killer Daemon
- Continuously monitors processes
- Kills if >90% CPU for 10+ minutes
- Renames real binary
- Runs as background service
"""
import os
import sys
import psutil
import shutil
import time
from datetime import datetime
from pathlib import Path
from collections import defaultdict

# ============================= CONFIG =============================
CPU_THRESHOLD = 90.0        # % CPU usage
TIME_THRESHOLD = 600        # seconds (10 minutes)
CHECK_INTERVAL = 30            # check every 30 sec
EXCEPTION_USERS = {'root', 'polkitd', '_chrony', 'postgres', 'nginx', 'devnull', 'zabbix', 'www-data', 'systemd-network', 'systemd-timesync', 'systemd-resolve'}  # <--- ADD YOUR USERS HERE
LOG_FILE = "/root/tools/cpu_killer.log"
DRY_RUN = False             # Set to True for testing (no kill/rename)
PID_FILE = "/root/tools/cpu_killer.pid"
# ==================================================================

def log(msg):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line)
    try:
        with open(LOG_FILE, "a") as f:
            f.write(line + "\n")
    except:
        pass

def daemonize():
    if os.fork(): sys.exit(0)
    os.setsid()
    if os.fork(): sys.exit(0)
    sys.stdout.flush()
    sys.stderr.flush()
    with open('/dev/null', 'r') as dev_null:
        os.dup2(dev_null.fileno(), sys.stdin.fileno())
    with open(LOG_FILE, "a") as f:
        os.dup2(f.fileno(), sys.stdout.fileno())
        os.dup2(f.fileno(), sys.stderr.fileno())

    # Write PID
    with open(PID_FILE, "w") as f:
        f.write(str(os.getpid()))

class ProcessTracker:
    def __init__(self):
        self.tracked = {}  # pid -> {start_time, cpu_samples}

    def update(self, proc):
        pid = proc.pid
        if pid not in self.tracked:
            self.tracked[pid] = {
                'start_time': time.time(),
                'cpu_samples': []
            }
        self.tracked[pid]['cpu_samples'].append(proc.cpu_percent())
        # Keep last 20 samples (~10 min at 30s interval)
        if len(self.tracked[pid]['cpu_samples']) > 20:
            self.tracked[pid]['cpu_samples'].pop(0)

    def is_offender(self, pid):
        data = self.tracked.get(pid)
        if not data or len(data['cpu_samples']) < 10:
            return False
        # Average CPU over last ~5-10 min
        avg_cpu = sum(data['cpu_samples']) / len(data['cpu_samples'])
        runtime = time.time() - data['start_time']
        return avg_cpu >= CPU_THRESHOLD and runtime >= TIME_THRESHOLD

    def cleanup(self, pid):
        self.tracked.pop(pid, None)

tracker = ProcessTracker()

def get_real_binary(proc):
    try:
        exe = f"/proc/{proc.pid}/exe"
        if not os.path.exists(exe):
            return None
        path = os.readlink(exe)
        if "(deleted)" in path:
            path = path.replace(" (deleted)", "").strip()
        return path if os.path.exists(path) else None
    except:
        return None

def rename_binary(bin_path):
    if not bin_path or not os.path.exists(bin_path):
        return
    path = Path(bin_path)
    new_name = path.parent / f"{path.stem}_too_much_cpu_time{path.suffix}"
    i = 1
    while new_name.exists():
        new_name = path.parent / f"{path.stem}_too_much_cpu_time.{i}{path.suffix}"
        i += 1
    log(f"  -> Renaming: {path.name} → {new_name.name}")
    if not DRY_RUN:
        try:
            shutil.move(str(path), str(new_name))
        except Exception as e:
            log(f"  -> Rename failed: {e}")

def kill_process(proc):
    try:
        cmd = " ".join(proc.cmdline()[:3]) + ("..." if len(proc.cmdline()) > 3 else "")
        log(f"KILLING PID {proc.pid} ({cmd}) | User: {proc.username()} | CPU: ~{proc.cpu_percent():.1f}%")
        if not DRY_RUN:
            proc.terminate()
            try:
                gone, alive = psutil.wait_procs([proc], timeout=5)
                if alive:
                    log(f"  -> Force kill")
                    for p in alive:
                        p.kill()
            except:
                pass
        return True
    except Exception as e:
        log(f"  -> Kill failed: {e}")
        return False

def main_loop():
    log("CPU Killer started (daemon mode)")
    seen_pids = set()

    while True:
        try:
            current_pids = {p.pid for p in psutil.process_iter()}
            # Cleanup dead
            for pid in list(tracker.tracked.keys()):
                if pid not in current_pids:
                    tracker.cleanup(pid)

            for proc in psutil.process_iter(['pid', 'name', 'username', 'cpu_percent', 'create_time', 'cmdline']):
                try:
                    username = proc.info['username']
                    if not username or username in EXCEPTION_USERS or username.startswith('_'):
                        continue
                    if proc.pid < 100:
                        continue

                    # Measure CPU (short interval)
                    cpu = proc.cpu_percent(interval=0.1)
                    if cpu <= 0:
                        continue

                    tracker.update(proc)

                    if tracker.is_offender(proc.pid):
                        bin_path = get_real_binary(proc)
                        if bin_path:
                            log(f"OFFENDER DETECTED: PID {proc.pid} | User: {username} | Binary: {bin_path}")
                        else:
                            log(f"OFFENDER DETECTED: PID {proc.pid} | User: {username} | Binary: [IN-MEMORY]")

                        if kill_process(proc):
                            if bin_path:
                                rename_binary(bin_path)
                            tracker.cleanup(proc.pid)

                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue

            time.sleep(CHECK_INTERVAL)

        except KeyboardInterrupt:
            log("CPU Killer stopped by user")
            break
        except Exception as e:
            log(f"Unexpected error: {e}")
            time.sleep(CHECK_INTERVAL)

    if os.path.exists(PID_FILE):
        os.remove(PID_FILE)

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "stop":
        if os.path.exists(PID_FILE):
            with open(PID_FILE) as f:
                pid = int(f.read().strip())
            try:
                os.kill(pid, 15)
                log(f"Stopped CPU Killer (PID {pid})")
            except:
                log("Failed to stop (not running?)")
            os.remove(PID_FILE)
        else:
            log("Not running")
        sys.exit(0)

    if os.path.exists(PID_FILE):
        log("Already running!")
        sys.exit(1)

    daemonize()
    main_loop()

Systemd unit'as:

[Unit]
Description=CPU Killer Daemon
After=network.target

[Service]
Type=simple
ExecStart=/root/tools/cpu_killer.py
ExecStop=/root/tools/cpu_killer.py stop
Restart=always
RestartSec=10
StandardOutput=null
StandardError=journal
User=root
PIDFile=/root/tools/cpu_killer.pid

[Install]
WantedBy=multi-user.target