Add demo file

Author: Olivier Roques

examples/demo.py

@@ -0,0 +1,100 @@
+#!/usr/bin/python3
+import logging
+import queue
+import sshcontroller
+import threading
+import time
+
+logging.basicConfig(format="[%(levelname)s] %(message)s", level=logging.INFO)
+
+HOST_IP = "93.184.216.34"  # an IPv4 or IPv6 address
+KEY_PWD = ""
+SSH_PWD = ""
+
+
+def demo_key():
+    ssh_controller = sshcontroller.SSHController(
+        host=HOST_IP,
+        user="olivier",
+        key_path="~/.ssh/id_rsa",  # if omitted, look in agent and in ~/.ssh
+        key_password=KEY_PWD,      # optional
+        key_type="rsa",            # rsa (default), dsa, ecdsa or ed25519
+        port=22,                   # 22 is the default
+    )
+
+    ssh_controller.connect()
+
+    return_code, output = ssh_controller.run(
+        command="echo 'Hello world!' > /tmp/hello.txt",
+        display=True,          # display output, false by default
+        combine_stderr=False,  # combine stderr and stdout, false by default
+        capture_output=True,   # return output, false by default
+        timeout=10,            # command timeout in seconds, 600s by default
+    )
+    logging.info(f"return code: {return_code}, output: {output}")
+
+    print(f"hello.txt exists: {ssh_controller.exists('/tmp/hello.txt')}")
+    print(f"bonjour.txt exists: {ssh_controller.exists('/tmp/bonjour.txt')}")
+
+    ssh_controller.get("/tmp/hello.txt", "/tmp/bonjour.txt")
+
+    with open("/tmp/bonjour.txt", 'r') as bonjour:
+        for line in bonjour:
+            print(line, end='')
+
+    ssh_controller.disconnect()
+
+
+def demo_pwd():
+    ssh_controller = sshcontroller.SSHController(
+        host=HOST_IP,
+        user="root",
+        ssh_password=SSH_PWD
+    )
+    output = queue.Queue()
+    stop_event_sleep = threading.Event()
+    stop_event_ping = threading.Event()
+    kwargs_sleep = {
+        "command": "echo 'thread sleep: sleeping for 10s' && sleep 10s",
+        "display": True,
+        "stop_event": stop_event_sleep,
+    }
+    kwargs_ping = {
+        "command": "echo 'thread ping: starting ping' && ping localhost",
+        "display": True,
+        "capture_output": True,
+        "stop_event": stop_event_ping,
+    }
+
+    def wrapper(kwargs):
+        return output.put(ssh_controller.run(**kwargs))
+
+    ssh_controller.connect()
+
+    thread_sleep = threading.Thread(
+        target=ssh_controller.run, name="thread_sleep", kwargs=kwargs_sleep)
+    thread_ping = threading.Thread(
+        target=wrapper, name="thread_ping", args=(kwargs_ping, ))
+
+    thread_ping.start()
+    thread_sleep.start()
+
+    try:
+        thread_sleep.join()
+    except KeyboardInterrupt:
+        logging.info("KeyboardInterrupt")
+    finally:
+        logging.info("Stopping threads")
+        stop_event_sleep.set()
+        stop_event_ping.set()
+        time.sleep(2)
+
+    return_code, ping_output = output.get()
+    logging.info(f"thread ping return code: {return_code}")
+    logging.info(f"thread ping output length: {len(ping_output)}")
+
+    ssh_controller.disconnect()
+
+
+demo_key()
+demo_pwd()

sshcontroller/sshcontroller.py

@@ -69,18 +69,22 @@ def __init__(
                 )
             )
         elif ssh_password is None:
+            self.keys = paramiko.Agent().get_keys()
+
             try:
-                self.keys.append(
-                    _KEY_TYPES[key_type].from_private_key(
-                        open(path.expanduser(f"~/.ssh/id_{key_type}"), 'r'),
-                        key_password
-                    )
+                key_file = _KEY_TYPES[key_type].from_private_key(
+                    open(path.expanduser(f"~/.ssh/id_{key_type}"), 'r'),
+                    key_password
                 )
-            except Exception as e:
-                agent_keys = paramiko.Agent().get_keys()
-                if not agent_keys:
-                    raise e
-                self.keys = agent_keys
+            except Exception:
+                pass
+            else:
+                self.keys.insert(
+                    len(self.keys) if key_password is None else 0, key_file
+                )
+
+            if not self.keys:
+                logging.error("No valid key found")
 
     def connect(self):
         try:
@@ -123,7 +127,7 @@ def __run_until_event(
         capture_output=False,
     ):
         channel = self.transport.open_session()
-        output = []
+        output = ""
         timeout = 2
 
         channel.settimeout(timeout)
@@ -145,10 +149,10 @@ def __run_until_event(
                 if not len(raw_data):
                     break
 
-                data = raw_data.decode("utf-8").splitlines()
+                data = raw_data.decode("utf-8")
 
                 if display:
-                    print('\n'.join(data))
+                    print(data, end='')
 
                 if capture_output:
                     output += data
@@ -157,7 +161,7 @@ def __run_until_event(
                     break
 
         channel.close()
-        return (channel.exit_status_ready(), output)
+        return (channel.exit_status_ready(), output.splitlines())
 
     def __run_until_exit(
         self,
@@ -168,7 +172,7 @@ def __run_until_exit(
         capture_output=False,
     ):
         channel = self.transport.open_session()
-        output = []
+        output = ""
 
         channel.settimeout(timeout)
         channel.set_combine_stderr(combine_stderr)
@@ -177,31 +181,31 @@ def __run_until_exit(
 
         try:
             if not display and not capture_output:
-                return (channel.recv_exit_status(), output)
+                return (channel.recv_exit_status(), output.splitlines())
             else:
                 while True:
                     raw_data = channel.recv(self.nb_bytes)
 
                     if not len(raw_data):
                         break
 
-                    data = raw_data.decode("utf-8").splitlines()
+                    data = raw_data.decode("utf-8")
 
                     if display:
-                        print('\n'.join(data))
+                        print(data, end='')
 
                     if capture_output:
                         output += data
         except socket.timeout:
             logging.warning(f"Timeout after {timeout}s")
-            return (1, output)
+            return (1, output.splitlines())
         except KeyboardInterrupt:
             logging.info("KeyboardInterrupt")
-            return (0, output)
+            return (0, output.splitlines())
         finally:
             channel.close()
 
-        return (channel.recv_exit_status(), output)
+        return (channel.recv_exit_status(), output.splitlines())
 
     def run(
         self,