from automatagl.helpers.config_operations import sanitize_username


class SSHKeyObject:
    """
    An SSH key storage object to make things a bit easier to manage.
    """

    def __init__(self, username: str = '') -> None:
        """
        Create the SSH key object
        :param username: the username of the account with the SSH keys
        """
        self.ssh_keys = list()
        self.username = sanitize_username(username)

    def add_keys(self, ssh_keys) -> None:
        """
        Add one or more keys to the SSH key object
        :param ssh_keys: The ssk_key(s) to add.
        :return:
        """
        if type(ssh_keys) is list:
            self.ssh_keys = ssh_keys
        elif type(ssh_keys) is str:
            self.ssh_keys = [ssh_keys]
        else:
            raise TypeError("SSH keys must either be type 'list' or 'str'.")

    def get_authorized_keys(self) -> str:
        """
        Generate the authorized keys file.
        :return: The contents of the authorized_keys file.
        """
        data = "# This file is automatically generated by Automata.\n"
        data += "\n".join(self.ssh_keys) + "\n"
        return data
