applied the except and raise fixers to the master branch to close the gap with py3
[nepi.git] / src / nepi / util / sshfuncs.py
index c00b494..1899de2 100644 (file)
@@ -3,9 +3,8 @@
 #    Copyright (C) 2013 INRIA
 #
 #    This program is free software: you can redistribute it and/or modify
-#    it under the terms of the GNU General Public License as published by
-#    the Free Software Foundation, either version 3 of the License, or
-#    (at your option) any later version.
+#    it under the terms of the GNU General Public License version 2 as
+#    published by the Free Software Foundation;
 #
 #    This program is distributed in the hope that it will be useful,
 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
@@ -18,6 +17,8 @@
 # Author: Alina Quereilhac <alina.quereilhac@inria.fr>
 #         Claudio Freire <claudio-daniel.freire@inria.fr>
 
+## TODO: This code needs reviewing !!!
+
 import base64
 import errno
 import hashlib
@@ -33,18 +34,17 @@ import threading
 import time
 import tempfile
 
+_re_inet = re.compile("\d+:\s+(?P<name>[a-z0-9_-]+)\s+inet6?\s+(?P<inet>[a-f0-9.:/]+)\s+(brd\s+[0-9.]+)?.*scope\s+global.*") 
+
 logger = logging.getLogger("sshfuncs")
 
-def log(msg, level, out = None, err = None):
+def log(msg, level = logging.DEBUG, out = None, err = None):
     if out:
         msg += " - OUT: %s " % out
-
     if err:
         msg += " - ERROR: %s " % err
-
     logger.log(level, msg)
 
-
 if hasattr(os, "devnull"):
     DEV_NULL = os.devnull
 else:
@@ -57,6 +57,7 @@ class STDOUT:
     Special value that when given to rspawn in stderr causes stderr to 
     redirect to whatever stdout was redirected to.
     """
+    pass
 
 class ProcStatus:
     """
@@ -74,6 +75,24 @@ class ProcStatus:
 hostbyname_cache = dict()
 hostbyname_cache_lock = threading.Lock()
 
+def resolve_hostname(host):
+    ip = None
+
+    if host in ["localhost", "127.0.0.1", "::1"]:
+        p = subprocess.Popen(
+            "ip -o addr list",
+            shell=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        stdout, stderr = p.communicate()
+        m = _re_inet.findall(stdout)
+        ip = m[0][1].split("/")[0]
+    else:
+        ip = socket.gethostbyname(host)
+
+    return ip
+
 def gethostbyname(host):
     global hostbyname_cache
     global hostbyname_cache_lock
@@ -81,7 +100,7 @@ def gethostbyname(host):
     hostbyname = hostbyname_cache.get(host)
     if not hostbyname:
         with hostbyname_cache_lock:
-            hostbyname = socket.gethostbyname(host)
+            hostbyname = resolve_hostname(host)
             hostbyname_cache[host] = hostbyname
 
             msg = " Added hostbyname %s - %s " % (host, hostbyname)
@@ -101,12 +120,15 @@ def openssh_has_persist():
     """
     global OPENSSH_HAS_PERSIST
     if OPENSSH_HAS_PERSIST is None:
-        proc = subprocess.Popen(["ssh","-v"],
-            stdout = subprocess.PIPE,
-            stderr = subprocess.STDOUT,
-            stdin = open("/dev/null","r") )
-        out,err = proc.communicate()
-        proc.wait()
+        with open("/dev/null") as null:
+            proc = subprocess.Popen(
+                ["ssh", "-v"],
+                stdout = subprocess.PIPE,
+                stderr = subprocess.STDOUT,
+                stdin = null,
+            )
+            out,err = proc.communicate()
+            proc.wait()
         
         vre = re.compile(r'OpenSSH_(?:[6-9]|5[.][8-9]|5[.][1-9][0-9]|[1-9][0-9]).*', re.I)
         OPENSSH_HAS_PERSIST = bool(vre.match(out))
@@ -144,9 +166,8 @@ def make_server_key_args(server_key, host, port):
     if os.environ.get('NEPI_STRICT_AUTH_MODE',"").lower() not in ('1','true','on'):
         user_hosts_path = '%s/.ssh/known_hosts' % (os.environ.get('HOME',""),)
         if os.access(user_hosts_path, os.R_OK):
-            f = open(user_hosts_path, "r")
-            tmp_known_hosts.write(f.read())
-            f.close()
+            with open(user_hosts_path, "r") as f:
+                tmp_known_hosts.write(f.read())
         
     tmp_known_hosts.flush()
     
@@ -173,12 +194,12 @@ def shell_escape(s):
         return s
     else:
         # unsafe string - escape
-        def escp(c):
+        def escape(c):
             if (32 <= ord(c) < 127 or c in ('\r','\n','\t')) and c not in ("'",'"'):
                 return c
             else:
                 return "'$'\\x%02x''" % (ord(c),)
-        s = ''.join(map(escp,s))
+        s = ''.join(map(escape, s))
         return "'%s'" % (s,)
 
 def eintr_retry(func):
@@ -190,12 +211,12 @@ def eintr_retry(func):
         for i in xrange(0 if retry else 4):
             try:
                 return func(*p, **kw)
-            except (select.error, socket.error), args:
+            except (select.error, socket.error) as args:
                 if args[0] == errno.EINTR:
                     continue
                 else:
                     raise 
-            except OSError, e:
+            except OSError as e:
                 if e.errno == errno.EINTR:
                     continue
                 else:
@@ -210,15 +231,12 @@ def rexec(command, host, user,
         gw = None, 
         agent = True,
         sudo = False,
-        stdin = None,
         identity = None,
         server_key = None,
         env = None,
         tty = False,
-        timeout = None,
-        retry = 3,
-        err_on_timeout = True,
         connect_timeout = 30,
+        retry = 3,
         persistent = True,
         forward_x11 = False,
         blocking = True,
@@ -226,7 +244,7 @@ def rexec(command, host, user,
     """
     Executes a remote command, returns ((stdout,stderr),process)
     """
-    
+
     tmp_known_hosts = None
     if not gw:
         hostip = gethostbyname(host)
@@ -253,10 +271,7 @@ def rexec(command, host, user,
         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
+        proxycommand = _proxy_command(gw, gwuser, identity)
         args.extend(['-o', proxycommand])
 
     if agent:
@@ -266,6 +281,7 @@ def rexec(command, host, user,
         args.append('-p%d' % port)
 
     if identity:
+        identity = os.path.expanduser(identity)
         args.extend(('-i', identity))
 
     if tty:
@@ -285,65 +301,25 @@ def rexec(command, host, user,
 
     args.append(command)
 
-    for x in xrange(retry):
-        # 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)
-       
-        # attach tempfile object to the process, to make sure the file stays
-        # alive until the process is finished with it
-        proc._known_hosts = tmp_known_hosts
-    
-        # by default, rexec calls _communicate which will block 
-        # until the process has exit. The argument block == False 
-        # forces to rexec to return immediately, without blocking 
-        try:
-            if blocking:
-                out, err = _communicate(proc, stdin, timeout, err_on_timeout)
-            else:
-                err = proc.stderr.read()
-                out = proc.stdout.read()
-
-            msg = " rexec - host %s - command %s " % (host, " ".join(args))
-            log(msg, logging.DEBUG, out, err)
+    log_msg = " rexec - host %s - command %s " % (str(host), " ".join(map(str, args))) 
 
-            if proc.poll():
-                skip = False
-
-                if err.strip().startswith('ssh: ') or err.strip().startswith('mux_client_hello_exchange: '):
-                    # SSH error, can safely retry
-                    skip = True 
-                elif retry:
-                    # Probably timed out or plain failed but can retry
-                    skip = True 
-                
-                if skip:
-                    t = x*2
-                    msg = "SLEEPING %d ... ATEMPT %d - host %s - command %s " % ( 
-                            t, x, host, " ".join(args))
-                    log(msg, logging.DEBUG)
+    stdout = stderr = stdin = subprocess.PIPE
+    if forward_x11:
+        stdout = stderr = stdin = None
 
-                    time.sleep(t)
-                    continue
-            break
-        except RuntimeError, e:
-            msg = " rexec EXCEPTION - host %s - command %s - TIMEOUT -> %s" % (host, " ".join(args), e.args)
-            log(msg, logging.DEBUG, out, err)
-
-            if retry <= 0:
-                raise
-            retry -= 1
-        
-    return ((out, err), proc)
+    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,
         gwuser = None,
         gw = None,
-        agent = True, 
         recursive = False,
         identity = None,
         server_key = None,
@@ -357,7 +333,7 @@ def rcopy(source, dest,
     
     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 string or a string list of colon-separeted files.
+    or a single file in a string.
     """
 
     # Parse destination as <user>@<server>:<path>
@@ -366,16 +342,20 @@ def rcopy(source, dest,
     elif isinstance(source, str) and ':' in source:
         remspec, path = source.split(':',1)
     else:
-        raise ValueError, "Both endpoints cannot be local"
+        raise ValueError("Both endpoints cannot be local")
     user,host = remspec.rsplit('@',1)
     
     # plain scp
     tmp_known_hosts = None
 
     args = ['scp', '-q', '-p', '-C',
+            # 2015-06-01 Thierry: I am commenting off blowfish
+            # as this is not available on a plain ubuntu 15.04 install
+            # this IMHO is too fragile, shoud be something the user
+            # decides explicitly (so he is at least aware of that dependency)
             # Speed up transfer using blowfish cypher specification which is 
             # faster than the default one (3des)
-            '-c', 'blowfish',
+            '-c', 'blowfish',
             # Don't bother with localhost. Makes test easier
             '-o', 'NoHostAuthenticationForLocalhost=yes',
             '-o', 'ConnectTimeout=60',
@@ -387,16 +367,14 @@ def rcopy(source, dest,
         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
+        proxycommand = _proxy_command(gw, gwuser, identity)
         args.extend(['-o', proxycommand])
 
     if recursive:
         args.append('-r')
 
     if identity:
+        identity = os.path.expanduser(identity)
         args.extend(('-i', identity))
 
     if server_key:
@@ -407,76 +385,45 @@ def rcopy(source, dest,
     if not strict_host_checking:
         # Do not check for Host key. Unsafe.
         args.extend(['-o', 'StrictHostKeyChecking=no'])
-
-    if openssh_has_persist():
-        args.extend([
-            '-o', 'ControlMaster=auto',
-            '-o', 'ControlPath=%s' % (make_control_path(agent, False),)
-            ])
-
-    if isinstance(dest, str):
-        dest = map(str.strip, dest.split(";"))
-
-    if isinstance(source, str):
-        source = map(str.strip, source.split(";"))
-
-    args.extend(source)
-
-    args.extend(dest)
-
-    for x in xrange(retry):
-        # connects to the remote host and starts a remote connection
-        proc = subprocess.Popen(args,
-                stdout = subprocess.PIPE,
-                stdin = subprocess.PIPE, 
-                stderr = subprocess.PIPE)
-        
-        # attach tempfile object to the process, to make sure the file stays
-        # alive until the process is finished with it
-        proc._known_hosts = tmp_known_hosts
     
-        try:
-            (out, err) = proc.communicate()
-            eintr_retry(proc.wait)()
-            msg = " rcopy - host %s - command %s " % (host, " ".join(args))
-            log(msg, logging.DEBUG, out, err)
-
-            if proc.poll():
-                t = x*2
-                msg = "SLEEPING %d ... ATEMPT %d - host %s - command %s " % ( 
-                        t, x, host, " ".join(args))
-                log(msg, logging.DEBUG)
-
-                time.sleep(t)
-                continue
-
-            break
-        except RuntimeError, e:
-            msg = " rcopy EXCEPTION - host %s - command %s - TIMEOUT -> %s" % (host, " ".join(args), e.args)
-            log(msg, logging.DEBUG, out, err)
+    if isinstance(source, list):
+        args.extend(source)
+    else:
+        if openssh_has_persist():
+            args.extend([
+                '-o', 'ControlMaster=auto',
+                '-o', 'ControlPath=%s' % (make_control_path(False, False),)
+                ])
+        args.append(source)
+
+    if isinstance(dest, list):
+        args.extend(dest)
+    else:
+        args.append(dest)
 
-            if retry <= 0:
-                raise
-            retry -= 1
-        
-    return ((out, err), proc)
+    log_msg = " rcopy - host %s - command %s " % (str(host), " ".join(map(str, args)))
+    
+    return _retry_rexec(args, log_msg, env = None, retry = retry, 
+            tmp_known_hosts = tmp_known_hosts,
+            blocking = True)
 
 def rspawn(command, pidfile, 
-        stdout = '/dev/null', 
-        stderr = STDOUT, 
-        stdin = '/dev/null',
-        home = None, 
-        create_home = False, 
-        sudo = False,
-        host = None, 
-        port = None, 
-        user = None, 
-        gwuser = None,
-        gw = None,
-        agent = None, 
-        identity = None, 
-        server_key = None,
-        tty = False):
+           stdout = '/dev/null', 
+           stderr = STDOUT, 
+           stdin = '/dev/null',
+           home = None, 
+           create_home = False, 
+           sudo = False,
+           host = None, 
+           port = None, 
+           user = None, 
+           gwuser = None,
+           gw = None,
+           agent = None, 
+           identity = None, 
+           server_key = None,
+           tty = False,
+           strict_host_checking = True):
     """
     Spawn a remote command such that it will continue working asynchronously in 
     background. 
@@ -511,7 +458,7 @@ def rspawn(command, pidfile,
         :param sudo: Flag forcing execution with sudo user
         :type sudo: bool
         
-        :rtype: touple
+        :rtype: tuple
 
         (stdout, stderr), process
         
@@ -551,24 +498,26 @@ def rspawn(command, pidfile,
         agent = agent,
         identity = identity,
         server_key = server_key,
-        tty = tty ,
+        tty = tty,
+        strict_host_checking = strict_host_checking ,
         )
     
     if proc.wait():
-        raise RuntimeError, "Failed to set up application on host %s: %s %s" % (host, out,err,)
+        raise RuntimeError("Failed to set up application on host %s: %s %s" % (host, out,err,))
 
     return ((out, err), proc)
 
 @eintr_retry
 def rgetpid(pidfile,
-        host = None, 
-        port = None, 
-        user = None, 
-        gwuser = None,
-        gw = None,
-        agent = None, 
-        identity = None,
-        server_key = None):
+            host = None, 
+            port = None, 
+            user = None, 
+            gwuser = None,
+            gw = None,
+            agent = None, 
+            identity = None,
+            server_key = None,
+            strict_host_checking = True):
     """
     Returns the pid and ppid of a process from a remote file where the 
     information was stored.
@@ -596,7 +545,8 @@ def rgetpid(pidfile,
         gw = gw,
         agent = agent,
         identity = identity,
-        server_key = server_key
+        server_key = server_key,
+        strict_host_checking = strict_host_checking
         )
         
     if proc.wait():
@@ -618,7 +568,8 @@ def rstatus(pid, ppid,
         gw = None,
         agent = None, 
         identity = None,
-        server_key = None):
+        server_key = None,
+        strict_host_checking = True):
     """
     Returns a code representing the the status of a remote process
 
@@ -644,7 +595,8 @@ def rstatus(pid, ppid,
         gw = gw,
         agent = agent,
         identity = identity,
-        server_key = server_key
+        server_key = server_key,
+        strict_host_checking = strict_host_checking
         )
     
     if proc.wait():
@@ -671,7 +623,8 @@ def rkill(pid, ppid,
         sudo = False,
         identity = None, 
         server_key = None, 
-        nowait = False):
+        nowait = False,
+        strict_host_checking = True):
     """
     Sends a kill signal to a remote process.
 
@@ -725,7 +678,8 @@ fi
         gw = gw,
         agent = agent,
         identity = identity,
-        server_key = server_key
+        server_key = server_key,
+        strict_host_checking = strict_host_checking
         )
     
     # wait, don't leave zombies around
@@ -733,15 +687,89 @@ fi
 
     return (out, err), proc
 
+def _retry_rexec(args,
+                 log_msg,
+                 stdout = subprocess.PIPE,
+                 stdin = subprocess.PIPE, 
+                 stderr = subprocess.PIPE,
+                 env = None,
+                 retry = 3,
+                 tmp_known_hosts = None,
+                 blocking = True):
+
+    for x in xrange(retry):
+        # display command actually invoked when debug is turned on
+        message = " ".join( [ "'{}'".format(arg) for arg in args ] )
+        log("sshfuncs: invoking {}".format(message), logging.DEBUG)
+        # connects to the remote host and starts a remote connection
+        proc = subprocess.Popen(
+            args,
+            env = env,
+            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
+        proc._known_hosts = tmp_known_hosts
+    
+        # The argument block == False forces to rexec to return immediately, 
+        # without blocking 
+        try:
+            err = out = " "
+            if blocking:
+                #(out, err) = proc.communicate()
+                # The method communicate was re implemented for performance issues
+                # when using python subprocess communicate method the ssh commands 
+                # last one minute each
+                out, err = _communicate(proc, input=None)
+
+            elif stdout:
+                out = proc.stdout.read()
+                if proc.poll() and stderr:
+                    err = proc.stderr.read()
+
+            log(log_msg, logging.DEBUG, out, err)
+
+            if proc.poll():
+                skip = False
+
+                if err.strip().startswith('ssh: ') or err.strip().startswith('mux_client_hello_exchange: '):
+                    # SSH error, can safely retry
+                    skip = True 
+                elif retry:
+                    # Probably timed out or plain failed but can retry
+                    skip = True 
+                
+                if skip:
+                    t = x*2
+                    msg = "SLEEPING %d ... ATEMPT %d - command %s " % ( 
+                            t, x, " ".join(args))
+                    log(msg, logging.DEBUG)
+
+                    time.sleep(t)
+                    continue
+            break
+        except RuntimeError as e:
+            msg = " rexec EXCEPTION - TIMEOUT -> %s \n %s" % ( e.args, log_msg )
+            log(msg, logging.DEBUG, out, err)
+
+            if retry <= 0:
+                raise
+            retry -= 1
+
+    return ((out, err), proc)
+
 # POSIX
+# Don't remove. The method communicate was re implemented for performance issues
 def _communicate(proc, input, timeout=None, err_on_timeout=True):
     read_set = []
     write_set = []
     stdout = None # Return
     stderr = None # Return
-    
+
     killed = False
-    
+
     if timeout is not None:
         timelimit = time.time() + timeout
         killtime = timelimit + 4
@@ -782,18 +810,18 @@ def _communicate(proc, input, timeout=None, err_on_timeout=True):
                 select_timeout = timelimit - curtime + 0.1
         else:
             select_timeout = 1.0
-        
+
         if select_timeout > 1.0:
             select_timeout = 1.0
-            
+
         try:
             rlist, wlist, xlist = select.select(read_set, write_set, [], select_timeout)
-        except select.error,e:
+        except select.error as e:
             if e[0] != 4:
                 raise
             else:
                 continue
-        
+
         if not rlist and not wlist and not xlist and proc.poll() is not None:
             # timeout and process exited, say bye
             break
@@ -823,7 +851,7 @@ def _communicate(proc, input, timeout=None, err_on_timeout=True):
                 proc.stderr.close()
                 read_set.remove(proc.stderr)
             stderr.append(data)
-    
+
     # All data exchanged.  Translate lists into strings.
     if stdout is not None:
         stdout = ''.join(stdout)
@@ -842,7 +870,7 @@ def _communicate(proc, input, timeout=None, err_on_timeout=True):
 
     if killed and err_on_timeout:
         errcode = proc.poll()
-        raise RuntimeError("Operation timed out", errcode, stdout, stderr)
+        raise RuntimeError("Operation timed out", errcode, stdout, stderr)
     else:
         if killed:
             proc.poll()
@@ -850,3 +878,33 @@ def _communicate(proc, input, timeout=None, err_on_timeout=True):
             proc.wait()
         return (stdout, stderr)
 
+def _proxy_command(gw, gwuser, gwidentity):
+    """
+    Constructs the SSH ProxyCommand option to add to the SSH command to connect
+    via a proxy
+        :param gw: SSH proxy hostname
+        :type gw:  str 
+       
+        :param gwuser: SSH proxy username
+        :type gwuser:  str
+
+        :param gwidentity: SSH proxy identity file 
+        :type gwidentity:  str
+
+  
+        :rtype: str 
+        
+        returns the SSH ProxyCommand option.
+    """
+
+    proxycommand = 'ProxyCommand=ssh -q '
+    if gwidentity:
+        proxycommand += '-i %s ' % os.path.expanduser(gwidentity)
+    if gwuser:
+        proxycommand += '%s' % gwuser
+    else:
+        proxycommand += '%r'
+    proxycommand += '@%s -W %%h:%%p' % gw
+
+    return proxycommand
+