AirlockTools/utils/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.config
import logging.handlers
import os
import platform
import sys
from pathlib import Path

from dotenv import load_dotenv, set_key

from utils.configmanager import PROTECTED_KEYS, load_protected_config


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"

    config = {
        "version": 1,  # Required key for dictConfig format version
        "disable_existing_loggers": False,  # Keeps existing loggers active
        "formatters": {
            "detailed": {
                "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
                # Includes timestamp, logger name, level, and message
            },
            "simple": {
                "format": "%(levelname)s - %(message)s"
                # Minimal format for console output
            },
        },
        "handlers": {
            "file": {
                "class": "logging.handlers.TimedRotatingFileHandler",
                "filename": str(log_file),
                "when": "midnight",       # Rotate logs at midnight
                "interval": 1,            # Every 1 day
                "backupCount": 7,         # Keep 7 days of logs
                "encoding": "utf-8",        # Ensure UTF-8 encoding
                "level": "DEBUG",           # Always log DEBUG and above
                "formatter": "detailed",    # Use detailed format
            },
            "console": {
                "class": "logging.StreamHandler",
                "level": log_level.upper(),  # Configurable log level
                "formatter": "simple",       # Use simple format
            },
        },
        "root": {
            "level": "DEBUG",  # Root logger level
            "handlers": ["file", "console"],  # Attach both handlers
        },
    }

    # Add Windows Event Log handler if on Windows
    if platform.system() == "Windows":
        try:
            config["handlers"]["eventlog"] = {
                "class": "logging.handlers.NTEventLogHandler",
                "appname": "AirlockTools",  # Event log source name
                "level": "CRITICAL",        # Only log critical errors
                "formatter": "simple",      # Use simple format
            }
            config["root"]["handlers"].append("eventlog")
        except Exception as e:
            logging.warning(f"Could not attach Windows Event Log handler: {e}")

    # Apply the logging configuration
    logging.config.dictConfig(config)
    logging.getLogger().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():
        if key in PROTECTED_KEYS:
            continue  # Skip protected keys
        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():
    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": []
    }

    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}

    protected_config = load_protected_config()
    merged_config.update(protected_config)

    # ✅ 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)