Jelajahi Sumber

Add port syntax to host definition and fix sudo issues.

Aurelien 5 tahun lalu
induk
melakukan
5373d98881
1 mengubah file dengan 6 tambahan dan 6 penghapusan
  1. 6 6
      byrd/main.py

+ 6 - 6
byrd/main.py

@@ -73,7 +73,9 @@ def connect(host, auth):
         password = get_password(host)
 
     host_info = host.split('@')
-    hostname = host_info[-1]
+    hostname_port = host_info[-1].split(':')
+    hostname = hostname_port[0]
+    port = hostname_port[1] if len(hostname_port) > 1 else 22
     if len(host_info) > 1:
         username = host_info[0]
     else:
@@ -81,7 +83,7 @@ def connect(host, auth):
 
     client = paramiko.SSHClient()
     client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    client.connect(hostname, username=username, password=password,
+    client.connect(hostname, port=port, username=username, password=password,
                    key_filename=private_key_file)
     logger.debug(f'Connected to {hostname} as {username}')
     CONNECTION_CACHE[host] = client
@@ -239,11 +241,9 @@ def run_helper(client, cmd, env=None, in_buff=None, sudo=False):
 
     if in_buff:
         # XXX use a real buff (not a simple str) ?
-        stdin.write(in_buff if in_buff[-1] == '\n' else in_buff + '\n')
+        in_buff = in_buff[:-1] if in_buff[-1] == '\n' else in_buff
+        stdin.write(in_buff + ('\n' if not sudo else ' && exit $?\n'))
         stdin.flush()
-        if sudo:
-            stdin.write('exit $?\n')
-            stdin.flush()
         stdin.close()
         chan.shutdown_write()