from ...Machine.PaedMLMachineCollection import MACHINE_COLLECTION
from ..SetupStep import SetupStep
from ..Playbooks.Playbook import FileCleanupTask, PlaybookCredentialCollection
from ...Logging.Logger import log, Severity

class SetFirewallPassword(SetupStep):
    def __init__(self, username, password):
        super().__init__()
        self.username = username
        self.password = password

    def get_crypt(self) -> str|None:
        stdin, stdout, stderr = MACHINE_COLLECTION.server.execute_shell_command_live("mkpasswd","--stdin", "-m", 'bcrypt')
        stdin.write(self.password.encode())
        stdin.write('\n'.encode())
        stdin.flush()

        if stderr.readable():
            error = stderr.read().strip().decode()
            if len(error) > 0:
                log(error, Severity.ERROR)
                return None

        hash = stdout.readline()
        if isinstance(hash, bytes):
            hash = hash.decode()
        return hash.strip()

    def run(self, file_cleanup_stack : list[FileCleanupTask], credentials: PlaybookCredentialCollection) -> bool:
        log(f'Setting firewall password for {self.username}...')
        crypt_hash = self.get_crypt()
        if crypt_hash is None:
            log("Unable to calculate hash.", Severity.ERROR)
            return False

        log(f"Crypt hash is {crypt_hash}", Severity.TRACE)

        master_passwd = MACHINE_COLLECTION.firewall.read_file('/etc/master.passwd').decode()
        lines = master_passwd.splitlines(True)
        new_lines = []
        user_found = False

        for line in lines:
            if line.startswith(f'{self.username}:'):
                parts = line.split(':')
                if len(parts) >= 2:
                    parts[1] = crypt_hash
                    new_lines.append(':'.join(parts))
                    user_found = True
                    continue
            new_lines.append(line)

        if not user_found:
            log(f"User {self.username} not found in /etc/master.passwd", Severity.ERROR)
            return False

        MACHINE_COLLECTION.firewall.write_file('/etc/master.passwd', ''.join(new_lines).encode())
        return True