Trying to make LinuxNS3Simulator to deploy remotely ....
[nepi.git] / src / nepi / util / sshfuncs.py
1 #
2 #    NEPI, a framework to manage network experiments
3 #    Copyright (C) 2013 INRIA
4 #
5 #    This program is free software: you can redistribute it and/or modify
6 #    it under the terms of the GNU General Public License as published by
7 #    the Free Software Foundation, either version 3 of the License, or
8 #    (at your option) any later version.
9 #
10 #    This program is distributed in the hope that it will be useful,
11 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
12 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 #    GNU General Public License for more details.
14 #
15 #    You should have received a copy of the GNU General Public License
16 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
17 #
18 # Author: Alina Quereilhac <alina.quereilhac@inria.fr>
19 #         Claudio Freire <claudio-daniel.freire@inria.fr>
20
21 ## TODO: This code needs reviewing !!!
22
23 import base64
24 import errno
25 import hashlib
26 import logging
27 import os
28 import os.path
29 import re
30 import select
31 import signal
32 import socket
33 import subprocess
34 import threading
35 import time
36 import tempfile
37
38 logger = logging.getLogger("sshfuncs")
39
40 def log(msg, level, out = None, err = None):
41     if out:
42         msg += " - OUT: %s " % out
43
44     if err:
45         msg += " - ERROR: %s " % err
46
47     logger.log(level, msg)
48
49
50 if hasattr(os, "devnull"):
51     DEV_NULL = os.devnull
52 else:
53     DEV_NULL = "/dev/null"
54
55 SHELL_SAFE = re.compile('^[-a-zA-Z0-9_=+:.,/]*$')
56
57 class STDOUT: 
58     """
59     Special value that when given to rspawn in stderr causes stderr to 
60     redirect to whatever stdout was redirected to.
61     """
62
63 class ProcStatus:
64     """
65     Codes for status of remote spawned process
66     """
67     # Process is still running
68     RUNNING = 1
69
70     # Process is finished
71     FINISHED = 2
72     
73     # Process hasn't started running yet (this should be very rare)
74     NOT_STARTED = 3
75
76 hostbyname_cache = dict()
77 hostbyname_cache_lock = threading.Lock()
78
79 def gethostbyname(host):
80     global hostbyname_cache
81     global hostbyname_cache_lock
82     
83     hostbyname = hostbyname_cache.get(host)
84     if not hostbyname:
85         with hostbyname_cache_lock:
86             hostbyname = socket.gethostbyname(host)
87             hostbyname_cache[host] = hostbyname
88
89             msg = " Added hostbyname %s - %s " % (host, hostbyname)
90             log(msg, logging.DEBUG)
91
92     return hostbyname
93
94 OPENSSH_HAS_PERSIST = None
95
96 def openssh_has_persist():
97     """ The ssh_config options ControlMaster and ControlPersist allow to
98     reuse a same network connection for multiple ssh sessions. In this 
99     way limitations on number of open ssh connections can be bypassed.
100     However, older versions of openSSH do not support this feature.
101     This function is used to determine if ssh connection persist features
102     can be used.
103     """
104     global OPENSSH_HAS_PERSIST
105     if OPENSSH_HAS_PERSIST is None:
106         proc = subprocess.Popen(["ssh","-v"],
107             stdout = subprocess.PIPE,
108             stderr = subprocess.STDOUT,
109             stdin = open("/dev/null","r") )
110         out,err = proc.communicate()
111         proc.wait()
112         
113         vre = re.compile(r'OpenSSH_(?:[6-9]|5[.][8-9]|5[.][1-9][0-9]|[1-9][0-9]).*', re.I)
114         OPENSSH_HAS_PERSIST = bool(vre.match(out))
115     return OPENSSH_HAS_PERSIST
116
117 def make_server_key_args(server_key, host, port):
118     """ Returns a reference to a temporary known_hosts file, to which 
119     the server key has been added. 
120     
121     Make sure to hold onto the temp file reference until the process is 
122     done with it
123
124     :param server_key: the server public key
125     :type server_key: str
126
127     :param host: the hostname
128     :type host: str
129
130     :param port: the ssh port
131     :type port: str
132
133     """
134     if port is not None:
135         host = '%s:%s' % (host, str(port))
136
137     # Create a temporary server key file
138     tmp_known_hosts = tempfile.NamedTemporaryFile()
139    
140     hostbyname = gethostbyname(host) 
141
142     # Add the intended host key
143     tmp_known_hosts.write('%s,%s %s\n' % (host, hostbyname, server_key))
144     
145     # If we're not in strict mode, add user-configured keys
146     if os.environ.get('NEPI_STRICT_AUTH_MODE',"").lower() not in ('1','true','on'):
147         user_hosts_path = '%s/.ssh/known_hosts' % (os.environ.get('HOME',""),)
148         if os.access(user_hosts_path, os.R_OK):
149             f = open(user_hosts_path, "r")
150             tmp_known_hosts.write(f.read())
151             f.close()
152         
153     tmp_known_hosts.flush()
154     
155     return tmp_known_hosts
156
157 def make_control_path(agent, forward_x11):
158     ctrl_path = "/tmp/nepi_ssh"
159
160     if agent:
161         ctrl_path +="_a"
162
163     if forward_x11:
164         ctrl_path +="_x"
165
166     ctrl_path += "-%r@%h:%p"
167
168     return ctrl_path
169
170 def shell_escape(s):
171     """ Escapes strings so that they are safe to use as command-line 
172     arguments """
173     if SHELL_SAFE.match(s):
174         # safe string - no escaping needed
175         return s
176     else:
177         # unsafe string - escape
178         def escp(c):
179             if (32 <= ord(c) < 127 or c in ('\r','\n','\t')) and c not in ("'",'"'):
180                 return c
181             else:
182                 return "'$'\\x%02x''" % (ord(c),)
183         s = ''.join(map(escp,s))
184         return "'%s'" % (s,)
185
186 def eintr_retry(func):
187     """Retries a function invocation when a EINTR occurs"""
188     import functools
189     @functools.wraps(func)
190     def rv(*p, **kw):
191         retry = kw.pop("_retry", False)
192         for i in xrange(0 if retry else 4):
193             try:
194                 return func(*p, **kw)
195             except (select.error, socket.error), args:
196                 if args[0] == errno.EINTR:
197                     continue
198                 else:
199                     raise 
200             except OSError, e:
201                 if e.errno == errno.EINTR:
202                     continue
203                 else:
204                     raise
205         else:
206             return func(*p, **kw)
207     return rv
208
209 def socat(local_socket_name, remote_socket_name,
210         host, user,
211         port = None, 
212         agent = True,
213         sudo = False,
214         identity = None,
215         server_key = None,
216         env = None,
217         tty = False,
218         connect_timeout = 30,
219         retry = 3,
220         strict_host_checking = True):
221     """
222     Executes a remote command, returns ((stdout,stderr),process)
223     """
224     
225     tmp_known_hosts = None
226     hostip = gethostbyname(host)
227
228
229     args = ["socat"]
230     args.append("UNIX-LISTEN:%s,unlink-early,fork" % local_socket_name)
231
232     ssh_args = ['ssh', '-C',
233             # Don't bother with localhost. Makes test easier
234             '-o', 'NoHostAuthenticationForLocalhost=yes',
235             '-o', 'ConnectTimeout=%d' % (int(connect_timeout),),
236             '-o', 'ConnectionAttempts=3',
237             '-o', 'ServerAliveInterval=30',
238             '-o', 'TCPKeepAlive=yes',
239             '-l', user, hostip or host]
240
241     if not strict_host_checking:
242         # Do not check for Host key. Unsafe.
243         ssh_args.extend(['-o', 'StrictHostKeyChecking=no'])
244
245     if agent:
246         ssh_args.append('-A')
247
248     if port:
249         ssh_args.append('-p%d' % port)
250
251     if identity:
252         ssh_args.extend(('-i', identity))
253
254     if tty:
255         ssh_args.append('-t')
256         ssh_args.append('-t')
257
258     if server_key:
259         # Create a temporary server key file
260         tmp_known_hosts = make_server_key_args(server_key, host, port)
261         ssh_args.extend(['-o', 'UserKnownHostsFile=%s' % (tmp_known_hosts.name,)])
262
263     ssh_cmd = " ".join(ssh_args)
264
265     exec_cmd = "EXEC:'%s socat STDIO UNIX-CONNECT\:%s'" % (ssh_cmd, 
266             remote_socket_name)
267
268     args.append(exec_cmd)
269     
270     log_msg = " socat - host %s - command %s " % (host, " ".join(args))
271
272     return _retry_rexec(args, log_msg, 
273             stdout = None,
274             stdin = None,
275             stderr = None,
276             env = env, 
277             retry = retry, 
278             tmp_known_hosts = tmp_known_hosts,
279             blocking = False)
280
281 def rexec(command, host, user, 
282         port = None, 
283         agent = True,
284         sudo = False,
285         identity = None,
286         server_key = None,
287         env = None,
288         tty = False,
289         connect_timeout = 30,
290         retry = 3,
291         persistent = True,
292         forward_x11 = False,
293         blocking = True,
294         strict_host_checking = True):
295     """
296     Executes a remote command, returns ((stdout,stderr),process)
297     """
298     
299     tmp_known_hosts = None
300     hostip = gethostbyname(host)
301
302     args = ['ssh', '-C',
303             # Don't bother with localhost. Makes test easier
304             '-o', 'NoHostAuthenticationForLocalhost=yes',
305             '-o', 'ConnectTimeout=%d' % (int(connect_timeout),),
306             '-o', 'ConnectionAttempts=3',
307             '-o', 'ServerAliveInterval=30',
308             '-o', 'TCPKeepAlive=yes',
309             '-l', user, hostip or host]
310
311     if persistent and openssh_has_persist():
312         args.extend([
313             '-o', 'ControlMaster=auto',
314             '-o', 'ControlPath=%s' % (make_control_path(agent, forward_x11),),
315             '-o', 'ControlPersist=60' ])
316
317     if not strict_host_checking:
318         # Do not check for Host key. Unsafe.
319         args.extend(['-o', 'StrictHostKeyChecking=no'])
320
321     if agent:
322         args.append('-A')
323
324     if port:
325         args.append('-p%d' % port)
326
327     if identity:
328         args.extend(('-i', identity))
329
330     if tty:
331         args.append('-t')
332         args.append('-t')
333
334     if forward_x11:
335         args.append('-X')
336
337     if server_key:
338         # Create a temporary server key file
339         tmp_known_hosts = make_server_key_args(server_key, host, port)
340         args.extend(['-o', 'UserKnownHostsFile=%s' % (tmp_known_hosts.name,)])
341
342     if sudo:
343         command = "sudo " + command
344
345     args.append(command)
346     
347     log_msg = " rexec - host %s - command %s " % (host, " ".join(args))
348
349     stdout = stderr = stdin = subprocess.PIPE
350     if forward_x11:
351         stdout = stderr = stdin = None
352
353     return _retry_rexec(args, log_msg, 
354             stderr = stderr,
355             stdin = stdin,
356             stdout = stdout,
357             env = env, 
358             retry = retry, 
359             tmp_known_hosts = tmp_known_hosts,
360             blocking = blocking)
361
362 def rcopy(source, dest,
363         port = None, 
364         agent = True, 
365         recursive = False,
366         identity = None,
367         server_key = None,
368         retry = 3,
369         strict_host_checking = True):
370     """
371     Copies from/to remote sites.
372     
373     Source and destination should have the user and host encoded
374     as per scp specs.
375     
376     Source can be a list of files to copy to a single destination,
377     in which case it is advised that the destination be a folder.
378     """
379     
380     # Parse destination as <user>@<server>:<path>
381     if isinstance(dest, basestring) and ':' in dest:
382         remspec, path = dest.split(':',1)
383     elif isinstance(source, basestring) and ':' in source:
384         remspec, path = source.split(':',1)
385     else:
386         raise ValueError, "Both endpoints cannot be local"
387     user,host = remspec.rsplit('@',1)
388     
389     # plain scp
390     tmp_known_hosts = None
391
392     args = ['scp', '-q', '-p', '-C',
393             # Speed up transfer using blowfish cypher specification which is 
394             # faster than the default one (3des)
395             '-c', 'blowfish',
396             # Don't bother with localhost. Makes test easier
397             '-o', 'NoHostAuthenticationForLocalhost=yes',
398             '-o', 'ConnectTimeout=60',
399             '-o', 'ConnectionAttempts=3',
400             '-o', 'ServerAliveInterval=30',
401             '-o', 'TCPKeepAlive=yes' ]
402             
403     if port:
404         args.append('-P%d' % port)
405
406     if recursive:
407         args.append('-r')
408
409     if identity:
410         args.extend(('-i', identity))
411
412     if server_key:
413         # Create a temporary server key file
414         tmp_known_hosts = make_server_key_args(server_key, host, port)
415         args.extend(['-o', 'UserKnownHostsFile=%s' % (tmp_known_hosts.name,)])
416
417     if not strict_host_checking:
418         # Do not check for Host key. Unsafe.
419         args.extend(['-o', 'StrictHostKeyChecking=no'])
420
421     if isinstance(source, list):
422         args.extend(source)
423     else:
424         if openssh_has_persist():
425             args.extend([
426                 '-o', 'ControlMaster=auto',
427                 '-o', 'ControlPath=%s' % (make_control_path(agent, False),)
428                 ])
429         args.append(source)
430
431     args.append(dest)
432
433     log_msg = " rcopy - host %s - command %s " % (host, " ".join(args))
434     
435     return _retry_rexec(args, log_msg, env = None, retry = retry, 
436             tmp_known_hosts = tmp_known_hosts,
437             blocking = True)
438
439 def rspawn(command, pidfile, 
440         stdout = '/dev/null', 
441         stderr = STDOUT, 
442         stdin = '/dev/null',
443         home = None, 
444         create_home = False, 
445         sudo = False,
446         host = None, 
447         port = None, 
448         user = None, 
449         agent = None, 
450         identity = None, 
451         server_key = None,
452         tty = False):
453     """
454     Spawn a remote command such that it will continue working asynchronously in 
455     background. 
456
457         :param command: The command to run, it should be a single line.
458         :type command: str
459
460         :param pidfile: Path to a file where to store the pid and ppid of the 
461                         spawned process
462         :type pidfile: str
463
464         :param stdout: Path to file to redirect standard output. 
465                        The default value is /dev/null
466         :type stdout: str
467
468         :param stderr: Path to file to redirect standard error.
469                        If the special STDOUT value is used, stderr will 
470                        be redirected to the same file as stdout
471         :type stderr: str
472
473         :param stdin: Path to a file with input to be piped into the command's standard input
474         :type stdin: str
475
476         :param home: Path to working directory folder. 
477                     It is assumed to exist unless the create_home flag is set.
478         :type home: str
479
480         :param create_home: Flag to force creation of the home folder before 
481                             running the command
482         :type create_home: bool
483  
484         :param sudo: Flag forcing execution with sudo user
485         :type sudo: bool
486         
487         :rtype: touple
488
489         (stdout, stderr), process
490         
491         Of the spawning process, which only captures errors at spawning time.
492         Usually only useful for diagnostics.
493     """
494     # Start process in a "daemonized" way, using nohup and heavy
495     # stdin/out redirection to avoid connection issues
496     if stderr is STDOUT:
497         stderr = '&1'
498     else:
499         stderr = ' ' + stderr
500     
501     daemon_command = '{ { %(command)s > %(stdout)s 2>%(stderr)s < %(stdin)s & } ; echo $! 1 > %(pidfile)s ; }' % {
502         'command' : command,
503         'pidfile' : shell_escape(pidfile),
504         'stdout' : stdout,
505         'stderr' : stderr,
506         'stdin' : stdin,
507     }
508     
509     cmd = "%(create)s%(gohome)s rm -f %(pidfile)s ; %(sudo)s nohup bash -c %(command)s " % {
510             'command' : shell_escape(daemon_command),
511             'sudo' : 'sudo -S' if sudo else '',
512             'pidfile' : shell_escape(pidfile),
513             'gohome' : 'cd %s ; ' % (shell_escape(home),) if home else '',
514             'create' : 'mkdir -p %s ; ' % (shell_escape(home),) if create_home and home else '',
515         }
516
517     (out,err),proc = rexec(
518         cmd,
519         host = host,
520         port = port,
521         user = user,
522         agent = agent,
523         identity = identity,
524         server_key = server_key,
525         tty = tty ,
526         )
527     
528     if proc.wait():
529         raise RuntimeError, "Failed to set up application on host %s: %s %s" % (host, out,err,)
530
531     return ((out, err), proc)
532
533 @eintr_retry
534 def rgetpid(pidfile,
535         host = None, 
536         port = None, 
537         user = None, 
538         agent = None, 
539         identity = None,
540         server_key = None):
541     """
542     Returns the pid and ppid of a process from a remote file where the 
543     information was stored.
544
545         :param home: Path to directory where the pidfile is located
546         :type home: str
547
548         :param pidfile: Name of file containing the pid information
549         :type pidfile: str
550         
551         :rtype: int
552         
553         A (pid, ppid) tuple useful for calling rstatus and rkill,
554         or None if the pidfile isn't valid yet (can happen when process is staring up)
555
556     """
557     (out,err),proc = rexec(
558         "cat %(pidfile)s" % {
559             'pidfile' : pidfile,
560         },
561         host = host,
562         port = port,
563         user = user,
564         agent = agent,
565         identity = identity,
566         server_key = server_key
567         )
568         
569     if proc.wait():
570         return None
571     
572     if out:
573         try:
574             return map(int,out.strip().split(' ',1))
575         except:
576             # Ignore, many ways to fail that don't matter that much
577             return None
578
579 @eintr_retry
580 def rstatus(pid, ppid, 
581         host = None, 
582         port = None, 
583         user = None, 
584         agent = None, 
585         identity = None,
586         server_key = None):
587     """
588     Returns a code representing the the status of a remote process
589
590         :param pid: Process id of the process
591         :type pid: int
592
593         :param ppid: Parent process id of process
594         :type ppid: int
595     
596         :rtype: int (One of NOT_STARTED, RUNNING, FINISHED)
597     
598     """
599     (out,err),proc = rexec(
600         # Check only by pid. pid+ppid does not always work (especially with sudo) 
601         " (( ps --pid %(pid)d -o pid | grep -c %(pid)d && echo 'wait')  || echo 'done' ) | tail -n 1" % {
602             'ppid' : ppid,
603             'pid' : pid,
604         },
605         host = host,
606         port = port,
607         user = user,
608         agent = agent,
609         identity = identity,
610         server_key = server_key
611         )
612     
613     if proc.wait():
614         return ProcStatus.NOT_STARTED
615     
616     status = False
617     if err:
618         if err.strip().find("Error, do this: mount -t proc none /proc") >= 0:
619             status = True
620     elif out:
621         status = (out.strip() == 'wait')
622     else:
623         return ProcStatus.NOT_STARTED
624     return ProcStatus.RUNNING if status else ProcStatus.FINISHED
625
626 @eintr_retry
627 def rkill(pid, ppid,
628         host = None, 
629         port = None, 
630         user = None, 
631         agent = None, 
632         sudo = False,
633         identity = None, 
634         server_key = None, 
635         nowait = False):
636     """
637     Sends a kill signal to a remote process.
638
639     First tries a SIGTERM, and if the process does not end in 10 seconds,
640     it sends a SIGKILL.
641  
642         :param pid: Process id of process to be killed
643         :type pid: int
644
645         :param ppid: Parent process id of process to be killed
646         :type ppid: int
647
648         :param sudo: Flag indicating if sudo should be used to kill the process
649         :type sudo: bool
650         
651     """
652     subkill = "$(ps --ppid %(pid)d -o pid h)" % { 'pid' : pid }
653     cmd = """
654 SUBKILL="%(subkill)s" ;
655 %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
656 %(sudo)s kill %(pid)d $SUBKILL || /bin/true
657 for x in 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 ; do 
658     sleep 0.2 
659     if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` == '0' ]; then
660         break
661     else
662         %(sudo)s kill -- -%(pid)d $SUBKILL || /bin/true
663         %(sudo)s kill %(pid)d $SUBKILL || /bin/true
664     fi
665     sleep 1.8
666 done
667 if [ `ps --pid %(pid)d -o pid | grep -c %(pid)d` != '0' ]; then
668     %(sudo)s kill -9 -- -%(pid)d $SUBKILL || /bin/true
669     %(sudo)s kill -9 %(pid)d $SUBKILL || /bin/true
670 fi
671 """
672     if nowait:
673         cmd = "( %s ) >/dev/null 2>/dev/null </dev/null &" % (cmd,)
674
675     (out,err),proc = rexec(
676         cmd % {
677             'ppid' : ppid,
678             'pid' : pid,
679             'sudo' : 'sudo -S' if sudo else '',
680             'subkill' : subkill,
681         },
682         host = host,
683         port = port,
684         user = user,
685         agent = agent,
686         identity = identity,
687         server_key = server_key
688         )
689     
690     # wait, don't leave zombies around
691     proc.wait()
692
693     return (out, err), proc
694
695 def _retry_rexec(args,
696         log_msg,
697         stdout = subprocess.PIPE,
698         stdin = subprocess.PIPE, 
699         stderr = subprocess.PIPE,
700         env = None,
701         retry = 3,
702         tmp_known_hosts = None,
703         blocking = True):
704
705     for x in xrange(retry):
706         # connects to the remote host and starts a remote connection
707         proc = subprocess.Popen(args,
708                 env = env,
709                 stdout = stdout,
710                 stdin = stdin, 
711                 stderr = stderr)
712         
713         # attach tempfile object to the process, to make sure the file stays
714         # alive until the process is finished with it
715         proc._known_hosts = tmp_known_hosts
716     
717         # The argument block == False forces to rexec to return immediately, 
718         # without blocking 
719         try:
720             err = out = " "
721             if blocking:
722                 (out, err) = proc.communicate()
723             elif stdout:
724                 out = proc.stdout.read()
725                 if proc.poll() and stderr:
726                     err = proc.stderr.read()
727
728             log(log_msg, logging.DEBUG, out, err)
729
730             if proc.poll():
731                 skip = False
732
733                 if err.strip().startswith('ssh: ') or err.strip().startswith('mux_client_hello_exchange: '):
734                     # SSH error, can safely retry
735                     skip = True 
736                 elif retry:
737                     # Probably timed out or plain failed but can retry
738                     skip = True 
739                 
740                 if skip:
741                     t = x*2
742                     msg = "SLEEPING %d ... ATEMPT %d - command %s " % ( 
743                             t, x, " ".join(args))
744                     log(msg, logging.DEBUG)
745
746                     time.sleep(t)
747                     continue
748             break
749         except RuntimeError, e:
750             msg = " rexec EXCEPTION - TIMEOUT -> %s \n %s" % ( e.args, log_msg )
751             log(msg, logging.DEBUG, out, err)
752
753             if retry <= 0:
754                 raise
755             retry -= 1
756         
757     return ((out, err), proc)
758
759