Skip to content

VM Tool SSH

SSHSetup

A class to set up SSH configuration and keys for a VM.

Attributes:

Name Type Description
hostname str

The hostname of the VM.
username str

The username for SSH login.
password str

The password for SSH login.
email str

The email for SSH key generation.
private_key_path str

The path to the private SSH key.
client SSHClient

The SSH client.
Source code in vm_tool/ssh.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
class SSHSetup:
    """
    A class to set up SSH configuration and keys for a VM.

    Attributes:
        hostname (str): The hostname of the VM.
        username (str): The username for SSH login.
        password (str): The password for SSH login.
        email (str): The email for SSH key generation.
        private_key_path (str): The path to the private SSH key.
        client (paramiko.SSHClient): The SSH client.
    """

    def __init__(self, hostname, username, password, email):
        """
        The constructor for SSHSetup class.

        Parameters:
            hostname (str): The hostname of the VM.
            username (str): The username for SSH login.
            password (str): The password for SSH login.
            email (str): The email for SSH key generation.
        """
        self.hostname = hostname
        self.username = username
        self.password = password
        self.email = email
        self.private_key_path = os.path.expanduser("~/.ssh/id_rsa")
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # nosec B507

    def __generate_ssh_key(self, email):
        """
        Generates an SSH key pair.

        Parameters:
            email (str): The email for SSH key generation.
        """

        logger.info(f"Generating SSH key for email: {email}")
        try:
            subprocess.run(
                [
                    "ssh-keygen",
                    "-t",
                    "rsa",
                    "-b",
                    "4096",
                    "-C",
                    email,
                    "-f",
                    self.private_key_path,
                    "-N",
                    "",
                ],
                check=True,
                capture_output=True,
            )
            logger.info("SSH key generated successfully.")
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to generate SSH key: {e.stderr.decode()}")
            raise RuntimeError(f"Failed to generate SSH key: {e.stderr.decode()}")

    def __read_or_generate_public_key(self, email) -> str:
        """
        Reads the public SSH key or generates a new one if it doesn't exist.

        Parameters:
            email (str): The email for SSH key generation.

        Returns:
            str: The public SSH key.
        """
        public_key_path = f"{self.private_key_path}.pub"
        if not os.path.exists(public_key_path):
            self.__generate_ssh_key(email)
        with open(public_key_path, "r") as file:
            logger.info(f"Reading public key from {public_key_path}")
            return file.read()

    def __configure_vm(self, vm_ip, vm_password, public_key):
        """
        Configures the VM by adding the public SSH key to the authorized keys.

        Parameters:
            vm_ip (str): The IP address of the VM.
            vm_password (str): The password for the VM.
            public_key (str): The public SSH key.
        """
        logger.info(f"Configuring VM at {vm_ip} with provided public key.")
        self.client.connect(vm_ip, username=self.username, password=vm_password)
        self.client.exec_command(
            f'echo "{public_key}" >> ~/.ssh/authorized_keys'
        )  # nosec B601
        self.client.close()

    def __update_ssh_config(self):
        """
        Updates the local SSH config file with the VM details.
        """
        config = f"""
Host {self.hostname}
  HostName {self.hostname}
  User {self.username}
  StrictHostKeyChecking no
  IdentityFile {self.private_key_path}
  ForwardAgent yes
"""
        logger.info(f"Updating SSH config for host {self.hostname}.")
        with open(f'{os.path.expanduser("~")}/.ssh/config', "a") as file:
            file.write(config)

    def __establish_connection(self):
        """
        Establishes an SSH connection to the VM.
        """
        logger.info(f"Establishing SSH connection to {self.hostname}.")
        self.client.connect(
            self.hostname, username=self.username, key_filename=self.private_key_path
        )

    def __close_connection(self):
        """
        Closes the SSH connection.
        """
        if self.client:
            logger.info("Closing SSH connection.")
            self.client.close()

    def setup(self):
        """
        Sets up the SSH configuration and keys for the VM.
        """
        logger.info("Starting SSH setup.")
        public_key = self.__read_or_generate_public_key(self.email)
        self.__configure_vm(self.hostname, self.password, public_key)
        self.__update_ssh_config()
        self.__establish_connection()
        self.__close_connection()
        logger.info("SSH setup completed.")

__close_connection()

Closes the SSH connection.

Source code in vm_tool/ssh.py 
132
133
134
135
136
137
138
def __close_connection(self):
    """
    Closes the SSH connection.
    """
    if self.client:
        logger.info("Closing SSH connection.")
        self.client.close()

__configure_vm(vm_ip, vm_password, public_key)

Configures the VM by adding the public SSH key to the authorized keys.

Parameters:

Name Type Description Default
vm_ip str

The IP address of the VM.

 required
vm_password str

The password for the VM.

 required
public_key str

The public SSH key.

 required
Source code in vm_tool/ssh.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def __configure_vm(self, vm_ip, vm_password, public_key):
    """
    Configures the VM by adding the public SSH key to the authorized keys.

    Parameters:
        vm_ip (str): The IP address of the VM.
        vm_password (str): The password for the VM.
        public_key (str): The public SSH key.
    """
    logger.info(f"Configuring VM at {vm_ip} with provided public key.")
    self.client.connect(vm_ip, username=self.username, password=vm_password)
    self.client.exec_command(
        f'echo "{public_key}" >> ~/.ssh/authorized_keys'
    )  # nosec B601
    self.client.close()

__establish_connection()

Establishes an SSH connection to the VM.

Source code in vm_tool/ssh.py 
123
124
125
126
127
128
129
130
def __establish_connection(self):
    """
    Establishes an SSH connection to the VM.
    """
    logger.info(f"Establishing SSH connection to {self.hostname}.")
    self.client.connect(
        self.hostname, username=self.username, key_filename=self.private_key_path
    )

__generate_ssh_key(email)

Generates an SSH key pair.

Parameters:

Name Type Description Default
email str

The email for SSH key generation.

 required
Source code in vm_tool/ssh.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def __generate_ssh_key(self, email):
    """
    Generates an SSH key pair.

    Parameters:
        email (str): The email for SSH key generation.
    """

    logger.info(f"Generating SSH key for email: {email}")
    try:
        subprocess.run(
            [
                "ssh-keygen",
                "-t",
                "rsa",
                "-b",
                "4096",
                "-C",
                email,
                "-f",
                self.private_key_path,
                "-N",
                "",
            ],
            check=True,
            capture_output=True,
        )
        logger.info("SSH key generated successfully.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to generate SSH key: {e.stderr.decode()}")
        raise RuntimeError(f"Failed to generate SSH key: {e.stderr.decode()}")

__init__(hostname, username, password, email)

The constructor for SSHSetup class.

Parameters:

Name Type Description Default
hostname str

The hostname of the VM.

 required
username str

The username for SSH login.

 required
password str

The password for SSH login.

 required
email str

The email for SSH key generation.

 required
Source code in vm_tool/ssh.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def __init__(self, hostname, username, password, email):
    """
    The constructor for SSHSetup class.

    Parameters:
        hostname (str): The hostname of the VM.
        username (str): The username for SSH login.
        password (str): The password for SSH login.
        email (str): The email for SSH key generation.
    """
    self.hostname = hostname
    self.username = username
    self.password = password
    self.email = email
    self.private_key_path = os.path.expanduser("~/.ssh/id_rsa")
    self.client = paramiko.SSHClient()
    self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # nosec B507

__read_or_generate_public_key(email)

Reads the public SSH key or generates a new one if it doesn't exist.

Parameters:

Name Type Description Default
email str

The email for SSH key generation.

 required

Returns:

Name Type Description
str str

The public SSH key.
Source code in vm_tool/ssh.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def __read_or_generate_public_key(self, email) -> str:
    """
    Reads the public SSH key or generates a new one if it doesn't exist.

    Parameters:
        email (str): The email for SSH key generation.

    Returns:
        str: The public SSH key.
    """
    public_key_path = f"{self.private_key_path}.pub"
    if not os.path.exists(public_key_path):
        self.__generate_ssh_key(email)
    with open(public_key_path, "r") as file:
        logger.info(f"Reading public key from {public_key_path}")
        return file.read()

__update_ssh_config()

Updates the local SSH config file with the VM details.

Source code in vm_tool/ssh.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
    def __update_ssh_config(self):
        """
        Updates the local SSH config file with the VM details.
        """
        config = f"""
Host {self.hostname}
  HostName {self.hostname}
  User {self.username}
  StrictHostKeyChecking no
  IdentityFile {self.private_key_path}
  ForwardAgent yes
"""
        logger.info(f"Updating SSH config for host {self.hostname}.")
        with open(f'{os.path.expanduser("~")}/.ssh/config', "a") as file:
            file.write(config)

setup()

Sets up the SSH configuration and keys for the VM.

Source code in vm_tool/ssh.py 
140
141
142
143
144
145
146
147
148
149
150
def setup(self):
    """
    Sets up the SSH configuration and keys for the VM.
    """
    logger.info("Starting SSH setup.")
    public_key = self.__read_or_generate_public_key(self.email)
    self.__configure_vm(self.hostname, self.password, public_key)
    self.__update_ssh_config()
    self.__establish_connection()
    self.__close_connection()
    logger.info("SSH setup completed.")