import os

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

class EnsureSshKeysExist(SetupStep):
    def __init__(self):
        super().__init__()

    def run(self, file_cleanup_stack : list[FileCleanupTask], credentials: PlaybookCredentialCollection) -> bool:
        if os.path.isfile(LMZ_COMMON_CREDENTIALS.get_private_keyfile()) and os.path.isfile(LMZ_COMMON_CREDENTIALS.get_public_keyfile()):
            return True
        else:
            log("No SSH keys found. Recreating using ssh-keygen...", Severity.DEBUG)
            stdout, stderr, retval = MACHINE_COLLECTION.local_machine.execute_shell_command("ssh-keygen", '-t', 'rsa', '-b', '4096', '-f', LMZ_COMMON_CREDENTIALS.get_private_keyfile(), '-N', '')
            if retval != 0:
                log(stderr, Severity.ERROR)
                return False
            else:
                log("SSH keys have been generated.", Severity.DEBUG)

            success = os.path.isfile(LMZ_COMMON_CREDENTIALS.get_private_keyfile()) and os.path.isfile(LMZ_COMMON_CREDENTIALS.get_public_keyfile())
            if not success:
                log ("No SSH keys exist on the local machine.", Severity.ERROR)
            return success