1"""Module for OpenWrt SSH authentication.
2
3This module provides a class, OpenWrtAuth, for managing SSH authentication for
4OpenWrt devices. It allows you to generate RSA key pairs, save them in a
5specified directory, and upload the public key to a remote host.
6
7Usage:
8  1. Create an instance of OpenWrtAuth with the required parameters.
9  2. Call generate_rsa_key() to generate RSA key pairs and save them.
10  3. Call send_public_key_to_remote_host() to upload the public key
11     to the remote host.
12"""
13
14import logging
15import os
16import paramiko
17import scp
18
19
20_REMOTE_PATH = "/etc/dropbear/authorized_keys"
21
22
23class OpenWrtAuth:
24  """Class for managing SSH authentication for OpenWrt devices."""
25
26  def __init__(self, hostname, username="root", password="root", port=22):
27    """Initializes a new instance of the OpenWrtAuth class.
28
29    Args:
30      hostname (str): The hostname or IP address of the remote device.
31      username (str): The username for authentication.
32      password (str): The password for authentication.
33      port (int): The port number for SSH.
34
35    Attributes:
36      public_key (str): The generated public key.
37      public_key_file (str): The path to the generated public key file.
38      private_key_file (str): The path to the generated private key file.
39    """
40    self.hostname = hostname
41    self.username = username
42    self.password = password
43    self.port = port
44    self.public_key = None
45    self.key_dir = "/tmp/openwrt/"
46    self.public_key_file = f"{self.key_dir}id_rsa_{self.hostname}.pub"
47    self.private_key_file = f"{self.key_dir}id_rsa_{self.hostname}"
48
49  def generate_rsa_key(self):
50    """Generates an RSA key pair and saves it to the specified directory.
51
52    Raises:
53      ValueError:
54        If an error occurs while generating the RSA key pair.
55      paramiko.SSHException:
56        If an error occurs while generating the RSA key pair.
57      FileNotFoundError:
58        If the directory for saving the private or public key does not exist.
59      PermissionError:
60        If there is a permission error while creating the directory
61        for saving the keys.
62      Exception:
63        If an unexpected error occurs while generating the RSA key pair.
64    """
65    # Checks if the private and public key files already exist.
66    if os.path.exists(self.private_key_file) and os.path.exists(
67        self.public_key_file
68    ):
69      logging.warning("RSA key pair already exists, skipping key generation.")
70      return
71
72    try:
73      # Generates an RSA key pair in /tmp/openwrt/ directory.
74      logging.info("Generating RSA key pair...")
75      key = paramiko.RSAKey.generate(bits=2048)
76      self.public_key = f"ssh-rsa {key.get_base64()}"
77      logging.debug("Public key: %s", self.public_key)
78
79      # Create /tmp/openwrt/ directory if it doesn't exist.
80      logging.info("Creating %s directory...", self.key_dir)
81      os.makedirs(self.key_dir, exist_ok=True)
82
83      # Saves the private key to a file.
84      key.write_private_key_file(self.private_key_file)
85      logging.debug("Saved private key to file: %s", self.private_key_file)
86
87      # Saves the public key to a file.
88      with open(self.public_key_file, "w") as f:
89        f.write(self.public_key)
90      logging.debug("Saved public key to file: %s", self.public_key_file)
91    except (ValueError, paramiko.SSHException, PermissionError) as e:
92      logging.error("An error occurred while generating "
93                    "the RSA key pair: %s", e)
94    except Exception as e:
95      logging.error("An unexpected error occurred while generating "
96                    "the RSA key pair: %s", e)
97
98  def send_public_key_to_remote_host(self):
99    """Uploads the public key to the remote host.
100
101    Raises:
102      paramiko.AuthenticationException:
103        If authentication to the remote host fails.
104      paramiko.SSHException:
105        If an SSH-related error occurs during the connection.
106      FileNotFoundError:
107        If the public key file or the private key file does not exist.
108      Exception: If an unexpected error occurs while sending the public key.
109    """
110    try:
111      # Connects to the remote host and uploads the public key.
112      logging.info("Uploading public key to remote host %s...", self.hostname)
113      with paramiko.SSHClient() as ssh:
114        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
115        ssh.connect(hostname=self.hostname,
116                    port=self.port,
117                    username=self.username,
118                    password=self.password)
119        scp_client = scp.SCPClient(ssh.get_transport())
120        scp_client.put(self.public_key_file, _REMOTE_PATH)
121      logging.info("Public key uploaded successfully.")
122    except (paramiko.AuthenticationException,
123            paramiko.SSHException,
124            FileNotFoundError) as e:
125      logging.error("An error occurred while sending the public key: %s", e)
126    except Exception as e:
127      logging.error("An unexpected error occurred while "
128                    "sending the public key: %s", e)
129