#!/opt/alt/python311/bin/python3

import os
import glob
from subprocess import run, CalledProcessError, DEVNULL

# ubuntu + plesk
# /etc/apache2/mods-available/security2.conf
# cpanel
# /etc/apache2/conf.d/modsec/modsec2.cpanel.conf
# centos + plesk
# /etc/httpd/conf.d/mod_security.conf:
# /etc/httpd/conf.d/security2.conf:
# Directadmin

MODSEC_CONFIG_LIST = [
    "/etc/apache2/mods-available/security2.conf",
    "/etc/apache2/conf.d/modsec/modsec2.cpanel.conf",
    "/etc/httpd/conf.d/mod_security.conf",
    "/etc/httpd/conf.d/security2.conf",
    "/etc/imunify360-webshield/modsecurity/modsecurity.conf",
]

SDBM_UTIL = "/opt/alt/modsec-sdbm-util/bin/modsec-sdbm-util"
DIR = ".dir"
PAG = ".pag"
NEW_DB = "new_db"
NEW_PAG = NEW_DB + PAG
NEW_DIR = NEW_DB + DIR


def get_persistent_path_list():
    MODSEC_DIRECTIVE = "SecDataDir"
    result = []
    for config in MODSEC_CONFIG_LIST:
        try:
            with open(config, "r") as c:
                for line in c:
                    if line.startswith(MODSEC_DIRECTIVE):
                        sec_dir = (
                            line.replace(MODSEC_DIRECTIVE, "")
                            .replace('"', "")
                            .strip()
                        )
                        result.append(sec_dir)
                        break
        except (OSError, IOError):
            pass
    return result


def _run_cmd(cmd):
    try:
        out = run(cmd, stderr=DEVNULL)
    except (FileNotFoundError, CalledProcessError):
        return None
    return out.returncode


def _unlink_old_db(path):
    try:
        os.unlink(path + PAG)
    except OSError:
        pass
    try:
        os.unlink(path + DIR)
    except OSError:
        pass


def move_with_drop_perm_and_owner(path, old_base_name, uid, gid, perm):
    try:
        old_pag = os.path.join(path, old_base_name + PAG)
        old_dir = os.path.join(path, old_base_name + DIR)
        os.rename(os.path.join(path, NEW_PAG), old_pag)
        os.rename(os.path.join(path, NEW_DIR), old_dir)
        os.chmod(old_pag, perm)
        os.chmod(old_dir, perm)
        os.chown(old_pag, uid, gid)
        os.chown(old_dir, uid, gid)
    except OSError:
        pass


def shrink_dir(path):
    for db in glob.glob(os.path.join(path, "*" + PAG)):
        db_base_name = os.path.splitext(db)[0]
        if os.path.exists(db_base_name + PAG) and os.path.exists(
            db_base_name + DIR
        ):
            db_file_stat = os.stat(db_base_name + PAG)
            rc = _run_cmd([SDBM_UTIL, "-D", path, "-v", "-n", db])
            if (
                rc is not None
                and not rc
                and os.path.exists(os.path.join(path, NEW_PAG))
                and os.path.exists(os.path.join(path, NEW_DIR))
            ):
                _unlink_old_db(os.path.join(path, db_base_name))
                move_with_drop_perm_and_owner(
                    path,
                    db_base_name,
                    db_file_stat.st_uid,
                    db_file_stat.st_gid,
                    db_file_stat.st_mode,
                )


def main():
    persistent_dirs = get_persistent_path_list()
    for path in persistent_dirs:
        shrink_dir(path)


if __name__ == "__main__":
    main()