AirlockTools/services/setup.py

#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.



import json
import logging
import logging.handlers
import os
import platform
import sys
from pathlib import Path

from dotenv import load_dotenv, set_key

PROTECTED_KEYS = [
    "APPNAME",
    "PATH_EXCLUSION_CONST",
    "MIN_FILES_FOR_PATH",
    "VT_THREAT_TOLERANCE",
    "POLICY_MAP_ENF_AUD"
]

def get_base_directory() -> Path:
    system = platform.system()
    home = Path.home()
    if system == 'Windows':
        return Path(os.getenv('APPDATA', home / 'AppData' / 'Roaming')) / "AirlockTools"
    elif system == 'Darwin':
        return home / 'Library' / 'Application Support' / "AirlockTools"
    else:
        return home / '.local' / 'share' / "AirlockTools"

def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
    log_file = log_dir / "airlocktools.log"
    logger = logging.getLogger()
    logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG))

    if not logger.handlers:
        file_handler = logging.handlers.RotatingFileHandler(
            log_file, maxBytes=5_000_000, backupCount=5, encoding='utf-8'
        )
        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        logger.addHandler(file_handler)

        console_handler = logging.StreamHandler()
        console_handler.setFormatter(logging.Formatter('%(levelname)s - %(message)s'))
        logger.addHandler(console_handler)

        if platform.system() == "Windows":
            try:
                event_handler = logging.handlers.NTEventLogHandler("AirlockTools")
                event_handler.setLevel(logging.CRITICAL)
                event_handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))
                logger.addHandler(event_handler)
            except Exception as e:
                logger.warning(f"Could not attach Windows Event Log handler: {e}")

    logger.debug("Logging configured.")

def get_system_config_path() -> Path:
    base_path = Path(getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__))))
    return base_path.parent / "system_config.json"

def load_system_config() -> dict:
    try:
        config_path = get_system_config_path()
        with open(config_path, "r") as f:
            return json.load(f)
    except FileNotFoundError:
        logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
        return {
            "APPNAME": "AirlockTools",
            "LOG_LEVEL": "DEBUG",
            "PATH_EXCLUSION_CONST": 4,
            "MIN_FILES_FOR_PATH": 4,
            "VT_THREAT_TOLERANCE": 4,
            "POLICY_MAP_ENF_AUD": {
                "enforced_id": "audit_id"
            }
        }

def load_user_config(config_dir: Path) -> dict:
    user_config_path = config_dir / "user_config.json"
    if not user_config_path.exists():
        default_user_config = {
            "URL": "",
            "LOG_LEVEL": "INFO"
        }
        with open(user_config_path, "w") as f:
            json.dump(default_user_config, f, indent=4)
        logging.debug(f"Created user config at {user_config_path}")
    with open(user_config_path, "r") as f:
        return json.load(f)

def write_config_to_env(config: dict, env_path: Path):
    for key, value in config.items():
        try:
            serialized = json.dumps(value) if isinstance(value, (list, dict)) else str(value)
            set_key(env_path, key, serialized)
        except Exception as e:
            logging.warning(f"Failed to write {key} to .env: {e}")

def setup() -> Path:
    base_dir = get_base_directory()
    dirs = {
        'config': base_dir / 'config',
        'cache': base_dir / 'cache',
        'logs': base_dir / 'logs',
    }

    for name, path in dirs.items():
        path.mkdir(parents=True, exist_ok=True)
        logging.debug(f"{name.capitalize()} directory ensured at: {path}")

    system_config = load_system_config()
    configure_logging(dirs['logs'], system_config.get("LOG_LEVEL", "DEBUG"))

    env_path = base_dir / ".env"
    if not env_path.exists():
        env_path.touch()
    load_dotenv(dotenv_path=env_path, override=True)

    working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data"))
    working_dir.mkdir(parents=True, exist_ok=True)
    set_key(env_path, "WORKING_DIR", str(working_dir))
    os.environ["WORKING_DIR"] = str(working_dir)
    logging.debug(f"Working directory set to: {working_dir}")

    folders_structure = {
        "Approved": [],
        "Needs_Review": ["Review_First", "Review_Second", "HTML"],
        "Preflight": ["HTML"],
        "Archived": ["HTML"],
        "Scheduling": []
    }

    for folder_name, subfolders in folders_structure.items():
        folder_path = working_dir / folder_name
        folder_path.mkdir(parents=True, exist_ok=True)
        logging.debug(f"'{folder_name}' folder ensured at: {folder_path}")
        for subfolder in subfolders:
            subfolder_path = folder_path / subfolder
            subfolder_path.mkdir(parents=True, exist_ok=True)
            logging.debug(f"  └─ '{subfolder}' subfolder created at: {subfolder_path}")

    user_config = load_user_config(dirs['config'])
    merged_config = {**system_config, **user_config}

    for key in PROTECTED_KEYS:
        merged_config[key] = system_config.get(key, "")

    # ✅ URL resolution order: system_config → .env → user prompt
    url = system_config.get("URL")
    if not url:
        url = os.getenv("URL")
    if not url:
        url = input("🌐 Enter the service URL (e.g., https://example.com/api): ").strip()

    merged_config["URL"] = url
    set_key(env_path, "URL", url)
    os.environ["URL"] = url
    logging.debug(f"Service URL set to: {url}")

    write_config_to_env(merged_config, env_path)

    return working_dir