# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Scheduling implementation -- persistence, threads, scheduling.

This module exports:

TaskScheduler -- schedule manager

Helper functions and classes:

timedelta_to_seconds() -- converts a timedelta to seconds
TaskThread    -- custom thread class for executing scheduler tasks
TaskState     -- makes the timestamp (when was a task last run) persistent

"""

# Python imports
import os
import sys
import copy
import time
import signal
import logging
import datetime
import threading

is_win = os.name == "nt"

# Sigma imports
from sigma.util import avalog
from sigma.util import config
from sigma.util import dance
from sigma.util import xmltime
from sigma.sched import timezones, tasks
from sigma.sched import process

# Global objects
logger = logging.getLogger("sched.scheduling")


def timedelta_to_seconds(d):
    """Converts a timedelta into a seconds duration."""
    return (d.days * 3600 * 24) + d.seconds + (d.microseconds * 1e-6)


class TaskState(object):

    """Persistent record of task state.

    This is used to keep track of the last time a task was run across
    scheduler invocations.  (This is currently the only state we want
    to persist.)

    There is one public read-only attribute:

    marktime -- datetime instance telling when mark() was last called; or
                None if never.  This is in UTC without an explicit timezone.

    Public methods:

    mark()   -- marks and saves the current time
    delete() -- deletes the state file

    """

    def __init__(self, id, sched=None):
        """Constructor.  Takes a task id and an optional TaskScheduler."""
        self.id = id
        self.sched = sched
        self.marktime = None
        self.filename = None
        if self.sched is not None and sched.taskdir is not None:
            self.filename = os.path.join(sched.taskdir, "task." + id)
            self.load()

    # API methods

    def mark(self):
        """Marks and saves the current time."""
        self.marktime = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
        self.save()

    def delete(self):
        """Deletes the state file."""
        if self.filename is not None:
            if os.path.exists(self.filename):
                try:
                    os.remove(self.filename)
                except os.error as err:
                    logger.error("can't delete state file %s: %s",
                                 self.filename, err)
                else:
                    logger.debug("deleted state file %s", self.filename)

    # Internal methods

    def load(self):
        """Loads the state from the file."""
        if self.filename is not None:
            try:
                f = open(self.filename, "r")
                try:
                    s = f.read()
                finally:
                    f.close()
            except IOError as err:
                # This is not an error!
                logger.debug("can't read state file %s", self.filename)
            else:
                s = s.strip()
                try:
                    self.marktime = xmltime.parse(s)
                except ValueError as err:
                    logger.error("can't parse mark time %r: %s", s, err)
                else:
                    logger.debug("loaded mark time %r from %s",
                                 self.marktime, self.filename)

    def save(self):
        """Writes the state to the file."""
        if self.filename is not None:
            try:
                f = open(self.filename, "w")
                try:
                    f.write("%s\n" % xmltime.formatdatetime(self.marktime))
                finally:
                    f.close()
            except IOError as err:
                logger.error("can't write state file %s", self.filename)
            else:
                logger.debug("saved %s to %s", self.marktime, self.filename)


class TaskThread(threading.Thread):

    """A Thread subclass used to manage a single task.

    For repetitive tasks, this mostly sits in a loop sleeping until
    the next time the task needs to be run, and then running it and
    waiting for its completion.

    Complications are caused by the requirement to be able to delete
    the task or shut down the thread, and the need to persist the most
    recent time the task was run.

    API methods:

    constructor
    shutdown()
    delete()
    doitnow()

    Also, these methods inherited from threading.Thread are useful:

    start()
    join()
    getName()
    isAlive()

    """

    def __init__(self, task, sched=None, options=None):
        """Constructor.  Takes a Task and an optional TaskScheduler."""
        assert task.command is not None, \
               "task.command not set (init): %r" % task
        assert task.freq is None or task.freq > 0, \
               "task.freq negative (init): %r" % task
        threading.Thread.__init__(self)
        self.name = task.id
        self.sched = sched
        if options is None:
            options = config.Options()
        self.options = options
        self.mutex = threading.Lock()
        # The following instance variables are protected by the mutex:
        self.task = task
        self.freq = task.freq
        self.state = TaskState(task.id, sched)
        self.interrupted = False
        self.runnow = False

    # API methods

    def shutdown(self):
        """Stops the thread without deleting the task state.

        This is used when shutting down the scheduling server.

        """
        self.lock()
        try:
            self.interrupted = True
            self.task = None
        finally:
            self.unlock()

    def delete(self):
        """Stops the thead and deletes the persistent state.

        This is used when deleting the task.

        """
        self.lock()
        try:
            self.interrupted = True
            self.task = None
            self.state.delete()
        finally:
            self.unlock()

    def doitnow(self):
        """Interrupts the sleep and causes the task to run right now."""
        self.lock()
        try:
            self.runnow = True
            self.interrupted = True
        finally:
            self.unlock()

    # Internal methods

    def lock(self):
        """Acquires the mutex."""
        self.mutex.acquire()

    def unlock(self):
        """Releases the mutex."""
        self.mutex.release()

    def run(self):
        """Run -- this code runs in the tread.

        It appears as if this holds on to the mutex at all times; but
        in fact the mutex is released (and reacquired before
        returning) by sleep() and execute() while they are sleeping or
        waiting.

        """
        delete_from_sched = False
        self.lock()
        try:
            while self.task is not None:
                task_id = self.task.id
                secs = self.sleeptime()
                if secs > 0:
                    if self.expired(secs):
                        logger.info("Task %s: expired", task_id)
                        delete_from_sched = True
                        break
                    logger.info("Task %s: sleep %s", task_id, int(secs))
                    if self.sleep(secs):
                        continue
                if self.expired():
                    logger.info("Task %s: expired", task_id)
                    delete_from_sched = True
                    break
                secs = self.blackoutcheck()
                if secs:
                    logger.info("Task %s: in blackout period; sleep %s",
                                task_id, secs)
                    self.sleep(secs)
                    continue
                # Reduce frequency in anticipation of failure;
                # this is reset to self.task.freq upon success
                if self.freq is not None:
                    limit = self.options.backoff_limit
                    if  self.freq < limit:
                        factor = self.options.backoff_factor
                        self.freq = min(self.freq*factor, limit)
                self.execute(self.task)
                if self.freq is None:
                    self.state.delete()
                    delete_from_sched = True
                    break
        finally:
            self.unlock()
        # Some work must be done after releasing the lock
        if delete_from_sched:
            if  self.sched is not None and self.task is not None:
                logger.info("Task %s: deleting", self.task.id)
                self.sched.add_list([tasks.newtask(self.task.id)])
                self.sched.save()

    def expired(self, secs=0):
        """Returns whether the current task is expired beyond its end date.

        The optional 'secs' argument gives a number of seconds to be
        added to 'now' to find the cut-off point (it indicates how far
        in the future the task's execution is scheduled).

        Note that the end date is inclusive.  Since toDatetime()
        returns a datetime object with the time set to zero, we must
        add a day to obtain the cut-off time.

        """
        if not self.task:
            return True
        if not self.task.end:
            return False
        now = datetime.datetime.now(timezones.Local)
        t = self.task.end.toDatetime()
        dt = datetime.timedelta(days=1, seconds=-secs)
        return t+dt <= now

    def blackoutcheck(self):
        """Checks whether the current task is not allowed to run right now.

        Returns a non-zero number of seconds to sleep before trying again.

        Returns zero if we're not in a blackout period.

        """
        if not self.task:
            return 0
        if not self.task.schedule:
            return 0
        now = time.localtime()
        wday, hour = now.tm_wday, now.tm_hour # 0 == Monday
        for sched in self.task.schedule:
            if sched.day != wday:
                continue
            if sched.start is not None and hour < sched.start:
                continue
            if sched.end is not None and hour >= sched.end:
                continue
            return 0
        # No schedule period allows us to run; try again when the hour
        # is over (intentionally overshooting by 1-60 seconds)
        secs = 3601 - 60*now.tm_min
        return secs

    def sleeptime(self):
        """Calculates sleep time."""
        if self.freq is None:
            return self.sleeptime_oneshot()
        else:
            return self.sleeptime_repeated()

    def sleeptime_oneshot(self):
        """Calculates sleep time for one-shot task."""
        assert self.freq is None, "self.freq is set (oneshot)"
        if self.task.when is None:
            # Tasks without specific start time run now
            return 0
        else:
            now = datetime.datetime.now(timezones.utc)
            secs = timedelta_to_seconds(self.task.when - now)
            # Tasks scheduled in the past run now
            return max(0, secs)

    def sleeptime_repeated(self):
        """Calculates sleep time for repeated task."""
        assert self.freq is not None, "self.freq not set (repeated)"
        # Use the freerunning algorithm when in exponential back-off mode
        if self.task.when is None or self.task.freq != self.freq:
            return self.sleeptime_freerunning()
        else:
            return self.sleeptime_synchronized()

    def sleeptime_synchronized(self):
        """Calculates sleep time for synchronized repeated task."""
        assert self.task.freq is not None, "task.freq not set (synchronized)"
        assert self.task.when is not None, "task.when not set (synchronized)"
        now = datetime.datetime.now(timezones.utc)
        secs = timedelta_to_seconds(self.task.when - now)
        if self.state.marktime is None:
            # First time; if start time is in the past, run it now
            return max(0, secs)
        else:
            # Repeat occurrence (expect secs < 0, but don't assert,
            # because the clock might have jumped)
            if secs < 0:
                secs = secs % self.task.freq # Always >= 0
            # The state uses unmarked UTC time
            now2 = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
            secs2 = timedelta_to_seconds(self.state.marktime - now2)
            if secs2 < secs - 1.5*self.task.freq:
                secs = 0 # Way overdue
            elif secs2 > secs - 0.5*self.task.freq:
                secs += self.task.freq # Did it rather recently
            return secs

    def sleeptime_freerunning(self):
        """Calculates sleep time for free-running repeated task."""
        if self.freq == self.task.freq:
            assert self.freq is not None, "self.freq not set (freerunning)"
            assert self.task.when is None, "task.when is set (freerunning)"
        if self.state.marktime is None:
            # First time
            return 0
        else:
            # Repeat occurrence
            now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
            secs = timedelta_to_seconds(self.state.marktime - now)
            secs += self.freq
            return max(0, secs)

    def sleep(self, secs):
        """Sleeps the given number of seconds, or until interrupted.

        CALL THIS WITH THE MUTEX LOCKED!

        Return True if interrupted without 'runnow' set, False if
        sleep time completed or interrupted with 'runnow' set.  (IOW,
        if this returns True, don't execute the command; if it returns
        False, do.)

        Implementation note: don't use the timed wait() methods on
        Condition or Event variables; those wake up 20 times per
        second, making them a bit of a resource hog.

        """
        task_id = "?"
        if self.task is not None:
            task_id = self.task.id
        limit = 1.5*secs
        self.unlock()
        try:
            then = time.time() + secs
            while secs > 0 and not self.interrupted:
                time.sleep(min(secs, 1))
                secs = then - time.time()
                if secs > limit:
                    # System clock moved way back; err on the side of
                    # sleeping too short
                    logger.info("Task %s: reset (clock moved back)", task_id)
                    break
        finally:
            self.lock()
        if not self.interrupted:
            return False
        self.interrupted = False
        if self.runnow:
            self.runnow = False
            return False
        return True

    def execute(self, task):
        """Wrapper around execute_raw() that handles the directive lock."""
        if not task.id.startswith("engine."):
            self.execute_raw(task)
            return
        # Only one engine.* task may run at a time
        logger.info("Task %s: special handling for directive tasks", task.id)
        if not self.lock_directive(task):
            return
        try:
            self.execute_raw(task)
        finally:
            self.unlock_directive(task)

    def lock_directive(self, task):
        """Attempt to acquire the directive semaphore.

        Return True if successful, False if giving up (because we've
        been interrupted).

        """
        sched = self.sched
        mlock = sched.directiveLock
        logtimer = -1
        while True:
            mlock.acquire()
            try:
                if sched.directiveTask is None:
                    sched.directiveTask = task
                    logger.info("Task %s: acquired directive semaphore",
                                task.id)
                    return True
            finally:
                mlock.release()
            if logtimer > 0:
                logtimer -= 1
            else:
                msg = (logtimer < 0) and "waiting" or "still waiting"
                logger.info("Task %s: %s for task %s",
                            task.id, msg, sched.directiveTask.id)
                logtimer = 60
            if self.sleep(1):
                logger.info("Task %s: interrupted, giving up", task.id)
                return False

    def unlock_directive(self, task):
        """Release the directive semaphore."""
        assert self.sched.directiveTask is task, "bad directiveTask (unlock)"
        self.sched.directiveTask = None
        logger.info("Task %s: released directive semaphore", task.id)

    def execute_raw(self, task):
        """Executes the task's command.

        The task is passed in as an argument because self.task may be
        set to None when another thread calls shutdown() on us.

        CALL THIS WITH THE MUTEX LOCKED!

        """
        logger.debug("Task %s: running command %r", task.id, task.command)
        # We intentionally split the command on spaces.  This
        # restricts us in the kind of arguments we can pass.  But on
        # Windows, this is how the command line is parsed anyway by
        # the system, so we might as well live with it.
        args = task.command.split()
        # Expand ~data
        args = [self.expand_tilde(arg) for arg in args]
        program = os.path.normcase(args[0])
        args[0] = os.path.join(self.options.agentbin, program)
        if program.endswith(".py"):
            args.insert(0, sys.executable)
        else:
            # For the poliwall executable, set the environment
            if task.id == "poliwall":
                os.environ["AVDATA"] = self.options.agentdata
                os.environ["AVLIB"] = self.options.agentlib
                # Also delete a pwkill file that might exist
                pwkill = self.options.datapath("pwkill")
                if os.path.exists(pwkill):
                    try:
                        os.remove(pwkill)
                    except os.error as err:
                        logger.warning("Can't delete %s: %s", pwkill, err)
        self.state.mark()
        self.unlock()
        try:
            self.howtokill = None
            pid = 0
            self.proc = process.createProcess(args)
            pid = self.proc.getpid()
            logger.info("Task %s: pid %d", task.id, pid)
            while True:
                try:
                    status = self.proc.getstatus()
                    if status >= 0: # -1 means still alive
                        break
                except process.Error as err:
                    logger.error("Task %s: pid %s: getstatus error: %s",
                                 task.id, pid, err)
                    status = 0
                    break
                if self.interrupted:
                    self.kill_it(task)
                time.sleep(1)
        finally:
            self.lock()
        self.decode_status(task.id, pid, status)
        if self.interrupted:
            self.interrupted = False

    def expand_tilde(self, arg):
        "Replace leading ~data/ with $AVDATA."
        if arg.startswith("~data/"):
            arg = self.options.datapath(arg[6:])
        return arg

    def kill_it(self, task):
        "Kill a subprocess using appropriate signals."
        proc = self.proc
        pid = proc.getpid()
        how = self.get_howtokill(task)
        if how:
            func = how[0]
            args = how[1:]
            callmsg = "%s(%s)" % (func.__name__, ", ".join(map(repr, args)))
            try:
                func(*args)
            except Exception as err:
                logger.error("Task %s: pid %d: %s failed: %s: %s",
                             task.id, pid, callmsg,
                             err.__class__.__name__, err)
            else:
                logger.info("Task %s: pid %d: %s called",
                            task.id, pid, callmsg)

    def get_howtokill(self, task):
        "Return the next value in the sequence of kill arguments."
        if self.howtokill is None:
            if is_win:
                # Default approach for Windows is a hard kill
                self.howtokill = [(self.proc.kill, 98)]
                # Handle idiosyncrasies
                if task.id == "poliwall":
                    # The poliwall can't intercepts signals on
                    # Windows, so we ask it to exit politely using
                    # "poliwall.py -k".  The extra Nones give it four
                    # seconds: up to two seconds for the poliwall to
                    # notice, and then two more seconds to dump its
                    # files.
                    import poliwall
                    self.howtokill = [(poliwall.main, ["-k"]),
                                      None,
                                      None,
                                      None] + self.howtokill
            else:
                # Unix does everything with signals.  First send
                # SIGTERM, four seconds later send SIGKILL.  (Four
                # seconds are needed for the poliwall, see above.)
                self.howtokill = [(self.proc.kill, signal.SIGTERM),
                                  None, None, None,
                                  (self.proc.kill, signal.SIGKILL)]
        how = self.howtokill.pop(0)
        if not self.howtokill:
            self.howtokill.append(how)
        return how

    def decode_status(self, id, pid, status):
        """Decode status returned by os.waitpid()."""
        if status == 0:
            logger.info("Task %s: pid %d: exited normally", id, pid)
            # Reset self.freq to original task.freq
            if self.task is not None:
                self.freq = self.task.freq
        elif not hasattr(os, "WIFSIGNALED"):
            # Windows
            if self.interrupted and status == 98:
                level = logging.INFO
            else:
                level = logging.ERROR
            logger.log(level, "Task %s: pid %d: exited with status %d",
                       id, pid, status)
        elif os.WIFSIGNALED(status):
            sig = os.WTERMSIG(status)
            if self.interrupted and sig in [signal.SIGTERM, signal.SIGKILL]:
                level = logging.INFO
            else:
                level = logging.ERROR
            logger.log(level, "Task %s: pid %d: exited, killed by signal %d",
                       id, pid, sig)
        elif os.WIFEXITED(status):
            sts = os.WEXITSTATUS(status)
            logger.error("Task %s: pid %d: exit status %d", id, pid, sts)


class TaskScheduler(object):

    """Schedule manager."""

    def __init__(self, schedule_file, taskdir=None, options=None):
        """Constructor.

        Takes a schedule file name and an optional task directory
        name.  The latter is where the task threads will persist their
        state; if no directory name is given, their state is not
        persistent.

        """
        self.schedule_file = schedule_file
        self.taskdir = taskdir
        if options is None:
            options = config.Options()
        self.options = options
        self.threads = {}
        self.tasks = {}
        self.mutex = threading.Lock()
        self.directiveLock = threading.Lock()
        self.directiveTask = None
        self.started = False
        # Create the directory where task threads store their state
        if taskdir is not None and not os.path.isdir(taskdir):
            logger.info("Creating task directory %s", taskdir)
            try:
                os.makedirs(taskdir)
            except os.error as err:
                if not os.path.isdir(taskdir):
                    logger.error("Failed to create %s: %s", taskdir, err)

    def lock(self):
        """Locks the schedule object for exclusive access."""
        self.mutex.acquire()

    def unlock(self):
        """Releases the lock on the schedule object."""
        self.mutex.release()

    def load(self, filename=None):
        """Loads the scheduler's list of tasks from disk."""
        if filename is None:
            filename = self.schedule_file
        self.lock()
        try:
            try:
                f = open(filename)
            except IOError as err:
                logger.error("Can't read schedule file %s: %s", filename, err)
                return
            xmltext = f.read()
            f.close()
            try:
                task_list = tasks.parse_tasks(xmltext)
            except Exception as err:
                logger.error("Exception parsing schedule file %s: %s: %s",
                             filename, err.__class__, err)
                return
            else:
                logger.info("Read %d tasks from %s", len(task_list), filename)
            for task in task_list:
                if task.command is None:
                    if task.id in self.tasks:
                        logger.info("Deleting task %s", task.id)
                        del self.tasks[task.id]
                else:
                    if task.id in self.tasks:
                        logger.info("Replacing task %s", task.id)
                    else:
                        logger.info("Adding task %s", task.id)
                    self.tasks[task.id] = task
        finally:
            self.unlock()

    def start(self):
        """Starts the threads corresponding to the scheduled tasks."""
        self.lock()
        try:
            assert not self.threads, "already started (start)"
            for task in self.tasks.values():
                thread = TaskThread(task, self, self.options)
                self.threads[task.id] = thread
                thread.start()
            logger.info("Started %d threads", len(self.tasks))
            self.started = True
        finally:
            self.unlock()

    def shutdown(self):
        """Cancels all pending tasks, and wait for them."""
        self.lock()
        try:
            threads = self.threads.values()
            self.threads = {}
            for thread in threads:
                thread.shutdown()
            self.started = False
        finally:
            self.unlock()
        for thread in threads:
            logger.info("Waiting for thread %s ...", thread.name)
            thread.join()
            logger.info("Thread %s finished", thread.name)
        return len(threads)

    def restart(self):
        """Restart the scheduler process.

        This isn't used when running as a Windows service.

        """
        count = self.shutdown()
        args = sys.argv[:]
        args[0] = os.path.join(self.options.agentbin, args[0])
        args.insert(0, sys.executable)
        try:
            proc = process.createProcess(args)
        except process.Error as err:
            logger.error("Can't restart scheduler: %s", err)
        else:
            logger.info("Scheduler restarted")
        return count

    def doitnow(self, predicate):
        self.lock()
        try:
            for thread in self.threads.values():
                if predicate(thread.task):
                    thread.doitnow()
        finally:
            self.unlock()

    def save(self, filename=None):
        """Saves the current schedule to disk."""
        self.lock()
        try:
            if filename is None:
                filename = self.schedule_file
            new = filename + ".new"
            old = filename + ".old"
            f = open(new, "w")
            try:
                self.write(f)
            finally:
                f.close()
            if dance.rename(new, filename, logger) is None:
                logger.info("Saved schedule to %s", filename)
        finally:
            self.unlock()

    def write(self, f, details=False):
        """Writes the current schedule to the given stream."""
        tasks.Tasks(tasks=self.list_tasks(details).values()).__render__(f)
        f.flush()

    def add_list(self, tasks):
        """Adds a list of tasks to the scheduler.

        This can also replace or delete tasks -- it will replace an
        existing task with the same id, and it will delete the task if
        the command is None.

        This returns the number of tasks processed.  If this is less
        than the length of the argument, it means that some deletions
        referred to unknown tasks.

        Note: this does not save the task list -- you must explicitly
        call save() to do that.

        """
        self.lock()
        skipped = 0
        try:
            for task in tasks:
                if task.command is None:
                    # Deleting task
                    if task.id not in self.tasks:
                        logger.debug("Can't delete unknown task %s", task.id)
                        skipped += 1
                    else:
                        thread = self.threads.get(task.id)
                        if thread is not None:
                            thread.delete()
                            del self.threads[task.id]
                            logger.debug("Deleted thread %s", task.id)
                        del self.tasks[task.id]
                        logger.debug("Deleted task %s", task.id)
                elif task.id not in self.tasks:
                    # New task
                    assert task.id not in self.threads, "dup task.id (add)"
                    self.tasks[task.id] = task
                    if self.started:
                        thread = TaskThread(task, self)
                        self.threads[task.id] = thread
                        thread.start()
                    logger.debug("New task %s", task.id)
                elif task != self.tasks[task.id]:
                    # Replacement task
                    thread = self.threads.get(task.id)
                    if thread is not None:
                        thread.shutdown()
                        logger.debug("Shut down thread %s", task.id)
                    self.tasks[task.id] = task
                    if self.started:
                        thread = TaskThread(task, self)
                        self.threads[task.id] = thread
                        thread.start()
                    logger.debug("Replaced task %s", task.id)
                else:
                    # No change
                    logger.debug("Unchanged task %s", task.id)
        finally:
            self.unlock()
        count = len(tasks) - skipped
        logger.info("Updated %d out of %d tasks", count, len(tasks))
        return count

    def list_tasks(self, details=False):
        """Returns a dict mapping task ids to Tasks."""
        d = {}
        for task in self.tasks.values():
            task = copy.copy(task)
            if details:
                th = self.threads.get(task.id)
                if th is not None:
                    ts = th.state
                    assert isinstance(ts, TaskState), "bogus TS (listtasks)"
                    when = ts.marktime
                    if when is not None:
                        ymdhms = (when.year, when.month, when.day,
                                  when.hour, when.minute, when.second)
                        task.detail = datetime.datetime(tzinfo=timezones.utc,
                                                        *ymdhms)
            d[task.id] = task
        return d
