3
\$\begingroup\$

I wrote a script to install public key of my system to remote host. I tested the script, and it is working fine for me. The script takes data like IP password and ssh port from CSV file. I do not want to use import os; only use CSV and paramiko imports. My aim is to copy public key to my remote host.

import csv
import os

from paramiko import SSHClient, AuthenticationException, AutoAddPolicy

username = "root"
key = open(os.path.expanduser('~/.ssh/id_rsa.pub')).read()

with open("hostinfo.csv") as h:
    host = csv.reader(h, delimiter=",")
    for i in host:
        # csv file field ip,pass,port
        HostIP, HostPass, HostPort = i[0], i[1], i[2]
        try:
            client = SSHClient()
            client.load_system_host_keys()  
            client.set_missing_host_key_policy(AutoAddPolicy)  
            client.connect(HostIP, HostPort, username, HostPass)
            print("Running commands on {0} port {1}".format(HostIP, HostPort))
            client.exec_command('mkdir -p ~/.ssh/')
            client.exec_command('echo "%s" >> ~/.ssh/authorized_keys' % key)
            client.exec_command('chmod 644 ~/.ssh/authorized_keys')
            client.exec_command('chmod 700 ~/.ssh/')
            print("Key is installed in {0}".format(HostIP))
        except AuthenticationException:
            raise AuthenticationException('Authentication failed: did you remember to create an SSH key?')
        finally:
            client.close()
            print("connection is closed")
\$\endgroup\$
1
  • 1
    \$\begingroup\$ Is this intended to be an alternative to ssh-copy-id? Perhaps explain what's different in your program. One obvious difference is that it only transfers RSA public keys and not any other. \$\endgroup\$ Commented Nov 21 at 9:52

3 Answers 3

4
\$\begingroup\$

Some suggestions to improve your code:

Unpacking

Instead of using indices to refer to the items in i, you can directly unpack them. This will make your code more concise and, in my opinion, readable:

for HostIP, HostPass, HostPort in host:  # <--
    # csv file field ip,pass,port
    try:
        client = SSHClient()
        client.load_system_host_keys()  
        ...

Use Functions

Even with such a simple script, packing chunks of code into their own function greatly helps with readability, for example:

for HostIP, HostPort, username, HostPass in host:
    try:
        connect_and_install_keys(HostIP, HostPort, username, HostPass)
    except:
        ...

def connect_and_install_keys(HostIP, HostPort, username, HostPass):
    try:
        client = connect_ssh(HostIP, HostPort, username, HostPass)
        install_keys(client)
    finally:
        client.close()

def connect_ssh(HostIP, HostPort, username, HostPass):
    client = SSHClient()
    client.load_system_host_keys()  
    client.set_missing_host_key_policy(AutoAddPolicy)  
    client.connect(HostIP, HostPort, username, HostPass)
    return client
    
def install_keys(client, ssh_dir="~/.ssh"):
    client.exec_command(f'mkdir -p {ssh_dir}')
    client.exec_command(f'echo "%s" >> {ssh_dir}/authorized_keys' % key)
    client.exec_command(f'chmod 644 {ssh_dir}/authorized_keys')
    client.exec_command(f'chmod 700 {ssh_dir}')

Spaces in File Names and Shell Injection

This is currently not a problem in your original code because you use hard-coded commands, but if you decide to construct command strings, you should also think about proper escaping command strings and shell injection vulnerabilities.

Wrap your logic with if __name__ == "__main__"

When someone imports your script, this code will get executed. If you define some functions, they may be useful for other scripts, too, so why not share them?

To avoid that your code runs when only importing functions, wrap your code logic in if __name__ == "__main__". Then, it will only run when this script is run as the main script:

if __name__ == "__main__":
    username = "root"
    key = open(os.path.expanduser('~/.ssh/id_rsa.pub')).read()

    with open("hostinfo.csv") as h:
        ...

Close files

key = open(os.path.expanduser('~/.ssh/id_rsa.pub')).read()

If there's no good reason to do otherwise, open should always be called in a with statement. Otherwise, the file may remain open and resources leak:

with open(os.path.expanduser('~/.ssh/id_rsa.pub')) as fh:
    key = fh.read()

Context Managers

Context managers are one of the bright sides of Python, in my opinion, a feature that I often miss in other languages. Among others, they allow you to clean up resources safely regardless of whether an error occurred. with open(...) as fh is such a context manager. If you don't know about context managers, I advise you to have a look. paramiko's SSHClient also can be used as a context manager, so it's very appropriate here:

def connect_and_install_keys(HostIP, HostPort, username, HostPass):
    with connect_ssh(HostIP, HostPort, username, HostPass) as client:
        install_keys(client)

So much easier, isn't it?

Record an Exception's Cause

In the body of your method, you catch an AuthenticationException, just to raise another with your own error message:

except AuthenticationException:
    raise AuthenticationException('Authentication failed: did you remember to create an SSH key?')

When doing this, I advise to record the original exception as cause of the exception you just thrown. The raise ... from ... construct of Python allows you to exactly that:

except AuthenticationException as err:
    raise AuthenticationException('Authentication failed: did you remember to create an SSH key?') from err

Then, the user will get to see both exceptions and no information is lost.

Argument Parser

Python makes it very easy to parse command-line arguments with its builtin argparse module. You could let the user input the username, making using the script a little easier:

from argparse import ArgumentParser

parser = ArgumentParser(description="Copy SSH keys")
parser.add_argument("hostfile", help="CSV file with hosts")  # positional argument
parser.add_argument("--username", default="root", help="User name of remote hosts")  # optional argument
\$\endgroup\$
4
\$\begingroup\$

Briefly, this -

        client.exec_command('mkdir -p ~/.ssh/')
        client.exec_command('echo "%s" >> ~/.ssh/authorized_keys' % key)
        client.exec_command('chmod 644 ~/.ssh/authorized_keys')
        client.exec_command('chmod 700 ~/.ssh/')

gives me security concerns. I'd be worried about some part of the pipeline logging or exposing the secret. For that reason, and also because it makes better use of the API, I strongly suggest performing direct file operations with paramiko.sftp instead. It supports all of these operations without needing a shell.

\$\endgroup\$
3
\$\begingroup\$

DRY

The following string is used in multiple places:

~/.ssh

You could store it in a variable named ssh_dir. That makes the code easier to modify. In fact, since you are already using os.path.expanduser, you may as well make use of it here too. Then you only do the expansion once.

Format

The following code:

print("Running commands on {0} port {1}".format(HostIP, HostPort))

can be simplified using an f-string:

print(f"Running commands on {HostIP} port {HostPort}")

Naming

The PEP 8 style guide recommends snake_case for variable names. Instead of:

HostIP, HostPass, HostPort

use:

host_ip, host_pass, host_port

The variable named h is too brief and conveys no meaning. I suggest a longer name:

with open("hostinfo.csv") as host_file_name:
    host = csv.reader(host_file_name, delimiter=",")

Documentation

Add a docstring to the top of the code to summarize its purpose. For example:

"""
Install public key of my system to remote host.
Take data like IP password and ssh port from CSV file.
"""
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.