tunnel-manager 0.0.5__py3-none-any.whl → 1.0.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of tunnel-manager might be problematic. Click here for more details.

@@ -1,73 +1,79 @@
1
+ #!/usr/bin/env python
2
+ # coding: utf-8
3
+
4
+ import sys
5
+ import argparse
6
+ import concurrent.futures
1
7
  import logging
2
8
  import os
3
9
  import paramiko
10
+ import yaml
4
11
 
5
12
 
6
13
  class Tunnel:
7
14
  def __init__(
8
15
  self,
9
16
  remote_host: str,
17
+ username: str = None,
18
+ password: str = None,
10
19
  port: int = 22,
11
20
  identity_file: str = None,
12
21
  certificate_file: str = None,
13
22
  proxy_command: str = None,
14
- log_file: str = None,
23
+ ssh_config_file: str = os.path.expanduser("~/.ssh/config"),
15
24
  ):
16
25
  """
17
26
  Initialize the Tunnel class.
18
27
 
19
28
  :param remote_host: The hostname or IP of the remote host.
29
+ :param username: The username for authentication (overrides config).
30
+ :param password: The password for authentication (if no identity_file).
31
+ :param port: The SSH port (default: 22).
20
32
  :param identity_file: Optional path to the private key file (overrides config).
21
- :param certificate_file: Optional path to the certificate file (overrides config, used for Teleport).
22
- :param proxy_command: Optional proxy command string (overrides config, used for Teleport proxying).
33
+ :param certificate_file: Optional path to the certificate file (overrides config).
34
+ :param proxy_command: Optional proxy command string (overrides config).
23
35
  :param log_file: Optional path to a log file for recording operations.
36
+ :param ssh_config_file: Optional path to a custom SSH config file (defaults to ~/.ssh/config).
24
37
  """
25
38
  self.remote_host = remote_host
39
+ self.username = username
40
+ self.password = password
26
41
  self.port = port
27
42
  self.ssh_client = None
28
43
  self.sftp = None
29
- self.logger = None
44
+ self.logger = logging.getLogger(__name__)
30
45
 
31
- # Load from ~/.ssh/config if not overridden
32
- ssh_config_path = os.path.expanduser("~/.ssh/config")
46
+ # Load SSH config from custom or default path
33
47
  self.ssh_config = paramiko.SSHConfig()
34
- if os.path.exists(ssh_config_path):
35
- with open(ssh_config_path) as f:
48
+ if os.path.exists(ssh_config_file) and os.path.isfile(ssh_config_file):
49
+ with open(ssh_config_file, "r") as f:
36
50
  self.ssh_config.parse(f)
51
+ self.logger.info(f"Loaded SSH config from: {ssh_config_file}")
52
+ else:
53
+ self.logger.warning(f"No SSH config found at: {ssh_config_file}")
37
54
  host_config = self.ssh_config.lookup(remote_host) or {}
38
55
 
56
+ self.username = username or host_config.get("user")
39
57
  self.identity_file = identity_file or (
40
- host_config.get("identityfile", [None])[0]
41
- if "identityfile" in host_config
58
+ host_config.get("identityfile")[0]
59
+ if host_config.get("identityfile")
42
60
  else None
43
61
  )
44
62
  self.certificate_file = certificate_file or host_config.get("certificatefile")
45
63
  self.proxy_command = proxy_command or host_config.get("proxycommand")
46
64
 
47
- if not self.identity_file:
48
- raise ValueError(
49
- "Identity file must be provided either via parameter or in ~/.ssh/config."
50
- )
51
-
52
- if log_file:
53
- logging.basicConfig(
54
- filename=log_file,
55
- level=logging.INFO,
56
- format="%(asctime)s - %(levelname)s - %(message)s",
57
- )
58
- self.logger = logging.getLogger(__name__)
59
- self.logger.info(f"Tunnel initialized for host: {remote_host}")
65
+ if not self.username:
66
+ raise ValueError("Username must be provided via parameter or SSH config.")
67
+ if not self.identity_file and not self.password:
68
+ raise ValueError("Either identity_file or password must be provided.")
60
69
 
61
70
  def connect(self):
62
- """
63
- Establish the SSH connection if not already connected.
64
- """
65
71
  if (
66
72
  self.ssh_client
67
73
  and self.ssh_client.get_transport()
68
74
  and self.ssh_client.get_transport().is_active()
69
75
  ):
70
- return # Already connected
76
+ return
71
77
 
72
78
  self.ssh_client = paramiko.SSHClient()
73
79
  self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
@@ -75,30 +81,49 @@ class Tunnel:
75
81
  proxy = None
76
82
  if self.proxy_command:
77
83
  proxy = paramiko.ProxyCommand(self.proxy_command)
78
- if self.logger:
79
- self.logger.info(f"Using proxy command: {self.proxy_command}")
80
-
81
- private_key = paramiko.RSAKey.from_private_key_file(self.identity_file)
82
- if self.certificate_file:
83
- private_key.load_certificate(self.certificate_file)
84
- if self.logger:
85
- self.logger.info(f"Loaded certificate: {self.certificate_file}")
84
+ self.logger.info(f"Using proxy command: {self.proxy_command}")
86
85
 
87
86
  try:
88
- self.ssh_client.connect(
89
- self.remote_host,
90
- port=self.port,
91
- pkey=private_key,
92
- sock=proxy,
93
- auth_timeout=30,
94
- look_for_keys=False,
95
- allow_agent=False,
96
- )
97
- if self.logger:
98
- self.logger.info(f"Connected to {self.remote_host}")
87
+ if self.identity_file:
88
+ # Try loading as ED25519 key first
89
+ try:
90
+ private_key = paramiko.Ed25519Key.from_private_key_file(
91
+ self.identity_file
92
+ )
93
+ self.logger.info(f"Loaded ED25519 key from: {self.identity_file}")
94
+ except paramiko.ssh_exception.SSHException:
95
+ # Fallback to RSA key
96
+ private_key = paramiko.RSAKey.from_private_key_file(
97
+ self.identity_file
98
+ )
99
+ self.logger.info(f"Loaded RSA key from: {self.identity_file}")
100
+ if self.certificate_file:
101
+ private_key.load_certificate(self.certificate_file)
102
+ self.logger.info(f"Loaded certificate: {self.certificate_file}")
103
+ self.ssh_client.connect(
104
+ self.remote_host,
105
+ port=self.port,
106
+ username=self.username,
107
+ pkey=private_key,
108
+ sock=proxy,
109
+ auth_timeout=30,
110
+ look_for_keys=False,
111
+ allow_agent=False,
112
+ )
113
+ else:
114
+ self.ssh_client.connect(
115
+ self.remote_host,
116
+ port=self.port,
117
+ username=self.username,
118
+ password=self.password,
119
+ sock=proxy,
120
+ auth_timeout=30,
121
+ look_for_keys=False,
122
+ allow_agent=False,
123
+ )
124
+ self.logger.info(f"Connected to {self.remote_host}")
99
125
  except Exception as e:
100
- if self.logger:
101
- self.logger.error(f"Connection failed: {str(e)}")
126
+ self.logger.error(f"Connection failed: {str(e)}")
102
127
  raise
103
128
 
104
129
  def run_command(self, command):
@@ -113,33 +138,71 @@ class Tunnel:
113
138
  stdin, stdout, stderr = self.ssh_client.exec_command(command)
114
139
  out = stdout.read().decode("utf-8").strip()
115
140
  err = stderr.read().decode("utf-8").strip()
116
- if self.logger:
117
- self.logger.info(
118
- f"Command executed: {command}\nOutput: {out}\nError: {err}"
119
- )
141
+ self.logger.info(
142
+ f"Command executed: {command}\nOutput: {out}\nError: {err}"
143
+ )
120
144
  return out, err
121
145
  except Exception as e:
122
- if self.logger:
123
- self.logger.error(f"Command execution failed: {str(e)}")
146
+ self.logger.error(f"Command execution failed: {str(e)}")
124
147
  raise
125
148
 
126
149
  def send_file(self, local_path, remote_path):
127
150
  """
128
151
  Send (upload) a file to the remote host.
129
-
130
152
  :param local_path: Path to the local file.
131
153
  :param remote_path: Path on the remote host.
132
154
  """
133
155
  self.connect()
134
156
  try:
157
+ # Normalize paths for consistency
158
+ local_path = os.path.abspath(os.path.expanduser(local_path))
159
+ remote_path = os.path.expanduser(
160
+ remote_path
161
+ ) # ~ expansion for remote, but paramiko handles it
162
+
163
+ self.logger.debug(
164
+ f"send_file: local_path='{local_path}', remote_path='{remote_path}'"
165
+ )
166
+ self.logger.debug(f"send_file: CWD={os.getcwd()}")
167
+
168
+ # Explicit checks before SFTP
169
+ if not os.path.exists(local_path):
170
+ err_msg = f"Local file does not exist: {local_path}"
171
+ self.logger.error(err_msg)
172
+ raise IOError(err_msg)
173
+ if not os.path.isfile(local_path):
174
+ err_msg = (
175
+ f"Local path is not a regular file (dir/symlink?): {local_path}"
176
+ )
177
+ self.logger.error(err_msg)
178
+ raise IOError(err_msg)
179
+ if not os.access(local_path, os.R_OK):
180
+ err_msg = f"No read permission for local file: {local_path}"
181
+ self.logger.error(err_msg)
182
+ raise PermissionError(err_msg)
183
+
184
+ # Test binary open (mimics what sftp.put does)
185
+ try:
186
+ with open(local_path, "rb") as f:
187
+ sample = f.read(1024) # Read a chunk to simulate transfer
188
+ self.logger.debug(
189
+ f"Binary open successful for {local_path}, sample size: {len(sample)} bytes"
190
+ )
191
+ except Exception as open_err:
192
+ err_msg = f"Failed to open {local_path} in binary mode: {str(open_err)}"
193
+ self.logger.error(err_msg)
194
+ raise IOError(err_msg)
195
+
135
196
  if not self.sftp:
136
197
  self.sftp = self.ssh_client.open_sftp()
198
+ self.logger.debug(f"Opening SFTP for put: {local_path} -> {remote_path}")
137
199
  self.sftp.put(local_path, remote_path)
138
- if self.logger:
139
- self.logger.info(f"File sent: {local_path} -> {remote_path}")
200
+ self.logger.info(f"File sent: {local_path} -> {remote_path}")
140
201
  except Exception as e:
141
- if self.logger:
142
- self.logger.error(f"File send failed: {str(e)}")
202
+ self.logger.error(f"File send failed: {str(e)} (type: {type(e).__name__})")
203
+ import traceback
204
+
205
+ self.logger.error(traceback.format_exc())
143
206
  raise
144
207
  finally:
145
208
  if self.sftp:
@@ -158,33 +221,770 @@ class Tunnel:
158
221
  if not self.sftp:
159
222
  self.sftp = self.ssh_client.open_sftp()
160
223
  self.sftp.get(remote_path, local_path)
161
- if self.logger:
162
- self.logger.info(f"File received: {remote_path} -> {local_path}")
224
+ self.logger.info(f"File received: {remote_path} -> {local_path}")
163
225
  except Exception as e:
164
- if self.logger:
165
- self.logger.error(f"File receive failed: {str(e)}")
226
+ self.logger.error(f"File receive failed: {str(e)}")
166
227
  raise
167
228
  finally:
168
229
  if self.sftp:
169
230
  self.sftp.close()
170
231
  self.sftp = None
171
232
 
233
+ def check_ssh_server(self):
234
+ """
235
+ Check if the SSH server is running and configured for key-based auth on the remote host.
236
+ :return: Tuple (bool, str) indicating if SSH server is running and any error message.
237
+ """
238
+ try:
239
+ self.connect()
240
+ out, err = self.run_command(
241
+ "systemctl status sshd || ps aux | grep '[s]shd'"
242
+ )
243
+ if "running" in out.lower() or "sshd" in out.lower():
244
+ out, err = self.run_command(
245
+ "grep '^PubkeyAuthentication' /etc/ssh/sshd_config"
246
+ )
247
+ if "PubkeyAuthentication yes" in out:
248
+ return True, "SSH server running with key-based auth enabled."
249
+ return False, "SSH server running but key-based auth not enabled."
250
+ return False, "SSH server not running."
251
+ except Exception as e:
252
+ self.logger.error(f"Failed to check SSH server: {str(e)}")
253
+ return False, f"Failed to check SSH server: {str(e)}"
254
+ finally:
255
+ self.close()
256
+
257
+ def test_key_auth(self, local_key_path):
258
+ """
259
+ Test if key-based authentication works for the remote host.
260
+ :param local_key_path: Path to the private key to test.
261
+ :return: Tuple (bool, str) indicating success and any error message.
262
+ """
263
+ local_key_path = os.path.expanduser(local_key_path)
264
+ try:
265
+ temp_tunnel = Tunnel(
266
+ remote_host=self.remote_host,
267
+ username=self.username,
268
+ identity_file=local_key_path,
269
+ )
270
+ temp_tunnel.connect()
271
+ temp_tunnel.close()
272
+ return True, "Key-based authentication successful."
273
+ except Exception as e:
274
+ self.logger.error(f"Key auth test failed: {str(e)}")
275
+ return False, f"Key auth test failed: {str(e)}"
276
+
172
277
  def close(self):
173
278
  """
174
279
  Close the SSH connection.
175
280
  """
176
281
  if self.ssh_client:
177
282
  self.ssh_client.close()
178
- if self.logger:
179
- self.logger.info(f"Connection closed for {self.remote_host}")
283
+ self.logger.info(f"Connection closed for {self.remote_host}")
180
284
  self.ssh_client = None
181
285
 
286
+ def setup_passwordless_ssh(
287
+ self, local_key_path=os.path.expanduser("~/.ssh/id_rsa"), key_type="ed25519"
288
+ ):
289
+ """
290
+ Set up passwordless SSH by copying a public key to the remote host.
291
+ Requires password-based authentication to be configured.
292
+
293
+ :param local_key_path: Path to the local private key (public key is assumed to be .pub).
294
+ :param key_type: Type of key to generate ('rsa' or 'ed25519', default: 'rsa').
295
+ """
296
+ if not self.password:
297
+ raise ValueError("Password-based authentication required for setup.")
298
+
299
+ local_key_path = os.path.expanduser(local_key_path)
300
+ pub_key_path = local_key_path + ".pub"
301
+
302
+ if key_type not in ["rsa", "ed25519"]:
303
+ raise ValueError("key_type must be 'rsa' or 'ed25519'")
304
+
305
+ if not os.path.exists(pub_key_path):
306
+ if key_type == "rsa":
307
+ os.system(f"ssh-keygen -t rsa -b 4096 -f {local_key_path} -N ''")
308
+ else: # ed25519
309
+ os.system(f"ssh-keygen -t ed25519 -f {local_key_path} -N ''")
310
+ self.logger.info(
311
+ f"Generated {key_type} key pair: {local_key_path}, {pub_key_path}"
312
+ )
313
+
314
+ with open(pub_key_path, "r") as f:
315
+ pub_key = f.read().strip()
316
+
317
+ try:
318
+ self.connect()
319
+ self.run_command("mkdir -p ~/.ssh && chmod 700 ~/.ssh")
320
+ self.run_command(f"echo '{pub_key}' >> ~/.ssh/authorized_keys")
321
+ self.run_command("chmod 600 ~/.ssh/authorized_keys")
322
+ self.logger.info(
323
+ f"Set up passwordless SSH for {self.username}@{self.remote_host} with {key_type} key"
324
+ )
325
+ except Exception as e:
326
+ self.logger.error(f"Failed to set up passwordless SSH: {str(e)}")
327
+ raise
328
+ finally:
329
+ self.close()
330
+
331
+ @staticmethod
332
+ def execute_on_inventory(
333
+ inventory, func, group="all", parallel=False, max_threads=5
334
+ ):
335
+ """
336
+ Execute a function on all hosts in the specified group of the YAML inventory, sequentially or in parallel.
337
+ :param inventory: Path to the YAML inventory file.
338
+ :param func: Function to execute, takes host dict as argument.
339
+ :param group: Inventory group to target (default: 'all').
340
+ :param parallel: Whether to run in parallel using threads.
341
+ :param max_threads: Maximum number of threads if parallel.
342
+ """
343
+ logger = logging.getLogger("Tunnel")
344
+ logger.info(f"Processing inventory '{inventory}' for group '{group}'")
345
+ print(f"Loading inventory '{inventory}' for group '{group}'...")
346
+
347
+ try:
348
+ with open(inventory, "r") as f:
349
+ inventory_data = yaml.safe_load(f)
350
+ logger.debug(f"Loaded inventory data: {inventory_data}")
351
+ except FileNotFoundError:
352
+ logger.error(f"Inventory file not found: {inventory}")
353
+ print(f"Error: Inventory file not found: {inventory}", file=sys.stderr)
354
+ raise
355
+ except yaml.YAMLError as e:
356
+ logger.error(f"Failed to parse inventory file: {str(e)}")
357
+ print(f"Error: Failed to parse inventory file: {str(e)}", file=sys.stderr)
358
+ raise
359
+
360
+ hosts = []
361
+ if (
362
+ group in inventory_data
363
+ and isinstance(inventory_data[group], dict)
364
+ and "hosts" in inventory_data[group]
365
+ and isinstance(inventory_data[group]["hosts"], dict)
366
+ ):
367
+ for host, vars in inventory_data[group]["hosts"].items():
368
+ host_entry = {
369
+ "hostname": vars.get("ansible_host", host),
370
+ "username": vars.get("ansible_user"),
371
+ "password": vars.get("ansible_ssh_pass"),
372
+ "key_path": vars.get("ansible_ssh_private_key_file"),
373
+ }
374
+ if not host_entry["username"]:
375
+ logger.error(
376
+ f"No username specified for host {host_entry['hostname']}"
377
+ )
378
+ print(
379
+ f"Error: No username specified for host {host_entry['hostname']}",
380
+ file=sys.stderr,
381
+ )
382
+ continue
383
+ logger.debug(f"Added host: {host_entry['hostname']}")
384
+ hosts.append(host_entry)
385
+ else:
386
+ logger.error(
387
+ f"Group '{group}' not found in inventory or invalid (hosts not a dict)"
388
+ )
389
+ print(
390
+ f"Error: Group '{group}' not found in inventory or invalid (hosts not a dict)",
391
+ file=sys.stderr,
392
+ )
393
+ raise ValueError(f"Group '{group}' not found in inventory or invalid")
394
+
395
+ logger.info(f"Found {len(hosts)} hosts in group '{group}'")
396
+ print(f"Found {len(hosts)} hosts in group '{group}'")
397
+
398
+ if not hosts:
399
+ logger.warning(f"No valid hosts found in group '{group}'")
400
+ print(f"Warning: No valid hosts found in group '{group}'")
401
+ return
402
+
403
+ if parallel:
404
+ with concurrent.futures.ThreadPoolExecutor(
405
+ max_workers=max_threads
406
+ ) as executor:
407
+ futures = [executor.submit(func, host) for host in hosts]
408
+ for future in concurrent.futures.as_completed(futures):
409
+ try:
410
+ future.result()
411
+ except Exception as e:
412
+ logger.error(f"Error in parallel execution: {str(e)}")
413
+ print(f"Error in parallel execution: {str(e)}", file=sys.stderr)
414
+ else:
415
+ for host in hosts:
416
+ func(host)
417
+ print(f"Completed processing group '{group}'")
418
+
419
+ def remove_host_key(
420
+ self, known_hosts_path=os.path.expanduser("~/.ssh/known_hosts")
421
+ ) -> str:
422
+ """
423
+ Remove the host key for the remote host from the known_hosts file.
424
+ :param known_hosts_path: Path to the known_hosts file (default: ~/.ssh/known_hosts).
425
+ """
426
+ known_hosts_path = os.path.expanduser(known_hosts_path)
427
+ kh = paramiko.HostKeys()
428
+ if os.path.exists(known_hosts_path):
429
+ kh.load(known_hosts_path)
430
+ if self.remote_host in kh:
431
+ del kh[self.remote_host]
432
+ kh.save(known_hosts_path)
433
+ self.logger.info(
434
+ f"Removed host key for {self.remote_host} from {known_hosts_path}"
435
+ )
436
+ return (
437
+ f"Removed host key for {self.remote_host} from {known_hosts_path}"
438
+ )
439
+ else:
440
+ self.logger.warning(
441
+ f"No host key found for {self.remote_host} in {known_hosts_path}"
442
+ )
443
+ return f"No host key found for {self.remote_host} in {known_hosts_path}"
444
+ else:
445
+ self.logger.warning(f"No known_hosts file at {known_hosts_path}")
446
+ return f"No known_hosts file at {known_hosts_path}"
447
+
448
+ def copy_ssh_config(
449
+ self, local_config_path, remote_config_path=os.path.expanduser("~/.ssh/config")
450
+ ):
451
+ """
452
+ Copy a local SSH config to the remote host’s ~/.ssh/config.
453
+ :param local_config_path: Path to the local config file.
454
+ :param remote_config_path: Path on remote (default ~/.ssh/config).
455
+ """
456
+ self.connect()
457
+ self.run_command("mkdir -p ~/.ssh && chmod 700 ~/.ssh")
458
+ self.send_file(local_config_path, remote_config_path)
459
+ self.run_command(f"chmod 600 {remote_config_path}")
460
+ self.logger.info(
461
+ f"Copied SSH config to {remote_config_path} on {self.remote_host}"
462
+ )
463
+
464
+ def rotate_ssh_key(self, new_key_path, key_type="ed25519"):
465
+ """
466
+ Rotate the SSH key by generating a new pair and updating authorized_keys.
467
+ :param new_key_path: Path for the new private key.
468
+ :param key_type: Type of key to generate ('rsa' or 'ed25519', default: 'rsa').
469
+ """
470
+ new_key_path = os.path.expanduser(new_key_path)
471
+ new_pub_path = new_key_path + ".pub"
472
+ if key_type not in ["rsa", "ed25519"]:
473
+ raise ValueError("key_type must be 'rsa' or 'ed25519'")
474
+
475
+ if not os.path.exists(new_key_path):
476
+ if key_type == "rsa":
477
+ os.system(f"ssh-keygen -t rsa -b 4096 -f {new_key_path} -N ''")
478
+ else: # ed25519
479
+ os.system(f"ssh-keygen -t ed25519 -f {new_key_path} -N ''")
480
+ self.logger.info(f"Generated new {key_type} key pair: {new_key_path}")
481
+
482
+ with open(new_pub_path, "r") as f:
483
+ new_pub = f.read().strip()
484
+
485
+ old_pub = None
486
+ if self.identity_file:
487
+ old_key_path = os.path.expanduser(self.identity_file)
488
+ old_pub_path = old_key_path + ".pub"
489
+ if os.path.exists(old_pub_path):
490
+ with open(old_pub_path, "r") as f:
491
+ old_pub = f.read().strip()
492
+
493
+ self.connect()
494
+ out, err = self.run_command("cat ~/.ssh/authorized_keys")
495
+ auth_keys = out.splitlines()
496
+ new_auth = [
497
+ line
498
+ for line in auth_keys
499
+ if line.strip() and (old_pub is None or line.strip() != old_pub)
500
+ ]
501
+ new_auth.append(new_pub)
502
+
503
+ temp_file = "/tmp/authorized_keys.new"
504
+ new_auth_joined = "\n".join(new_auth)
505
+ self.run_command(f"echo '{new_auth_joined}' > {temp_file}")
506
+ self.run_command(f"mv {temp_file} ~/.ssh/authorized_keys")
507
+ self.run_command("chmod 600 ~/.ssh/authorized_keys")
508
+
509
+ self.identity_file = new_key_path
510
+ self.password = None
511
+ self.logger.info(
512
+ f"Rotated {key_type} key to {new_key_path} on {self.remote_host}"
513
+ )
514
+ logging.info(
515
+ f"Please update SSH config for {self.remote_host} IdentityFile to {new_key_path}"
516
+ )
517
+
518
+ @staticmethod
519
+ def setup_all_passwordless_ssh(
520
+ inventory,
521
+ shared_key_path=os.path.expanduser("~/.ssh/id_shared"),
522
+ key_type="ed25519",
523
+ group="all",
524
+ parallel=False,
525
+ max_threads=5,
526
+ ):
527
+ """
528
+ Set up passwordless SSH for all hosts in the specified group of the YAML inventory.
529
+ :param inventory: Path to the YAML inventory file.
530
+ :param shared_key_path: Path to a shared private key (optional, generates if missing).
531
+ :param key_type: Type of key to generate ('rsa' or 'ed25519', default: 'rsa').
532
+ :param group: Inventory group to target (default: 'all').
533
+ :param parallel: Run in parallel.
534
+ :param max_threads: Max threads for parallel.
535
+ """
536
+ shared_key_path = os.path.expanduser(shared_key_path)
537
+ shared_pub_key_path = shared_key_path + ".pub"
538
+ if key_type not in ["rsa", "ed25519"]:
539
+ raise ValueError("key_type must be 'rsa' or 'ed25519'")
540
+
541
+ if not os.path.exists(shared_key_path):
542
+ if key_type == "rsa":
543
+ os.system(f"ssh-keygen -t rsa -b 4096 -f {shared_key_path} -N ''")
544
+ else: # ed25519
545
+ os.system(f"ssh-keygen -t ed25519 -f {shared_key_path} -N ''")
546
+ logging.info(
547
+ f"Generated shared {key_type} key pair: {shared_key_path}, {shared_pub_key_path}"
548
+ )
549
+
550
+ with open(shared_pub_key_path, "r") as f:
551
+ shared_pub_key = f.read().strip()
552
+
553
+ def setup_host(host):
554
+ hostname = host["hostname"]
555
+ username = host["username"]
556
+ password = host["password"]
557
+ key_path = host.get("key_path", shared_key_path)
558
+
559
+ logging.info(f"\nSetting up {username}@{hostname}...")
560
+
561
+ tunnel = Tunnel(
562
+ remote_host=hostname,
563
+ username=username,
564
+ password=password,
565
+ )
566
+ tunnel.remove_host_key()
567
+ tunnel.setup_passwordless_ssh(local_key_path=key_path, key_type=key_type)
568
+
569
+ try:
570
+ tunnel.connect()
571
+ tunnel.run_command(f"echo '{shared_pub_key}' >> ~/.ssh/authorized_keys")
572
+ tunnel.run_command("chmod 600 ~/.ssh/authorized_keys")
573
+ logging.info(f"Added shared {key_type} key to {username}@{hostname}")
574
+ except Exception as e:
575
+ logging.error(
576
+ f"Failed to add shared key to {username}@{hostname}: {str(e)}"
577
+ )
578
+ finally:
579
+ tunnel.close()
580
+
581
+ result, msg = tunnel.test_key_auth(key_path)
582
+ logging.info(f"Key auth test for {username}@{hostname}: {msg}")
583
+
584
+ Tunnel.execute_on_inventory(inventory, setup_host, group, parallel, max_threads)
585
+
586
+ @staticmethod
587
+ def run_command_on_inventory(
588
+ inventory, command, group="all", parallel=False, max_threads=5
589
+ ):
590
+ """
591
+ Run a shell command on all hosts in the specified group of the YAML inventory.
592
+ :param inventory: Path to the YAML inventory file.
593
+ :param command: The shell command to run.
594
+ :param group: Inventory group to target (default: 'all').
595
+ :param parallel: Run in parallel.
596
+ :param max_threads: Max threads for parallel.
597
+ """
598
+ logger = logging.getLogger("Tunnel")
599
+ logger.info(f"Running command '{command}' on group '{group}'")
600
+ print(f"Executing command '{command}' on group '{group}'...")
601
+
602
+ def run_host(host):
603
+ try:
604
+ tunnel = Tunnel(
605
+ remote_host=host["hostname"],
606
+ username=host["username"],
607
+ password=host.get("password"),
608
+ identity_file=host.get("key_path"),
609
+ )
610
+ out, err = tunnel.run_command(command)
611
+ logger.info(
612
+ f"Host {host['hostname']}: In: {command}, Out: {out}, Err: {err}"
613
+ )
614
+ print(
615
+ f"Host {host['hostname']}:\nInput: {command}\nOutput: {out}\nError: {err}"
616
+ )
617
+ tunnel.close()
618
+ except Exception as e:
619
+ logger.error(f"Failed to run command on {host['hostname']}: {str(e)}")
620
+ print(f"Error on {host['hostname']}: {str(e)}", file=sys.stderr)
621
+
622
+ try:
623
+ Tunnel.execute_on_inventory(
624
+ inventory, run_host, group, parallel, max_threads
625
+ )
626
+ print(f"Completed command execution on group '{group}'")
627
+ except Exception as e:
628
+ logger.error(f"Failed to execute command on group '{group}': {str(e)}")
629
+ print(
630
+ f"Error executing command on group '{group}': {str(e)}", file=sys.stderr
631
+ )
632
+ raise
633
+
634
+ @staticmethod
635
+ def copy_ssh_config_on_inventory(
636
+ inventory,
637
+ local_config_path,
638
+ remote_config_path=os.path.expanduser("~/.ssh/config"),
639
+ group="all",
640
+ parallel=False,
641
+ max_threads=5,
642
+ ):
643
+ """
644
+ Copy local SSH config to all hosts in the specified group of the YAML inventory.
645
+ :param inventory: Path to the YAML inventory file.
646
+ :param local_config_path: Local SSH config path.
647
+ :param remote_config_path: Remote path (default ~/.ssh/config).
648
+ :param group: Inventory group to target (default: 'all').
649
+ :param parallel: Run in parallel.
650
+ :param max_threads: Max threads for parallel.
651
+ """
652
+
653
+ def copy_host(host):
654
+ tunnel = Tunnel(
655
+ remote_host=host["hostname"],
656
+ username=host["username"],
657
+ password=host.get("password"),
658
+ identity_file=host.get("key_path"),
659
+ )
660
+ tunnel.copy_ssh_config(local_config_path, remote_config_path)
661
+ tunnel.close()
662
+
663
+ Tunnel.execute_on_inventory(inventory, copy_host, group, parallel, max_threads)
664
+
665
+ @staticmethod
666
+ def rotate_ssh_key_on_inventory(
667
+ inventory,
668
+ key_prefix=os.path.expanduser("~/.ssh/id_"),
669
+ key_type="ed25519",
670
+ group="all",
671
+ parallel=False,
672
+ max_threads=5,
673
+ ):
674
+ """
675
+ Rotate SSH keys for all hosts in the specified group of the YAML inventory.
676
+ :param inventory: Path to the YAML inventory file.
677
+ :param key_prefix: Prefix for new key paths (appends hostname).
678
+ :param key_type: Type of key to generate ('rsa' or 'ed25519', default: 'rsa').
679
+ :param group: Inventory group to target (default: 'all').
680
+ :param parallel: Run in parallel.
681
+ :param max_threads: Max threads for parallel.
682
+ """
683
+
684
+ def rotate_host(host):
685
+ new_key_path = os.path.expanduser(key_prefix + host["hostname"])
686
+ tunnel = Tunnel(
687
+ remote_host=host["hostname"],
688
+ username=host["username"],
689
+ password=host.get("password"),
690
+ identity_file=host.get("key_path"),
691
+ )
692
+ tunnel.rotate_ssh_key(new_key_path, key_type=key_type)
693
+ logging.info(
694
+ f"Rotated {key_type} key for {host['hostname']}. Update inventory key_path to {new_key_path} if needed."
695
+ )
696
+ tunnel.close()
697
+
698
+ Tunnel.execute_on_inventory(
699
+ inventory, rotate_host, group, parallel, max_threads
700
+ )
701
+
702
+ @staticmethod
703
+ def send_file_on_inventory(
704
+ inventory,
705
+ local_path,
706
+ remote_path,
707
+ group="all",
708
+ parallel=False,
709
+ max_threads=5,
710
+ ):
711
+ """
712
+ Upload a file to all hosts in the specified group of the YAML inventory.
713
+ :param inventory: Path to the YAML inventory file.
714
+ :param local_path: Path to the local file to upload.
715
+ :param remote_path: Path on the remote hosts to save the file.
716
+ :param group: Inventory group to target (default: 'all').
717
+ :param parallel: Run in parallel.
718
+ :param max_threads: Max threads for parallel execution.
719
+ """
720
+
721
+ def send_host(host):
722
+ tunnel = Tunnel(
723
+ remote_host=host["hostname"],
724
+ username=host["username"],
725
+ password=host.get("password"),
726
+ identity_file=host.get("key_path"),
727
+ )
728
+ tunnel.send_file(local_path, remote_path)
729
+ logging.info(f"Host {host['hostname']}: File uploaded to {remote_path}")
730
+ tunnel.close()
731
+
732
+ if not os.path.exists(local_path):
733
+ raise ValueError(f"Local file does not exist: {local_path}")
734
+
735
+ Tunnel.execute_on_inventory(inventory, send_host, group, parallel, max_threads)
736
+
737
+ @staticmethod
738
+ def receive_file_on_inventory(
739
+ inventory,
740
+ remote_path: str,
741
+ local_path_prefix,
742
+ group="all",
743
+ parallel=False,
744
+ max_threads=5,
745
+ ):
746
+ """
747
+ Download a file from all hosts in the specified group of the YAML inventory.
748
+ :param inventory: Path to the YAML inventory file.
749
+ :param remote_path: Path on the remote hosts to download the file from.
750
+ :param local_path_prefix: Local directory path prefix to save files (creates host-specific subdirectories).
751
+ :param group: Inventory group to target (default: 'all').
752
+ :param parallel: Run in parallel.
753
+ :param max_threads: Max threads for parallel execution.
754
+ """
755
+
756
+ def receive_host(host):
757
+ host_dir = os.path.join(local_path_prefix, host["hostname"])
758
+ os.makedirs(host_dir, exist_ok=True)
759
+ local_path = os.path.join(f"{host_dir}", os.path.basename(remote_path))
760
+ tunnel = Tunnel(
761
+ remote_host=host["hostname"],
762
+ username=host["username"],
763
+ password=host.get("password"),
764
+ identity_file=host.get("key_path"),
765
+ )
766
+ tunnel.receive_file(remote_path, local_path)
767
+ logging.info(f"Host {host['hostname']}: File downloaded to {local_path}")
768
+ tunnel.close()
769
+
770
+ os.makedirs(local_path_prefix, exist_ok=True)
771
+ Tunnel.execute_on_inventory(
772
+ inventory, receive_host, group, parallel, max_threads
773
+ )
774
+
775
+
776
+ def tunnel_manager():
777
+ parser = argparse.ArgumentParser(description="Tunnel Manager CLI")
778
+ parser.add_argument("--log-file", help="Log to this file (default: console output)")
779
+
780
+ subparsers = parser.add_subparsers(dest="command", required=True)
781
+
782
+ # Setup-all command
783
+ setup_parser = subparsers.add_parser("setup-all", help="Setup passwordless for all")
784
+ setup_parser.add_argument("--inventory", help="YAML inventory path")
785
+ setup_parser.add_argument(
786
+ "--shared-key-path",
787
+ default="~/.ssh/id_shared",
788
+ help="Path to shared private key",
789
+ )
790
+ setup_parser.add_argument(
791
+ "--key-type",
792
+ choices=["rsa", "ed25519"],
793
+ default="ed25519",
794
+ help="Key type to generate (rsa or ed25519, default: ed25519)",
795
+ )
796
+ setup_parser.add_argument(
797
+ "--group", default="all", help="Inventory group to target (default: all)"
798
+ )
799
+ setup_parser.add_argument("--parallel", action="store_true", help="Run in parallel")
800
+ setup_parser.add_argument(
801
+ "--max-threads", type=int, default=5, help="Max threads for parallel execution"
802
+ )
803
+
804
+ # Run-command command
805
+ run_parser = subparsers.add_parser("run-command", help="Run command on all")
806
+ run_parser.add_argument("--inventory", help="YAML inventory path")
807
+ run_parser.add_argument("--remote-command", help="Shell command to run")
808
+ run_parser.add_argument(
809
+ "--group", default="all", help="Inventory group to target (default: all)"
810
+ )
811
+ run_parser.add_argument("--parallel", action="store_true", help="Run in parallel")
812
+ run_parser.add_argument(
813
+ "--max-threads", type=int, default=5, help="Max threads for parallel execution"
814
+ )
815
+
816
+ # Copy-config command
817
+ copy_parser = subparsers.add_parser("copy-config", help="Copy SSH config to all")
818
+ copy_parser.add_argument("--inventory", help="YAML inventory path")
819
+ copy_parser.add_argument(
820
+ "--local-config-path", default="~/.ssh/config", help="Local SSH config path"
821
+ )
822
+ copy_parser.add_argument(
823
+ "--remote-config-path",
824
+ default="~/.ssh/config",
825
+ help="Remote path (default ~/.ssh/config)",
826
+ )
827
+ copy_parser.add_argument(
828
+ "--group", default="all", help="Inventory group to target (default: all)"
829
+ )
830
+ copy_parser.add_argument("--parallel", action="store_true", help="Run in parallel")
831
+ copy_parser.add_argument(
832
+ "--max-threads", type=int, default=5, help="Max threads for parallel execution"
833
+ )
834
+
835
+ # Rotate-key command
836
+ rotate_parser = subparsers.add_parser("rotate-key", help="Rotate keys for all")
837
+ rotate_parser.add_argument("--inventory", help="YAML inventory path")
838
+ rotate_parser.add_argument(
839
+ "--key-prefix",
840
+ default="~/.ssh/id_",
841
+ help="Prefix for new key paths (appends hostname)",
842
+ )
843
+ rotate_parser.add_argument(
844
+ "--key-type",
845
+ choices=["rsa", "ed25519"],
846
+ default="ed25519",
847
+ help="Key type to generate (rsa or ed25519, default: ed25519)",
848
+ )
849
+ rotate_parser.add_argument(
850
+ "--group", default="all", help="Inventory group to target (default: all)"
851
+ )
852
+ rotate_parser.add_argument(
853
+ "--parallel", action="store_true", help="Run in parallel"
854
+ )
855
+ rotate_parser.add_argument(
856
+ "--max-threads", type=int, default=5, help="Max threads for parallel execution"
857
+ )
858
+
859
+ # Send-file command
860
+ send_parser = subparsers.add_parser(
861
+ "send-file", help="Upload file to all hosts in inventory"
862
+ )
863
+ send_parser.add_argument("--inventory", help="YAML inventory path")
864
+ send_parser.add_argument("--local-path", help="Local file path to upload")
865
+ send_parser.add_argument("--remote-path", help="Remote destination path")
866
+ send_parser.add_argument(
867
+ "--group", default="all", help="Inventory group to target (default: all)"
868
+ )
869
+ send_parser.add_argument("--parallel", action="store_true", help="Run in parallel")
870
+ send_parser.add_argument(
871
+ "--max-threads", type=int, default=5, help="Max threads for parallel execution"
872
+ )
873
+
874
+ # Receive-file command
875
+ receive_parser = subparsers.add_parser(
876
+ "receive-file", help="Download file from all hosts in inventory"
877
+ )
878
+ receive_parser.add_argument("--inventory", help="YAML inventory path")
879
+ receive_parser.add_argument("--remote-path", help="Remote file path to download")
880
+ receive_parser.add_argument(
881
+ "--local-path-prefix", help="Local directory path prefix to save files"
882
+ )
883
+ receive_parser.add_argument(
884
+ "--group", default="all", help="Inventory group to target (default: all)"
885
+ )
886
+ receive_parser.add_argument(
887
+ "--parallel", action="store_true", help="Run in parallel"
888
+ )
889
+ receive_parser.add_argument(
890
+ "--max-threads", type=int, default=5, help="Max threads for parallel execution"
891
+ )
892
+
893
+ args = parser.parse_args()
894
+
895
+ # Ensure log file directory exists
896
+ if args.log_file:
897
+ log_dir = (
898
+ os.path.dirname(os.path.abspath(args.log_file))
899
+ if os.path.dirname(args.log_file)
900
+ else os.getcwd()
901
+ )
902
+ os.makedirs(log_dir, exist_ok=True)
903
+ try:
904
+ logging.basicConfig(
905
+ filename=args.log_file,
906
+ level=logging.DEBUG,
907
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
908
+ )
909
+ except PermissionError as e:
910
+ print(
911
+ f"Error: Cannot write to log file '{args.log_file}': {str(e)}",
912
+ file=sys.stderr,
913
+ )
914
+ sys.exit(1)
915
+ else:
916
+ logging.basicConfig(
917
+ level=logging.DEBUG,
918
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
919
+ )
920
+
921
+ logger = logging.getLogger("Tunnel")
922
+ logger.debug(
923
+ f"Starting Tunnel Automation with command: {args.command}, args: {vars(args)}"
924
+ )
925
+ print(f"Starting Tunnel Automation with command: {args.command}")
926
+
927
+ try:
928
+ if args.command == "setup-all":
929
+ Tunnel.setup_all_passwordless_ssh(
930
+ args.inventory,
931
+ args.shared_key_path,
932
+ args.key_type,
933
+ args.group,
934
+ args.parallel,
935
+ args.max_threads,
936
+ )
937
+ elif args.command == "run-command":
938
+ Tunnel.run_command_on_inventory(
939
+ args.inventory,
940
+ args.remote_command,
941
+ args.group,
942
+ args.parallel,
943
+ args.max_threads,
944
+ )
945
+ elif args.command == "copy-config":
946
+ Tunnel.copy_ssh_config_on_inventory(
947
+ args.inventory,
948
+ args.local_config_path,
949
+ args.remote_config_path,
950
+ args.group,
951
+ args.parallel,
952
+ args.max_threads,
953
+ )
954
+ elif args.command == "rotate-key":
955
+ Tunnel.rotate_ssh_key_on_inventory(
956
+ args.inventory,
957
+ args.key_prefix,
958
+ args.key_type,
959
+ args.group,
960
+ args.parallel,
961
+ args.max_threads,
962
+ )
963
+ elif args.command == "send-file":
964
+ Tunnel.send_file_on_inventory(
965
+ args.inventory,
966
+ args.local_path,
967
+ args.remote_path,
968
+ args.group,
969
+ args.parallel,
970
+ args.max_threads,
971
+ )
972
+ elif args.command == "receive-file":
973
+ Tunnel.receive_file_on_inventory(
974
+ args.inventory,
975
+ args.remote_path,
976
+ args.local_path_prefix,
977
+ args.group,
978
+ args.parallel,
979
+ args.max_threads,
980
+ )
981
+ logger.debug("Automation Complete")
982
+ print("Automation Complete")
983
+ except Exception as e:
984
+ logger.error(f"Automation failed: {str(e)}")
985
+ print(f"Error: Automation failed: {str(e)}", file=sys.stderr)
986
+ sys.exit(1)
987
+
182
988
 
183
- # Example usage (commented out):
184
- # tunnel = Tunnel("your-remote-host.example.com", log_file="tunnel.log")
185
- # tunnel.connect()
186
- # out, err = tunnel.run_command("ls -la")
187
- # print(out)
188
- # tunnel.send_file("/local/file.txt", "/remote/file.txt")
189
- # tunnel.receive_file("/remote/file.txt", "/local/downloaded.txt")
190
- # tunnel.close()
989
+ if __name__ == "__main__":
990
+ tunnel_manager()