cast --port argument to int

check SSHTunnel variables before referencing fixes #9
ThreatResponse · Sep 24, 2016 · f2a2b97 · f2a2b97
1 parent daef4e1
commit f2a2b97
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/margaritashotgun/cli.py b/margaritashotgun/cli.py
@@ -163,7 +163,7 @@ def configure_args(self, arguments):
 
         if arguments.server is not None:
             host = dict(zip(host_allowed_keys,
-                            [arguments.server, arguments.port,
+                            [arguments.server, int(arguments.port),
                              arguments.username, arguments.password,
                              module, key, arguments.filename]))
             args_config['hosts'] = []

diff --git a/margaritashotgun/remote_host.py b/margaritashotgun/remote_host.py
@@ -136,7 +136,7 @@ def connect(self, username, password, key, address, port):
         if port is None:
             self.remote_port = 22
         else:
-            self.remote_port = port
+            self.remote_port = int(port)
         auth = Auth(username=username, password=password, key=key)
         self.tunnel.connect(auth, address, self.remote_port)
         self.shell.connect(auth, address, self.remote_port)

diff --git a/margaritashotgun/ssh_tunnel.py b/margaritashotgun/ssh_tunnel.py
@@ -21,6 +21,7 @@ class SSHTunnel():
     def __init__(self):
         self.transport = None
         self.forward = None
+        self.local_port = None
 
     # TODO raise exceptions if address==None
     def connect(self, auth, address, port, hostkey=None):
@@ -112,10 +113,12 @@ def cleanup(self):
         """
         Cleanup resources used during execution
         """
-        logger.debug(("Stopping ssh tunnel {0}:{1}:{2} for "
-                      "{3}@{4}".format(self.local_port,
-                                       self.remote_address, self.remote_port,
-                                       self.username, self.address)))
+        if self.local_port is not None:
+            logger.debug(("Stopping ssh tunnel {0}:{1}:{2} for "
+                          "{3}@{4}".format(self.local_port,
+                                           self.remote_address,
+                                           self.remote_port,
+                                           self.username, self.address)))
         if self.forward is not None:
             self.forward.stop()
             self.forward.join()