from pathlib import Path

from pyinfra.api import deploy
from pyinfra.operations import files
from pyinfra.operations import systemd

BASE_DIR = Path(__file__).parent.parent


def deploy_ssh_keys():
    files.file(
        name="Create authorized_keys file",
        path="/root/.ssh/authorized_keys",
    )

    for key_path in BASE_DIR.glob("pubkeys/*.pub"):
        with open(key_path, "r") as f:
            key = f.read().strip()
        files.line(
            name=f"Adding key {key_path.name} to /root/.ssh/authorized_keys",
            path="/root/.ssh/authorized_keys",
            line=key,
        )


def reconfigure_ssh():
    config_changed = files.line(
        name="Disable password login",
        path="/etc/ssh/sshd_config",
        line="PasswordAuthentication .+",
        replace="PasswordAuthentication no",
    ).changed

    systemd.service(
        name="Restart SSHD service",
        service="ssh",
        restarted=config_changed,
    )


@deploy
def setup_ssh():
    deploy_ssh_keys()
    reconfigure_ssh()