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
import stat
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
SAFE_RENAME_PREFIXES = ('/home/',) # Only rename files living under these prefixes
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 _is_under_safe_prefix(p: str) -> bool:
    try:
        rp = Path(p).resolve().as_posix()
    except Exception:
        return False
    return any(rp.startswith(prefix) for prefix in SAFE_RENAME_PREFIXES)

def _abs_from_proc_cwd(proc: psutil.Process, path: str) -> str:
    """Make path absolute using the process' cwd if it's relative."""
    if os.path.isabs(path):
        return path
    try:
        cwd = os.readlink(f"/proc/{proc.pid}/cwd")
        return os.path.normpath(os.path.join(cwd, path))
    except Exception:
        return path

def get_safe_rename_target(proc: psutil.Process) -> str | None:
    """
    Return a file path that is safe to rename:
    - Prefer the launched script (for Python or shebang scripts)
    - Must be under /home/
    - Never return the interpreter/binary in /usr, /bin, etc.
    """
    try:
        cmd = proc.cmdline()
        if not cmd:
            return None

        candidates = []

        # If invoked as: python /home/user/app.py ...
        exe_base = os.path.basename(cmd[0])
        if 'python' in exe_base.lower() and len(cmd) >= 2:
            candidates.append(cmd[1])

        # If invoked directly: /home/user/app.py (shebang)
        candidates.append(cmd[0])

        # Consider only first existing file under /home/
        for c in candidates:
            c_abs = _abs_from_proc_cwd(proc, c)
            if not c_abs:
                continue
            try:
                st = os.stat(c_abs)
                if not stat.S_ISREG(st.st_mode):
                    continue
            except Exception:
                continue
            if _is_under_safe_prefix(c_abs):
                return c_abs

        # As a last resort, fall back to real binary only if it's under /home/
        real_bin = get_real_binary(proc)
        if real_bin and _is_under_safe_prefix(real_bin):
            return real_bin

        return None
    except Exception:
        return None



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
    if not _is_under_safe_prefix(bin_path):
        log(f"  -> Skip rename (outside safe prefixes): {bin_path}")
        return

    path = Path(bin_path)
    # keep suffix for .py, .sh, etc.
    suffix = path.suffix
    stem = path.stem
    new_name = path.parent / f"{stem}_too_much_cpu_time{suffix}"

    i = 1
    while new_name.exists():
        new_name = path.parent / f"{stem}_too_much_cpu_time.{i}{suffix}"
        i += 1

    log(f"  -> Renaming: {path} → {new_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):
                        target = get_safe_rename_target(proc)
                        bin_path = get_real_binary(proc)
                        log_bin = bin_path if bin_path else "[IN-MEMORY]"
                        log(f"OFFENDER DETECTED: PID {proc.pid} | User: {username} | Exe: {log_bin} | Target: {target or '[none]'}")

                        if kill_process(proc):
                            if target:
                                rename_binary(target)
                            else:
                                log("  -> No safe /home/ target to rename; skipping rename.")
                            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