Fix #128 - [NS3] Test ns-3 with localhost

[nepi.git] / src / nepi / util / sshfuncs.py

diff --git a/src/nepi/util/sshfuncs.py b/src/nepi/util/sshfuncs.py
index 8c8ccf1..d1952bc 100644 (file)
--- a/src/nepi/util/sshfuncs.py
+++ b/src/nepi/util/sshfuncs.py
@@ -46,7 +46,6 @@ def log(msg, level, out = None, err = None):
 
     logger.log(level, msg)
 
-
 if hasattr(os, "devnull"):
     DEV_NULL = os.devnull
 else:
@@ -207,10 +206,11 @@ def eintr_retry(func):
     return rv
 
 def rexec(command, host, user, 
-        port = None, 
+        port = None,
+        gwuser = None,
+        gw = None, 
         agent = True,
         sudo = False,
-        stdin = None,
         identity = None,
         server_key = None,
         env = None,
@@ -226,7 +226,9 @@ def rexec(command, host, user,
     """
     
     tmp_known_hosts = None
-    hostip = gethostbyname(host)
+    if not gw:
+        hostip = gethostbyname(host)
+    else: hostip = None
 
     args = ['ssh', '-C',
             # Don't bother with localhost. Makes test easier
@@ -235,6 +237,7 @@ def rexec(command, host, user,
             '-o', 'ConnectionAttempts=3',
             '-o', 'ServerAliveInterval=30',
             '-o', 'TCPKeepAlive=yes',
+            '-o', 'Batchmode=yes',
             '-l', user, hostip or host]
 
     if persistent and openssh_has_persist():
@@ -247,6 +250,13 @@ def rexec(command, host, user,
         # Do not check for Host key. Unsafe.
         args.extend(['-o', 'StrictHostKeyChecking=no'])
 
+    if gw:
+        if gwuser:
+            proxycommand = 'ProxyCommand=ssh %s@%s -W %%h:%%p' % (gwuser, gw)
+        else:
+            proxycommand = 'ProxyCommand=ssh %%r@%s -W %%h:%%p' % gw
+        args.extend(['-o', proxycommand])
+
     if agent:
         args.append('-A')
 
@@ -254,6 +264,7 @@ def rexec(command, host, user,
         args.append('-p%d' % port)
 
     if identity:
+        identity = os.path.expanduser(identity)
         args.extend(('-i', identity))
 
     if tty:
@@ -275,13 +286,23 @@ def rexec(command, host, user,
     
     log_msg = " rexec - host %s - command %s " % (host, " ".join(args))
 
-    return _retry_rexec(args, log_msg, env = env, retry = retry, 
+    stdout = stderr = stdin = subprocess.PIPE
+    if forward_x11:
+        stdout = stderr = stdin = None
+
+    return _retry_rexec(args, log_msg, 
+            stderr = stderr,
+            stdin = stdin,
+            stdout = stdout,
+            env = env, 
+            retry = retry, 
             tmp_known_hosts = tmp_known_hosts,
             blocking = blocking)
 
 def rcopy(source, dest,
-        port = None, 
-        agent = True, 
+        port = None,
+        gwuser = None,
+        gw = None,
         recursive = False,
         identity = None,
         server_key = None,
@@ -293,14 +314,16 @@ def rcopy(source, dest,
     Source and destination should have the user and host encoded
     as per scp specs.
     
-    Source can be a list of files to copy to a single destination,
-    in which case it is advised that the destination be a folder.
+    Source can be a list of files to copy to a single destination, 
+    (in which case it is advised that the destination be a folder),
+    a single file in a string or a semi-colon separated list of files
+    in a string.
     """
-    
+
     # Parse destination as <user>@<server>:<path>
-    if isinstance(dest, basestring) and ':' in dest:
+    if isinstance(dest, str) and ':' in dest:
         remspec, path = dest.split(':',1)
-    elif isinstance(source, basestring) and ':' in source:
+    elif isinstance(source, str) and ':' in source:
         remspec, path = source.split(':',1)
     else:
         raise ValueError, "Both endpoints cannot be local"
@@ -323,10 +346,18 @@ def rcopy(source, dest,
     if port:
         args.append('-P%d' % port)
 
+    if gw:
+        if gwuser:
+            proxycommand = 'ProxyCommand=ssh %s@%s -W %%h:%%p' % (gwuser, gw)
+        else:
+            proxycommand = 'ProxyCommand=ssh %%r@%s -W %%h:%%p' % gw
+        args.extend(['-o', proxycommand])
+
     if recursive:
         args.append('-r')
 
     if identity:
+        identity = os.path.expanduser(identity)
         args.extend(('-i', identity))
 
     if server_key:
@@ -337,18 +368,21 @@ def rcopy(source, dest,
     if not strict_host_checking:
         # Do not check for Host key. Unsafe.
         args.extend(['-o', 'StrictHostKeyChecking=no'])
-
-    if isinstance(source,list):
+    
+    if isinstance(source, list):
         args.extend(source)
     else:
         if openssh_has_persist():
             args.extend([
                 '-o', 'ControlMaster=auto',
-                '-o', 'ControlPath=%s' % (make_control_path(agent, False),)
+                '-o', 'ControlPath=%s' % (make_control_path(False, False),)
                 ])
         args.append(source)
 
-    args.append(dest)
+    if isinstance(dest, list):
+        args.extend(dest)
+    else:
+        args.append(dest)
 
     log_msg = " rcopy - host %s - command %s " % (host, " ".join(args))
     
@@ -366,6 +400,8 @@ def rspawn(command, pidfile,
         host = None, 
         port = None, 
         user = None, 
+        gwuser = None,
+        gw = None,
         agent = None, 
         identity = None, 
         server_key = None,
@@ -439,6 +475,8 @@ def rspawn(command, pidfile,
         host = host,
         port = port,
         user = user,
+        gwuser = gwuser,
+        gw = gw,
         agent = agent,
         identity = identity,
         server_key = server_key,
@@ -455,6 +493,8 @@ def rgetpid(pidfile,
         host = None, 
         port = None, 
         user = None, 
+        gwuser = None,
+        gw = None,
         agent = None, 
         identity = None,
         server_key = None):
@@ -481,6 +521,8 @@ def rgetpid(pidfile,
         host = host,
         port = port,
         user = user,
+        gwuser = gwuser,
+        gw = gw,
         agent = agent,
         identity = identity,
         server_key = server_key
@@ -501,6 +543,8 @@ def rstatus(pid, ppid,
         host = None, 
         port = None, 
         user = None, 
+        gwuser = None,
+        gw = None,
         agent = None, 
         identity = None,
         server_key = None):
@@ -525,6 +569,8 @@ def rstatus(pid, ppid,
         host = host,
         port = port,
         user = user,
+        gwuser = gwuser,
+        gw = gw,
         agent = agent,
         identity = identity,
         server_key = server_key
@@ -548,6 +594,8 @@ def rkill(pid, ppid,
         host = None, 
         port = None, 
         user = None, 
+        gwuser = None,
+        gw = None,
         agent = None, 
         sudo = False,
         identity = None, 
@@ -602,6 +650,8 @@ fi
         host = host,
         port = port,
         user = user,
+        gwuser = gwuser,
+        gw = gw,
         agent = agent,
         identity = identity,
         server_key = server_key
@@ -614,6 +664,9 @@ fi
 
 def _retry_rexec(args,
         log_msg,
+        stdout = subprocess.PIPE,
+        stdin = subprocess.PIPE, 
+        stderr = subprocess.PIPE,
         env = None,
         retry = 3,
         tmp_known_hosts = None,
@@ -623,9 +676,9 @@ def _retry_rexec(args,
         # connects to the remote host and starts a remote connection
         proc = subprocess.Popen(args,
                 env = env,
-                stdout = subprocess.PIPE,
-                stdin = subprocess.PIPE, 
-                stderr = subprocess.PIPE)
+                stdout = stdout,
+                stdin = stdin, 
+                stderr = stderr)
         
         # attach tempfile object to the process, to make sure the file stays
         # alive until the process is finished with it
@@ -634,11 +687,13 @@ def _retry_rexec(args,
         # The argument block == False forces to rexec to return immediately, 
         # without blocking 
         try:
+            err = out = " "
             if blocking:
                 (out, err) = proc.communicate()
-            else:
-                err = proc.stderr.read()
+            elif stdout:
                 out = proc.stdout.read()
+                if proc.poll() and stderr:
+                    err = proc.stderr.read()
 
             log(log_msg, logging.DEBUG, out, err)
 
@@ -671,4 +726,3 @@ def _retry_rexec(args,
         
     return ((out, err), proc)
 
-