gentoo-legion-python/install/users.py

"""User configuration and shell setup."""

import re
import shutil
from pathlib import Path

from .utils import info, success, error, run, run_quiet, emerge, prompt


# Shell and portage tool packages
SHELL_PACKAGES = [
    # Shell
    "app-shells/zsh",
    "app-shells/gentoo-zsh-completions",
    "app-shells/zsh-autosuggestions",
    "app-shells/zsh-syntax-highlighting",
    # Prompt & tools
    "app-shells/starship",
    "app-shells/zoxide",
    "app-shells/fzf",
    "sys-apps/lsd",
    # Portage tools
    "app-portage/gentoolkit",
    "app-portage/portage-utils",
    "app-portage/eix",
]

# User groups for desktop use
# Minimal set - elogind grants device access to active session
USER_GROUPS = [
    "users",
    "wheel",   # sudo access
    "input",   # Wayland compositor input devices
]


def emerge_shell_packages() -> None:
    """Emerge shell and portage tool packages."""
    info("=== Installing Shell & Portage Tools ===")

    print()
    info("Packages to install:")
    for pkg in SHELL_PACKAGES:
        print(f"  - {pkg}")
    print()

    emerge(*SHELL_PACKAGES)

    success("Shell packages installed.")


def install_sudo() -> None:
    """Install and configure sudo."""
    info("=== Installing sudo ===")

    emerge("app-admin/sudo")

    # Configure sudoers - enable wheel group
    sudoers = Path("/etc/sudoers")
    if sudoers.exists():
        content = sudoers.read_text()

        # Check if wheel group already has sudo access (any format)
        if re.search(r"^%wheel\s+ALL=", content, re.MULTILINE):
            success("Wheel group already enabled in sudoers")
        elif re.search(r"^#\s*%wheel\s+ALL=", content, re.MULTILINE):
            # Uncomment the existing commented line
            content = re.sub(
                r"^#\s*(%wheel\s+ALL=.*$)",
                r"\1",
                content,
                flags=re.MULTILINE
            )
            sudoers.write_text(content)
            success("Enabled wheel group in sudoers")
        else:
            # Append if not present at all
            with open(sudoers, "a") as f:
                f.write("\n# Allow wheel group to use sudo\n")
                f.write("%wheel ALL=(ALL:ALL) ALL\n")
            success("Added wheel group to sudoers")

    success("sudo configured.")


def copy_shell_configs(source_dir: Path, target_home: Path, owner: str | None = None) -> None:
    """Copy shell configuration files to a home directory."""
    # .zshrc
    zshrc_src = source_dir / ".zshrc"
    if zshrc_src.exists():
        shutil.copy2(zshrc_src, target_home / ".zshrc")
        success(f"Copied .zshrc to {target_home}")

    # starship.toml
    starship_src = source_dir / "starship.toml"
    if starship_src.exists():
        config_dir = target_home / ".config"
        config_dir.mkdir(parents=True, exist_ok=True)
        shutil.copy2(starship_src, config_dir / "starship.toml")
        success(f"Copied starship.toml to {config_dir}")

    # Fix ownership if specified
    if owner:
        run("chown", "-R", f"{owner}:{owner}", str(target_home), check=False)


def setup_root_shell(source_dir: Path) -> None:
    """Configure root's shell environment."""
    info("=== Configuring Root Shell ===")

    copy_shell_configs(source_dir, Path("/root"))

    # Change root's shell to zsh
    run("chsh", "-s", "/bin/zsh", "root")
    success("Root shell set to zsh")


def set_root_password() -> None:
    """Prompt to set root password."""
    info("=== Set Root Password ===")
    print()
    run("passwd")
    print()


def create_user(username: str, source_dir: Path) -> None:
    """Create a new user with proper groups and shell config."""
    info(f"=== Creating User: {username} ===")

    # Build groups string
    groups = ",".join(USER_GROUPS)

    # Check if user already exists using id command (more reliable than string matching)
    result = run_quiet("id", username, check=False)
    if result.ok:
        info(f"User {username} already exists, ensuring groups are correct...")
        run("usermod", "-aG", groups, username, check=False)
    else:
        # Create user
        result = run(
            "useradd", "-m",
            "-G", groups,
            "-s", "/bin/zsh",
            username,
            check=False
        )
        if not result.ok:
            error(f"Failed to create user: {result.stderr}")
            return

    success(f"User {username} configured with groups: {groups}")

    # Copy shell configs
    user_home = Path(f"/home/{username}")
    copy_shell_configs(source_dir, user_home, owner=username)

    # Set user password
    print()
    info(f"Set password for {username}:")
    run("passwd", username)
    print()

    success(f"User {username} configured.")


def setup_users(source_dir: Path | None = None) -> None:
    """Full user setup workflow."""
    if source_dir is None:
        source_dir = Path("/root/gentoo")

    emerge_shell_packages()
    print()
    install_sudo()
    print()
    setup_root_shell(source_dir)
    print()
    set_root_password()

    # Prompt for username
    print()
    info("=== Create User Account ===")
    username = prompt("Enter username to create: ").strip()

    if username:
        create_user(username, source_dir)
    else:
        info("Skipped user creation.")

    print()
    success("=== User Setup Complete ===")
    print()
    info("Next steps:")
    print("  1. Install @hyprland set for desktop environment")
    print("  2. Configure display manager (SDDM)")
    print("  3. Reboot and login as your user")