3 # NEPI, a framework to manage network experiments
4 # Copyright (C) 2013 INRIA
6 # This program is free software: you can redistribute it and/or modify
7 # it under the terms of the GNU General Public License as published by
8 # the Free Software Foundation, either version 3 of the License, or
9 # (at your option) any later version.
11 # This program is distributed in the hope that it will be useful,
12 # but WITHOUT ANY WARRANTY; without even the implied warranty of
13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14 # GNU General Public License for more details.
16 # You should have received a copy of the GNU General Public License
17 # along with this program. If not, see <http://www.gnu.org/licenses/>.
19 # Author: Alina Quereilhac <alina.quereilhac@inria.fr>
22 from nepi.util.sshfuncs import rexec, rcopy, rspawn, rgetpid, rstatus, rkill,\
37 def find_bin(name, extra_path = None):
39 if "PATH" in os.environ:
40 search += os.environ["PATH"].split(":")
41 for pref in ("/", "/usr/", "/usr/local/"):
42 for d in ("bin", "sbin"):
43 search.append(pref + d)
49 os.stat(d + "/" + name)
52 if e.errno != os.errno.ENOENT:
56 def find_bin_or_die(name, extra_path = None):
59 raise RuntimeError(("Cannot find `%s' command, impossible to " +
63 def gen_ssh_keypair(filename):
64 ssh_keygen = find_bin_or_die("ssh-keygen")
65 args = [ssh_keygen, '-q', '-N', '', '-f', filename]
66 assert subprocess.Popen(args).wait() == 0
67 return filename, "%s.pub" % filename
69 def add_key_to_agent(filename):
70 ssh_add = find_bin_or_die("ssh-add")
71 args = [ssh_add, filename]
72 null = file("/dev/null", "w")
73 assert subprocess.Popen(args, stderr = null).wait() == 0
78 s.bind(("127.0.0.1", 0))
79 port = s.getsockname()[1]
82 _SSH_CONF = """ListenAddress 127.0.0.1:%d
85 UsePrivilegeSeparation no
86 PubkeyAuthentication yes
87 PasswordAuthentication no
90 AllowAgentForwarding yes
93 PermitUserEnvironment yes
96 def gen_sshd_config(filename, port, server_key, auth_keys):
97 conf = open(filename, "w")
98 text = _SSH_CONF % (port, server_key, auth_keys)
103 def gen_auth_keys(pubkey, output, environ):
104 #opts = ['from="127.0.0.1/32"'] # fails in stupid yans setup
106 for k, v in environ.items():
107 opts.append('environment="%s=%s"' % (k, v))
109 lines = file(pubkey).readlines()
110 pubkey = lines[0].split()[0:2]
111 out = file(output, "w")
112 out.write("%s %s %s\n" % (",".join(opts), pubkey[0], pubkey[1]))
116 def start_ssh_agent():
117 ssh_agent = find_bin_or_die("ssh-agent")
118 proc = subprocess.Popen([ssh_agent], stdout = subprocess.PIPE)
119 (out, foo) = proc.communicate()
120 assert proc.returncode == 0
122 for l in out.split("\n"):
123 match = re.search("^(\w+)=([^ ;]+);.*", l)
126 k, v = match.groups()
131 def stop_ssh_agent(data):
132 # No need to gather the pid, ssh-agent knows how to kill itself; after we
133 # had set up the environment
134 ssh_agent = find_bin_or_die("ssh-agent")
135 null = file("/dev/null", "w")
136 proc = subprocess.Popen([ssh_agent, "-k"], stdout = null)
138 assert proc.wait() == 0
142 class test_environment(object):
144 sshd = find_bin_or_die("sshd")
146 self.dir = tempfile.mkdtemp()
147 self.server_keypair = gen_ssh_keypair(
148 os.path.join(self.dir, "server_key"))
149 self.client_keypair = gen_ssh_keypair(
150 os.path.join(self.dir, "client_key"))
151 self.authorized_keys = gen_auth_keys(self.client_keypair[1],
152 os.path.join(self.dir, "authorized_keys"), environ)
153 self.port = get_free_port()
154 self.sshd_conf = gen_sshd_config(
155 os.path.join(self.dir, "sshd_config"),
156 self.port, self.server_keypair[0], self.authorized_keys)
158 self.sshd = subprocess.Popen([sshd, '-q', '-D', '-f', self.sshd_conf])
159 self.ssh_agent_vars = start_ssh_agent()
160 add_key_to_agent(self.client_keypair[0])
164 os.kill(self.sshd.pid, signal.SIGTERM)
166 if self.ssh_agent_vars:
167 stop_ssh_agent(self.ssh_agent_vars)
168 shutil.rmtree(self.dir)
170 class SSHfuncsTestCase(unittest.TestCase):
171 def test_rexec(self):
172 env = test_environment()
173 user = getpass.getuser()
178 plocal = subprocess.Popen(command, stdout=subprocess.PIPE,
179 stdin=subprocess.PIPE)
180 outlocal, errlocal = plocal.communicate()
182 (outremote, errrmote), premote = rexec(command, host, user,
183 port = env.port, agent = True)
185 self.assertEquals(outlocal, outremote)
187 def test_rcopy_list(self):
188 env = test_environment()
189 user = getpass.getuser()
192 # create some temp files and directories to copy
193 dirpath = tempfile.mkdtemp()
194 f = tempfile.NamedTemporaryFile(dir=dirpath, delete=False)
197 f1 = tempfile.NamedTemporaryFile(delete=False)
201 source = [dirpath, f1.name]
202 destdir = tempfile.mkdtemp()
203 dest = "%s@%s:%s" % (user, host, destdir)
204 rcopy(source, dest, port = env.port, agent = True, recursive = True)
207 def recls(files, dirname, names):
209 os.path.walk(destdir, recls, files)
211 origfiles = map(lambda s: os.path.basename(s), [dirpath, f.name, f1.name])
213 self.assertEquals(sorted(origfiles), sorted(files))
216 shutil.rmtree(dirpath)
218 def test_rcopy_list(self):
219 env = test_environment()
220 user = getpass.getuser()
223 # create some temp files and directories to copy
224 dirpath = tempfile.mkdtemp()
225 f = tempfile.NamedTemporaryFile(dir=dirpath, delete=False)
228 f1 = tempfile.NamedTemporaryFile(delete=False)
232 # Copy a list of files
233 source = [dirpath, f1.name]
234 destdir = tempfile.mkdtemp()
235 dest = "%s@%s:%s" % (user, host, destdir)
236 ((out, err), proc) = rcopy(source, dest, port = env.port, agent = True, recursive = True)
239 def recls(files, dirname, names):
241 os.path.walk(destdir, recls, files)
243 origfiles = map(lambda s: os.path.basename(s), [dirpath, f.name, f1.name])
245 self.assertEquals(sorted(origfiles), sorted(files))
247 def test_rproc_manage(self):
248 env = test_environment()
249 user = getpass.getuser()
251 command = "ping localhost"
253 f = tempfile.NamedTemporaryFile(delete=False)
256 (out,err), proc = rspawn(
266 (pid, ppid) = rgetpid(pidfile,
272 status = rstatus(pid, ppid,
278 self.assertEquals(status, ProcStatus.RUNNING)
286 status = rstatus(pid, ppid,
292 self.assertEquals(status, ProcStatus.FINISHED)
295 if __name__ == '__main__':