<reponame>reviforks/CivilServant
#!/usr/bin/env python3

import argparse
import email.message
import json
import os
import socket
import subprocess
from pathlib import Path
from smtplib import SMTP

from psutil import Process, process_iter, wait_procs
from psutil import ZombieProcess, NoSuchProcess, AccessDenied

BASE_DIR = Path(__file__).resolve().parents[1]
EMAIL_CONFIG_PATH = Path(BASE_DIR, "config", "email_worker_restart.json")
ENV = os.environ["CS_ENV"]
HOSTNAME = socket.gethostname()
SMTP_HOST = "localhost"
SMTP_PORT = 25


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--cmdline", required=True,
                        help="The command to be executed.")
    parser.add_argument("-e", "--exclude-forks", action="store_true",
                        default=False,
                        help="Exclude forked child processes in the count.")
    parser.add_argument("-m", "--min-count", type=int, required=True,
                        help="The minimum count of processes required.")
    parser.add_argument("-n", "--notify", action="store_true", default=False,
                        help="Send email notifications as needed.")
    parser.add_argument("-s", "--screen", required=True,
                        help="The name of the target screen.")

    mutex = parser.add_mutually_exclusive_group()
    mutex.add_argument("-d", "--daemon", action="store_true", default=False,
                       help="Run the monitor continuously.")
    mutex.add_argument("-D", "--dry-run", action="store_true", default=False,
                       help="Just output commands rather than running them.")

    return parser.parse_args()


class WorkerMonitor:
    """Monitor and respawn worker processes when necessary."""

    def __init__(self, cmdline, min_count, screen, exclude_forks=False,
                 notify=False, daemon=False, dry_run=False):
        self._current_process = Process()
        self.cmdline = cmdline
        self.daemon = daemon
        self.dry_run = dry_run
        self.exclude_forks = exclude_forks
        self.min_count = min_count
        self.notify = notify
        self.screen = screen

    def __repr__(self):
        process_name = cmdline.split(" ")[0] if self.cmdline else ""
        return (f"<WorkerMonitor process_name={process_name} "
                f"min_count={self.min_count}>")

    def _notify(self, count):
        """Send a notification email that workers needed to be restarted."""
        from_, recipients = self._read_email_config()
        if recipients:
            subject = f"{count} worker(s) were just restarted on {HOSTNAME}."
            body = f"{subject}\n\nWorker process: {self.cmdline}"
            for recipient in recipients:
                with SMTP(SMTP_HOST, SMTP_PORT) as smtp:
                    message = email.message.EmailMessage()
                    message["From"] = from_
                    message["To"] = recipient
                    message["Subject"] = subject
                    message.set_content(body)
                    smtp.send_message(message)

    def _read_email_config(self):
        """Read the email configuration file."""
        with open(str(EMAIL_CONFIG_PATH)) as f:
            email_config = json.loads(f.read())
            return email_config["fromaddr"], email_config["toaddrs"]

    def _screen_init(self):
        """Initialize the screen session and run the first process."""
        num_running = len(list(self.fetch_processes()))
        num_needed = max(self.min_count - num_running, 0)
        if num_needed == self.min_count:
            cmd = f"CS_ENV={ENV} screen -d -m -S {self.screen} {self.cmdline}"
            if self.dry_run:
                cmd = f"echo {cmd}"
            subprocess.Popen(cmd, shell=True).wait()

    def _screen_add_processes(self, count):
        """Add new processes to an existing screen session."""
        cmd = f"screen -S {self.screen} -X screen {self.cmdline}"
        if self.dry_run:
            cmd = f"echo {cmd}"
        for _ in range(count):
            subprocess.Popen(cmd, shell=True).wait()

    def _screen_update_env(self):
        """Ensure CS_ENV is set on all new screen windows by default."""
        cmd = f"screen -S {self.screen} -X setenv CS_ENV {ENV}"
        if self.dry_run:
            cmd = f"echo {cmd}"
        subprocess.Popen(cmd, shell=True).wait()

    def _validate_process(self, process):
        """Ensure a process meets the criteria to be included."""
        if process == self._current_process:
            return False
        if " ".join(process.cmdline()) != self.cmdline:
            return False
        if self.exclude_forks and process.parent().cmdline() == self.cmdline:
            return False
        return True

    def fetch_processes(self):
        """Fetch the running target processes."""
        for process in process_iter():
            try:
                if self._validate_process(process):
                    yield process
            except (ZombieProcess, NoSuchProcess, AccessDenied):
                pass

    def run(self):
        """Run the worker monitor."""
        if self.daemon:
            self.run_daemon()
        else:
            self.run_once()

    def run_daemon(self):
        """Run the worker monitor continuously."""
        self._screen_init()
        self._screen_update_env()
        while True:
            # Retrieves processes twice to ensure the list contains the newly
            # added processes since there is no way to return their individual
            # pids due to launching them wrapped in a screen
            processes = list(self.fetch_processes())
            num_needed = max(self.min_count - len(processes), 0)
            self._screen_add_processes(num_needed)
            if self.notify and num_needed > 0:
                self._notify(num_needed)
            processes = list(self.fetch_processes())
            wait_procs(processes, timeout=3)

    def run_once(self):
        """Perform a single run of the worker monitor."""
        self._screen_init()
        self._screen_update_env()
        num_running = len(list(self.fetch_processes()))
        num_needed = max(self.min_count - num_running, 0)
        self._screen_add_processes(num_needed)
        if self.notify and num_needed > 0:
            self._notify(num_needed)


if __name__ == "__main__":
    try:
        args = parse_args()
        manager = WorkerMonitor(**vars(args))
        manager.run()
    except KeyboardInterrupt:
        pass


#!/usr/bin/env python3

import argparse
import email.message
import json
import os
import socket
import subprocess
from pathlib import Path
from smtplib import SMTP

from psutil import Process, process_iter, wait_procs
from psutil import ZombieProcess, NoSuchProcess, AccessDenied

BASE_DIR = Path(__file__).resolve().parents[1]
EMAIL_CONFIG_PATH = Path(BASE_DIR, "config", "email_worker_restart.json")
ENV = os.environ["CS_ENV"]
HOSTNAME = socket.gethostname()
SMTP_HOST = "localhost"
SMTP_PORT = 25


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--cmdline", required=True,
                        help="The command to be executed.")
    parser.add_argument("-e", "--exclude-forks", action="store_true",
                        default=False,
                        help="Exclude forked child processes in the count.")
    parser.add_argument("-m", "--min-count", type=int, required=True,
                        help="The minimum count of processes required.")
    parser.add_argument("-n", "--notify", action="store_true", default=False,
                        help="Send email notifications as needed.")
    parser.add_argument("-s", "--screen", required=True,
                        help="The name of the target screen.")

    mutex = parser.add_mutually_exclusive_group()
    mutex.add_argument("-d", "--daemon", action="store_true", default=False,
                       help="Run the monitor continuously.")
    mutex.add_argument("-D", "--dry-run", action="store_true", default=False,
                       help="Just output commands rather than running them.")

    return parser.parse_args()


class WorkerMonitor:
    """Monitor and respawn worker processes when necessary."""

    def __init__(self, cmdline, min_count, screen, exclude_forks=False,
                 notify=False, daemon=False, dry_run=False):
        self._current_process = Process()
        self.cmdline = cmdline
        self.daemon = daemon
        self.dry_run = dry_run
        self.exclude_forks = exclude_forks
        self.min_count = min_count
        self.notify = notify
        self.screen = screen

    def __repr__(self):
        process_name = cmdline.split(" ")[0] if self.cmdline else ""
        return (f"<WorkerMonitor process_name={process_name} "
                f"min_count={self.min_count}>")

    def _notify(self, count):
        """Send a notification email that workers needed to be restarted."""
        from_, recipients = self._read_email_config()
        if recipients:
            subject = f"{count} worker(s) were just restarted on {HOSTNAME}."
            body = f"{subject}\n\nWorker process: {self.cmdline}"
            for recipient in recipients:
                with SMTP(SMTP_HOST, SMTP_PORT) as smtp:
                    message = email.message.EmailMessage()
                    message["From"] = from_
                    message["To"] = recipient
                    message["Subject"] = subject
                    message.set_content(body)
                    smtp.send_message(message)

    def _read_email_config(self):
        """Read the email configuration file."""
        with open(str(EMAIL_CONFIG_PATH)) as f:
            email_config = json.loads(f.read())
            return email_config["fromaddr"], email_config["toaddrs"]

    def _screen_init(self):
        """Initialize the screen session and run the first process."""
        num_running = len(list(self.fetch_processes()))
        num_needed = max(self.min_count - num_running, 0)
        if num_needed == self.min_count:
            cmd = f"CS_ENV={ENV} screen -d -m -S {self.screen} {self.cmdline}"
            if self.dry_run:
                cmd = f"echo {cmd}"
            subprocess.Popen(cmd, shell=True).wait()

    def _screen_add_processes(self, count):
        """Add new processes to an existing screen session."""
        cmd = f"screen -S {self.screen} -X screen {self.cmdline}"
        if self.dry_run:
            cmd = f"echo {cmd}"
        for _ in range(count):
            subprocess.Popen(cmd, shell=True).wait()

    def _screen_update_env(self):
        """Ensure CS_ENV is set on all new screen windows by default."""
        cmd = f"screen -S {self.screen} -X setenv CS_ENV {ENV}"
        if self.dry_run:
            cmd = f"echo {cmd}"
        subprocess.Popen(cmd, shell=True).wait()

    def _validate_process(self, process):
        """Ensure a process meets the criteria to be included."""
        if process == self._current_process:
            return False
        if " ".join(process.cmdline()) != self.cmdline:
            return False
        if self.exclude_forks and process.parent().cmdline() == self.cmdline:
            return False
        return True

    def fetch_processes(self):
        """Fetch the running target processes."""
        for process in process_iter():
            try:
                if self._validate_process(process):
                    yield process
            except (ZombieProcess, NoSuchProcess, AccessDenied):
                pass

    def run(self):
        """Run the worker monitor."""
        if self.daemon:
            self.run_daemon()
        else:
            self.run_once()

    def run_daemon(self):
        """Run the worker monitor continuously."""
        self._screen_init()
        self._screen_update_env()
        while True:
            # Retrieves processes twice to ensure the list contains the newly
            # added processes since there is no way to return their individual
            # pids due to launching them wrapped in a screen
            processes = list(self.fetch_processes())
            num_needed = max(self.min_count - len(processes), 0)
            self._screen_add_processes(num_needed)
            if self.notify and num_needed > 0:
                self._notify(num_needed)
            processes = list(self.fetch_processes())
            wait_procs(processes, timeout=3)

    def run_once(self):
        """Perform a single run of the worker monitor."""
        self._screen_init()
        self._screen_update_env()
        num_running = len(list(self.fetch_processes()))
        num_needed = max(self.min_count - num_running, 0)
        self._screen_add_processes(num_needed)
        if self.notify and num_needed > 0:
            self._notify(num_needed)


if __name__ == "__main__":
    try:
        args = parse_args()
        manager = WorkerMonitor(**vars(args))
        manager.run()
    except KeyboardInterrupt:
        pass
#!/usr/bin/env python3 Parse command line arguments. Monitor and respawn worker processes when necessary. Send a notification email that workers needed to be restarted. Read the email configuration file. Initialize the screen session and run the first process. Add new processes to an existing screen session. Ensure CS_ENV is set on all new screen windows by default. Ensure a process meets the criteria to be included. Fetch the running target processes. Run the worker monitor. Run the worker monitor continuously. # Retrieves processes twice to ensure the list contains the newly # added processes since there is no way to return their individual # pids due to launching them wrapped in a screen Perform a single run of the worker monitor.