From f7368c6421328bced59c3817d79c8a6acbee6c50 Mon Sep 17 00:00:00 2001 From: alfiechen Date: Thu, 23 Nov 2023 05:03:41 -0500 Subject: [openwrt_authentication] Fix some readability issues.. - Remove unused import - Add module level docstring. - Fix "Quote delimiter ' is inconsistent with the rest of the file" - Fix "Use lazy % formatting in logging functions" Bug: None Test: act.py Change-Id: I0a84b64cdcecda01be3901c2f08536c6ad94bf27 --- .../openwrt_lib/openwrt_authentication.py | 93 +++++++++++++--------- 1 file changed, 57 insertions(+), 36 deletions(-) diff --git a/acts/framework/acts/controllers/openwrt_lib/openwrt_authentication.py b/acts/framework/acts/controllers/openwrt_lib/openwrt_authentication.py index a300fff4d..5cab36f6c 100644 --- a/acts/framework/acts/controllers/openwrt_lib/openwrt_authentication.py +++ b/acts/framework/acts/controllers/openwrt_lib/openwrt_authentication.py @@ -1,20 +1,30 @@ +"""Module for OpenWrt SSH authentication. + +This module provides a class, OpenWrtAuth, for managing SSH authentication for +OpenWrt devices. It allows you to generate RSA key pairs, save them in a +specified directory, and upload the public key to a remote host. + +Usage: + 1. Create an instance of OpenWrtAuth with the required parameters. + 2. Call generate_rsa_key() to generate RSA key pairs and save them. + 3. Call send_public_key_to_remote_host() to upload the public key + to the remote host. +""" + import logging import os import paramiko import scp -import subprocess -_REMOTE_PATH = '/etc/dropbear/authorized_keys' +_REMOTE_PATH = "/etc/dropbear/authorized_keys" class OpenWrtAuth: - """ - A class for managing SSH authentication for OpenWrt devices. - """ - def __init__(self, hostname, username='root', password='root', port=22): - """ - Initializes a new instance of the OpenWrtAuth class. + """Class for managing SSH authentication for OpenWrt devices.""" + + def __init__(self, hostname, username="root", password="root", port=22): + """Initializes a new instance of the OpenWrtAuth class. Args: hostname (str): The hostname or IP address of the remote device. @@ -32,23 +42,30 @@ class OpenWrtAuth: self.password = password self.port = port self.public_key = None - self.key_dir = '/tmp/openwrt/' - self.public_key_file = f'{self.key_dir}id_rsa_{self.hostname}.pub' - self.private_key_file = f'{self.key_dir}id_rsa_{self.hostname}' + self.key_dir = "/tmp/openwrt/" + self.public_key_file = f"{self.key_dir}id_rsa_{self.hostname}.pub" + self.private_key_file = f"{self.key_dir}id_rsa_{self.hostname}" def generate_rsa_key(self): - """ - Generates an RSA key pair and saves it to the specified directory. + """Generates an RSA key pair and saves it to the specified directory. Raises: - ValueError: If an error occurs while generating the RSA key pair. - paramiko.SSHException: If an error occurs while generating the RSA key pair. - FileNotFoundError: If the directory for saving the private or public key does not exist. - PermissionError: If there is a permission error while creating the directory for saving the keys. - Exception: If an unexpected error occurs while generating the RSA key pair. + ValueError: + If an error occurs while generating the RSA key pair. + paramiko.SSHException: + If an error occurs while generating the RSA key pair. + FileNotFoundError: + If the directory for saving the private or public key does not exist. + PermissionError: + If there is a permission error while creating the directory + for saving the keys. + Exception: + If an unexpected error occurs while generating the RSA key pair. """ # Checks if the private and public key files already exist. - if os.path.exists(self.private_key_file) and os.path.exists(self.public_key_file): + if os.path.exists(self.private_key_file) and os.path.exists( + self.public_key_file + ): logging.warning("RSA key pair already exists, skipping key generation.") return @@ -57,38 +74,42 @@ class OpenWrtAuth: logging.info("Generating RSA key pair...") key = paramiko.RSAKey.generate(bits=2048) self.public_key = f"ssh-rsa {key.get_base64()}" - logging.debug(f"Public key: {self.public_key}") + logging.debug("Public key: %s", self.public_key) # Create /tmp/openwrt/ directory if it doesn't exist. - logging.info(f"Creating {self.key_dir} directory...") + logging.info("Creating %s directory...", self.key_dir) os.makedirs(self.key_dir, exist_ok=True) # Saves the private key to a file. key.write_private_key_file(self.private_key_file) - logging.debug(f"Saved private key to file: {self.private_key_file}") + logging.debug("Saved private key to file: %s", self.private_key_file) # Saves the public key to a file. with open(self.public_key_file, "w") as f: - f.write(self.public_key) - logging.debug(f"Saved public key to file: {self.public_key_file}") + f.write(self.public_key) + logging.debug("Saved public key to file: %s", self.public_key_file) except (ValueError, paramiko.SSHException, PermissionError) as e: - logging.error(f"An error occurred while generating the RSA key pair: {e}") + logging.error("An error occurred while generating " + "the RSA key pair: %s", e) except Exception as e: - logging.error(f"An unexpected error occurred while generating the RSA key pair: {e}") + logging.error("An unexpected error occurred while generating " + "the RSA key pair: %s", e) def send_public_key_to_remote_host(self): - """ - Uploads the public key to the remote host. + """Uploads the public key to the remote host. Raises: - paramiko.AuthenticationException: If authentication to the remote host fails. - paramiko.SSHException: If an SSH-related error occurs during the connection. - FileNotFoundError: If the public key file or the private key file does not exist. + paramiko.AuthenticationException: + If authentication to the remote host fails. + paramiko.SSHException: + If an SSH-related error occurs during the connection. + FileNotFoundError: + If the public key file or the private key file does not exist. Exception: If an unexpected error occurs while sending the public key. """ try: # Connects to the remote host and uploads the public key. - logging.info(f"Uploading public key to remote host {self.hostname}...") + logging.info("Uploading public key to remote host %s...", self.hostname) with paramiko.SSHClient() as ssh: ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) ssh.connect(hostname=self.hostname, @@ -97,11 +118,11 @@ class OpenWrtAuth: password=self.password) scp_client = scp.SCPClient(ssh.get_transport()) scp_client.put(self.public_key_file, _REMOTE_PATH) - logging.info('Public key uploaded successfully.') + logging.info("Public key uploaded successfully.") except (paramiko.AuthenticationException, paramiko.SSHException, FileNotFoundError) as e: - logging.error(f"An error occurred while sending the public key: {e}") + logging.error("An error occurred while sending the public key: %s", e) except Exception as e: - logging.error(f"An unexpected error occurred while " - f"sending the public key: {e}") + logging.error("An unexpected error occurred while " + "sending the public key: %s", e) -- cgit v1.2.3