#!/usr/bin/env python2.7
import sys
import os
import subprocess
import time
import shutil


def print_help():
    print "Usage: suricata-jobs [OPTIONS]."
    print ""
    print "Options:"
    print ""
    print "   --restart - Restarts the service"
    print "   --reload  - Reloads the configurations"
    print "   --update  - Downloads updates and applies changes to signatures"
    print "   --status  - Checks if the Suricata service is running"


def modify_sid():
    input_file = "/var/efw/suricata/modify-sid"
    output_file = "/etc/suricata/modify.conf"

    if not os.path.exists(input_file):
        return

    with open(output_file, "w") as f_out, open(input_file, "r") as f_in:
        for line in f_in:
            if ',' not in line:
                continue
            parts = line.split(',', 1)
            action = parts[0].replace(" ", "").strip().upper()
            sid = parts[1].replace(" ", "").strip()

            if not action or not sid:
                continue
            if action.startswith("#"):
                continue

            if action == "ALERT":
                f_out.write('modifysid {} "drop" | "alert"\n'.format(sid))
            elif action == "DROP":
                f_out.write('modifysid {} "alert" | "drop"\n'.format(sid))


def suppress_sid():
    input_file = "/var/efw/suricata/suppress"
    output_file = "/etc/suricata/threshold.config"

    if not os.path.exists(input_file):
        return


    with open(output_file, "w") as f_out, open(input_file, "r") as f_in:
        for line in f_in:
            if ',' not in line:
                continue
            parts = line.split(',')
            if len(parts) < 3:
                continue
            direction = parts[0].replace(" ", "").strip().upper()
            ip = parts[1].replace(" ", "").strip()
            sid = parts[2].replace(" ", "").strip()

            if not direction or not ip or not sid:
                continue
            if direction.startswith("#"):
                continue

            if direction == "ORIGEM":
                f_out.write("suppress gen_id 1, sig_id {}, track by_src, ip {}\n".format(sid, ip))
            elif direction == "DESTINO":
                f_out.write("suppress gen_id 1, sig_id {}, track by_dst, ip {}\n".format(sid, ip))


def apply_category_rules():
    input_file = "/var/efw/suricata/category-rules"
    drop_file = "/etc/suricata/drop.conf"
    enable_file = "/etc/suricata/enable.conf"
    disable_file = "/etc/suricata/disable.conf"

    for f in [drop_file, enable_file, disable_file]:
        with open(f, "w") as _:
            pass

    if not os.path.exists(input_file):
        return

    with open(drop_file, "a") as f_drop, open(enable_file, "a") as f_enable, open(disable_file, "a") as f_disable, open(input_file, "r") as f_in:
        for line in f_in:
            if '=' not in line:
                continue
            parts = line.split('=', 1)
            category = parts[0].strip()
            value = parts[1].strip()

            if not category or category.startswith("#"):
                continue

            if value == "drop":
                f_drop.write("group:{}\n".format(category))
                f_enable.write("group:{}\n".format(category))
            elif value == "off":
                f_disable.write("group:{}\n".format(category))
            elif value == "default":
                f_enable.write("group:{}\n".format(category))
            else:
                print "Unknown value '{}' for category {}".format(value, category)


def set_update_schedule():
    settings_file = "/var/efw/suricata/settings"
    cron_file = "/etc/cron.d/suricata-update"

    enabled_rules = ""
    schedule = ""

    if os.path.exists(settings_file):
        with open(settings_file, "r") as f:
            for line in f:
                if line.startswith("ENABLED_RULES="):
                    enabled_rules = line.split("=", 1)[1].strip()
                elif line.startswith("UPDATE_SCHEDULE="):
                    schedule = line.split("=", 1)[1].strip()

    if enabled_rules != "auto":
        print "ENABLED_RULES={}, removing schedule if it exists...".format(enabled_rules)
        if os.path.exists(cron_file):
            try:
                os.remove(cron_file)
            except OSError:
                pass
        
        with open(os.devnull, 'w') as devnull:
            subprocess.call("/etc/init.d/fcron restart", shell=True, stdout=devnull, stderr=devnull)
        return

    with open(cron_file, "w") as f_cron:
        if schedule == "daily":
            f_cron.write("0 23 * * * /usr/local/bin/suricata-jobs --update\n")
            with open(os.devnull, 'w') as devnull:
                subprocess.call("/etc/init.d/fcron restart", shell=True, stdout=devnull, stderr=devnull)
        elif schedule == "weekly":
            f_cron.write("0 23 * * 0 /usr/local/bin/suricata-jobs --update\n")
            with open(os.devnull, 'w') as devnull:
                subprocess.call("/etc/init.d/fcron restart", shell=True, stdout=devnull, stderr=devnull)
        elif schedule == "monthly":
            f_cron.write("0 23 1 * * /usr/local/bin/suricata-jobs --update\n")
            with open(os.devnull, 'w') as devnull:
                subprocess.call("/etc/init.d/fcron restart", shell=True, stdout=devnull, stderr=devnull)
        else:
            print "Unknown value '{}' for UPDATE_SCHEDULE.".format(schedule)
            sys.exit(1)


def update():
    settings_file = "/var/efw/suricata/settings"
    enabled = ""

    if os.path.exists(settings_file):
        with open(settings_file, "r") as f:
            for line in f:
                if line.startswith("ENABLED="):
                    enabled = line.split("=", 1)[1].strip()

    if enabled != "1":
        timestamp = time.strftime("%d/%m/%Y -- %H:%M:%S")
        with open("/var/log/suricata/suricata.log", "a") as log:
            log.write("{} - <Warn> -- Suricata is disabled (ENABLED={}), skipping update\n".format(timestamp, enabled))
        print "Suricata is disabled. Skipping update."
        return 0

    try:
        with open(os.devnull, 'w') as devnull:
            pgrep_res = subprocess.call(["pgrep", "-f", "/usr/bin/suricata-update"], stdout=devnull)
        if pgrep_res == 0:
            timestamp = time.strftime("%d/%m/%Y -- %H:%M:%S")
            with open("/var/log/suricata/suricata.log", "a") as log:
                log.write("{} - <Warn> -- suricata-update is already running\n".format(timestamp))
            print "suricata-update is already running."
            return 0
    except OSError:
        pass

    print "Running suricata-update..."
    timestamp = time.strftime("%d/%m/%Y -- %H:%M:%S")
    with open("/var/log/suricata/suricata.log", "a") as log:
        log.write("{} - <Info> -- Rules update started.....\n".format(timestamp))

    with open("/var/log/suricata/suricata.log", "a") as log:
        status_code = subprocess.call("/usr/bin/suricata-update", shell=True, stdout=log, stderr=log)

    if not os.path.exists("/var/efw/download"):
        try:
            os.makedirs("/var/efw/download")
        except OSError:
            pass

    try:
        shutil.chown("/var/efw/download", user="nobody", group="nogroup")
    except (LookupError, AttributeError):
        try:
            import pwd
            import grp
            uid = pwd.getpwnam("nobody").pw_uid
            gid = grp.getgrnam("nogroup").gr_gid
            os.chown("/var/efw/download", uid, gid)
        except Exception:
            pass
    except Exception:
        pass

    if status_code == 0:
        ts_file = "/var/efw/download/timestamps"
        key = "SURICATA_SURICATA_RULES_URL"
        value = str(int(time.time()))

        lines = []
        key_found = False
        if os.path.exists(ts_file):
            with open(ts_file, "r") as f:
                for line in f:
                    # Remove any line that contains "SNORT"
                    if "SNORT" in line:
                        continue
                    if line.startswith("{}=".format(key)):
                        lines.append("{}={}\n".format(key, value))
                        key_found = True
                    else:
                        lines.append(line)
        
        if not key_found:
            lines.append("{}={}\n".format(key, value))

        with open(ts_file, "w") as f:
            f.writelines(lines)

        revision_file = "/var/lib/suricata/rules/suricatarules.revision"
        if not os.path.exists(os.path.dirname(revision_file)):
            try:
                os.makedirs(os.path.dirname(revision_file))
            except OSError:
                pass
        with open(revision_file, "a") as _:
            try:
                os.utime(revision_file, None)
            except OSError:
                pass

        timestamp = time.strftime("%d/%m/%Y -- %H:%M:%S")
        with open("/var/log/suricata/suricata.log", "a") as log:
            log.write("{} - <Info> -- suricata-update completed successfully\n".format(timestamp))
        print "suricata-update executed successfully."
    else:
        timestamp = time.strftime("%d/%m/%Y -- %H:%M:%S")
        with open("/var/log/suricata/suricata.log", "a") as log:
            log.write("{} - <Error> -- Failed to execute suricata-update\n".format(timestamp))
        print "Error executing suricata-update."
        sys.exit(1)


def set_startup_hook():
    settings_file = "/var/efw/suricata/settings"
    hook_file = "/var/efw/inithooks/start.local"
    enabled = ""

    if os.path.exists(settings_file):
        with open(settings_file, "r") as f:
            for line in f:
                if line.startswith("ENABLED="):
                    enabled = line.split("=", 1)[1].strip()

    if not os.path.exists(os.path.dirname(hook_file)):
        try:
            os.makedirs(os.path.dirname(hook_file))
        except OSError:
            pass

    with open(hook_file, "a") as _:
        pass
    os.chmod(hook_file, 0o755)

    if enabled == "1":
        subprocess.call("/etc/init.d/suricata restart", shell=True)
        subprocess.call("monit monitor suricata", shell=True)

        with open(hook_file, "r") as f:
            lines = f.readlines()

        if not lines or not lines[0].startswith("#!/bin/bash"):
            lines.insert(0, "#!/bin/bash\n")

        cmd = "/usr/local/bin/suricata-jobs --restart\n"
        if cmd not in lines and cmd.strip() not in [l.strip() for l in lines]:
            lines.append(cmd)

        with open(hook_file, "w") as f:
            f.writelines(lines)
    else:
        print "Suricata is disabled!"
        if os.path.exists(hook_file):
            with open(hook_file, "r") as f:
                lines = f.readlines()
            lines = [l for l in lines if "/usr/local/bin/suricata-jobs --restart" not in l]
            with open(hook_file, "w") as f:
                f.writelines(lines)

        subprocess.call("monit unmonitor suricata", shell=True)
        subprocess.call("/etc/init.d/suricata stop", shell=True)


def restart():
    print "Restarting Suricata..."
    set_startup_hook()
    modify_sid()
    suppress_sid()
    apply_category_rules()
    set_update_schedule()
    
    res = subprocess.call("/etc/rc.d/rc.firewall restart", shell=True)
    if res == 0:
        print "Suricata restarted successfully."
    else:
        print "Error restarting Suricata."
        sys.exit(1)


def reload():
    print "Reloading Suricata configurations..."
    modify_sid()
    suppress_sid()
    apply_category_rules()
    
    res = subprocess.call("/etc/init.d/suricata reload", shell=True)
    if res == 0:
        print "Configurations reloaded successfully."
    else:
        print "Error reloading configurations."
        sys.exit(1)


def status():
    print "Checking Suricata status..."
    res = subprocess.call("/etc/init.d/suricata status", shell=True)
    if res == 0:
        print "Suricata is active."
    else:
        print "Suricata is not running."
        sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print_help()
        sys.exit(0)

    arg = sys.argv[1]
    if arg == "--help":
        print_help()
        sys.exit(0)
    elif arg == "--update":
        update()
        sys.exit(0)
    elif arg == "--restart":
        restart()
        sys.exit(0)
    elif arg == "--reload":
        reload()
        sys.exit(0)
    elif arg == "--status":
        status()
        sys.exit(0)
    else:
        print_help()
        sys.exit(0)
