tunnel-manager 1.0.9__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.
@@ -0,0 +1,990 @@
1
+ #!/usr/bin/env python
2
+ # coding: utf-8
3
+
4
+ import sys
5
+ import argparse
6
+ import concurrent.futures
7
+ import logging
8
+ import os
9
+ import paramiko
10
+ import yaml
11
+
12
+
13
+ class Tunnel:
14
+ def __init__(
15
+ self,
16
+ remote_host: str,
17
+ username: str = None,
18
+ password: str = None,
19
+ port: int = 22,
20
+ identity_file: str = None,
21
+ certificate_file: str = None,
22
+ proxy_command: str = None,
23
+ ssh_config_file: str = os.path.expanduser("~/.ssh/config"),
24
+ ):
25
+ """
26
+ Initialize the Tunnel class.
27
+
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).
32
+ :param identity_file: Optional path to the private key file (overrides config).
33
+ :param certificate_file: Optional path to the certificate file (overrides config).
34
+ :param proxy_command: Optional proxy command string (overrides config).
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).
37
+ """
38
+ self.remote_host = remote_host
39
+ self.username = username
40
+ self.password = password
41
+ self.port = port
42
+ self.ssh_client = None
43
+ self.sftp = None
44
+ self.logger = logging.getLogger(__name__)
45
+
46
+ # Load SSH config from custom or default path
47
+ self.ssh_config = paramiko.SSHConfig()
48
+ if os.path.exists(ssh_config_file) and os.path.isfile(ssh_config_file):
49
+ with open(ssh_config_file, "r") as f:
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}")
54
+ host_config = self.ssh_config.lookup(remote_host) or {}
55
+
56
+ self.username = username or host_config.get("user")
57
+ self.identity_file = identity_file or (
58
+ host_config.get("identityfile")[0]
59
+ if host_config.get("identityfile")
60
+ else None
61
+ )
62
+ self.certificate_file = certificate_file or host_config.get("certificatefile")
63
+ self.proxy_command = proxy_command or host_config.get("proxycommand")
64
+
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.")
69
+
70
+ def connect(self):
71
+ if (
72
+ self.ssh_client
73
+ and self.ssh_client.get_transport()
74
+ and self.ssh_client.get_transport().is_active()
75
+ ):
76
+ return
77
+
78
+ self.ssh_client = paramiko.SSHClient()
79
+ self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
80
+
81
+ proxy = None
82
+ if self.proxy_command:
83
+ proxy = paramiko.ProxyCommand(self.proxy_command)
84
+ self.logger.info(f"Using proxy command: {self.proxy_command}")
85
+
86
+ try:
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}")
125
+ except Exception as e:
126
+ self.logger.error(f"Connection failed: {str(e)}")
127
+ raise
128
+
129
+ def run_command(self, command):
130
+ """
131
+ Run a shell command on the remote host.
132
+
133
+ :param command: The command to execute.
134
+ :return: Tuple of (stdout, stderr) as strings.
135
+ """
136
+ self.connect()
137
+ try:
138
+ stdin, stdout, stderr = self.ssh_client.exec_command(command)
139
+ out = stdout.read().decode("utf-8").strip()
140
+ err = stderr.read().decode("utf-8").strip()
141
+ self.logger.info(
142
+ f"Command executed: {command}\nOutput: {out}\nError: {err}"
143
+ )
144
+ return out, err
145
+ except Exception as e:
146
+ self.logger.error(f"Command execution failed: {str(e)}")
147
+ raise
148
+
149
+ def send_file(self, local_path, remote_path):
150
+ """
151
+ Send (upload) a file to the remote host.
152
+ :param local_path: Path to the local file.
153
+ :param remote_path: Path on the remote host.
154
+ """
155
+ self.connect()
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
+
196
+ if not self.sftp:
197
+ self.sftp = self.ssh_client.open_sftp()
198
+ self.logger.debug(f"Opening SFTP for put: {local_path} -> {remote_path}")
199
+ self.sftp.put(local_path, remote_path)
200
+ self.logger.info(f"File sent: {local_path} -> {remote_path}")
201
+ except Exception as 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())
206
+ raise
207
+ finally:
208
+ if self.sftp:
209
+ self.sftp.close()
210
+ self.sftp = None
211
+
212
+ def receive_file(self, remote_path, local_path):
213
+ """
214
+ Receive (download) a file from the remote host.
215
+
216
+ :param remote_path: Path on the remote host.
217
+ :param local_path: Path to save the local file.
218
+ """
219
+ self.connect()
220
+ try:
221
+ if not self.sftp:
222
+ self.sftp = self.ssh_client.open_sftp()
223
+ self.sftp.get(remote_path, local_path)
224
+ self.logger.info(f"File received: {remote_path} -> {local_path}")
225
+ except Exception as e:
226
+ self.logger.error(f"File receive failed: {str(e)}")
227
+ raise
228
+ finally:
229
+ if self.sftp:
230
+ self.sftp.close()
231
+ self.sftp = None
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
+
277
+ def close(self):
278
+ """
279
+ Close the SSH connection.
280
+ """
281
+ if self.ssh_client:
282
+ self.ssh_client.close()
283
+ self.logger.info(f"Connection closed for {self.remote_host}")
284
+ self.ssh_client = None
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
+
988
+
989
+ if __name__ == "__main__":
990
+ tunnel_manager()