parameters.py

"""Command-line and environment parameter handling for SMBLoot."""

from __future__ import annotations

import argparse
import os
import re
from collections.abc import Sequence

HASH_PATTERN = re.compile(r"^[0-9a-fA-F]{32}$")


def _parse_tree_workers(value: str) -> int:
    try:
        workers = int(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("must be an integer") from exc
    if not 1 <= workers <= 16:
        raise argparse.ArgumentTypeError("must be between 1 and 16")
    return workers


def _parse_hash(value: str) -> str:
    parts = value.split(":")
    valid_nt_hash = len(parts) == 1 and bool(HASH_PATTERN.fullmatch(parts[0]))
    valid_hash_pair = (
        len(parts) == 2
        and (not parts[0] or bool(HASH_PATTERN.fullmatch(parts[0])))
        and bool(HASH_PATTERN.fullmatch(parts[1]))
    )
    if not valid_nt_hash and not valid_hash_pair:
        raise argparse.ArgumentTypeError(
            "hash must be a 32-character NT hash or LMHASH:NTHASH"
        )
    return value.casefold()


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Browse SMB shares and display remote file contents."
    )
    server = os.getenv("SMB_TUI_SERVER")
    username = os.getenv("SMB_TUI_USERNAME")
    environment_password = os.getenv("SMB_TUI_PASSWORD")
    environment_hash = os.getenv("SMB_TUI_HASH")

    parser.add_argument(
        "--server",
        type=str,
        default=server,
        required=server is None,
    )
    parser.add_argument(
        "--remote-name",
        type=str,
        default=os.getenv("SMB_TUI_REMOTE_NAME"),
        help="NetBIOS server name when it differs from --server",
    )
    parser.add_argument(
        "--mode",
        type=str.upper,
        choices=("LIVE", "OPSEC"),
        default=os.getenv("SMB_TUI_MODE", "LIVE").upper(),
        help="operating mode displayed in the interface (default: LIVE)",
    )
    parser.add_argument(
        "--username",
        type=str,
        default=username,
    )
    parser.add_argument(
        "--kerberos",
        action="store_true",
        default=False,
        help="authenticate using kerberos",
    )
    authentication = parser.add_mutually_exclusive_group()
    authentication.add_argument(
        "--password",
        type=str,
        default=None,
        help="authenticate with a password",
    )
    authentication.add_argument(
        "--hash",
        dest="hash_value",
        type=_parse_hash,
        default=None,
        help="authenticate with an NT hash or LMHASH:NTHASH",
    )
    authentication.add_argument(
        "--no-pass",
        action="store_true",
        default=False,
        help="Don't ask for a password (useful for --kerberos)",
    )
    parser.add_argument(
        "--domain",
        type=str,
        default=os.getenv("SMB_TUI_DOMAIN", ""),
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("SMB_TUI_PORT", "445")),
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=int(os.getenv("SMB_TUI_TIMEOUT", "30")),
    )
    parser.add_argument(
        "--tree-workers",
        type=_parse_tree_workers,
        default=_parse_tree_workers(os.getenv("SMB_TUI_TREE_WORKERS", "4")),
        help="parallel SMB sessions for tree scans/downloads (default: 4)",
    )
    parser.add_argument(
        "--share",
        type=str,
        default=os.getenv("SMB_TUI_SHARE"),
        help="optional share to open initially",
    )
    parser.add_argument(
        "--path",
        type=str,
        default=os.getenv("SMB_TUI_PATH", ""),
        help="initial path within --share",
    )
    parser.add_argument(
        "--loot-dir",
        type=str,
        default=os.getenv("SMB_TUI_LOOT_DIR", "loot"),
        help="local download directory (default: ./loot)",
    )
    parser.add_argument(
        "--auto-loot-preview",
        action=argparse.BooleanOptionalAction,
        default=os.getenv("SMB_TUI_AUTO_LOOT_PREVIEW", "").casefold()
        in {"1", "true", "yes", "on"},
        help="automatically download files opened for preview (default: disabled)",
    )
    parser.add_argument(
        "--highlight-sensitive-files",
        action=argparse.BooleanOptionalAction,
        default=os.getenv("SMB_TUI_HIGHLIGHT_SENSITIVE_FILES", "true").casefold()
        in {"1", "true", "yes", "on"},
        help="display potentially sensitive filenames in orange (default: enabled)",
    )
    args = parser.parse_args(argv)
    if args.username is None and args.kerberos is True:
        # Username need to be an empty string and not None else it will fail in impacket/krb5/types.py name.setComponentByName('name-type', int(self.type))
        args.username = ""
    if args.password is None and args.hash_value is None:
        if environment_password is not None and environment_hash is not None:
            parser.error("SMB_TUI_PASSWORD and SMB_TUI_HASH are mutually exclusive")
        if environment_password is not None:
            args.password = environment_password
        elif environment_hash is not None:
            try:
                args.hash_value = _parse_hash(environment_hash)
            except argparse.ArgumentTypeError as exc:
                parser.error(f"SMB_TUI_HASH: {exc}")
        elif args.no_pass:
            args.password = None
        else:
            parser.error("one of --password, --hash or --no-pass is required")
    return args