1
# Copyright 2012 Canonical Ltd.
3
# This file is part of python-shell-toolbox.
5
# python-shell-toolbox is free software: you can redistribute it and/or modify
6
# it under the terms of the GNU General Public License as published by the
7
# Free Software Foundation, version 3 of the License.
9
# python-shell-toolbox is distributed in the hope that it will be useful, but
10
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
11
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
14
# You should have received a copy of the GNU General Public License
15
# along with python-shell-toolbox. If not, see <http://www.gnu.org/licenses/>.
17
"""Helper functions for accessing shell commands in Python."""
33
'install_extra_repositories',
43
'wait_for_page_contents',
46
from collections import namedtuple
47
from contextlib import contextmanager
48
from email.Utils import parseaddr
58
from textwrap import dedent
63
Env = namedtuple('Env', 'uid gid home')
66
def apt_get_install(*args, **kwargs):
67
"""Install given packages using apt.
69
It is possible to pass environment variables to be set during install
70
using keyword arguments.
72
:raises: subprocess.CalledProcessError
74
caller = kwargs.pop('caller', run)
75
debian_frontend = kwargs.pop('DEBIAN_FRONTEND', 'noninteractive')
76
with environ(DEBIAN_FRONTEND=debian_frontend, **kwargs):
77
cmd = ('apt-get', '-y', 'install') + args
82
"""Return full name and email of bzr `user`.
84
Return None if the given `user` does not have a bzr user id.
88
whoami = run('bzr', 'whoami')
89
except (subprocess.CalledProcessError, OSError):
91
return parseaddr(whoami)
96
"""A context manager to temporarily change current working dir, e.g.::
100
>>> with cd('/bin'): print os.getcwd()
102
>>> print os.getcwd()
113
def command(*base_args):
114
"""Return a callable that will run the given command with any arguments.
116
The first argument is the path to the command to run, subsequent arguments
117
are command-line arguments to "bake into" the returned callable.
119
The callable runs the given executable and also takes arguments that will
120
be appeneded to the "baked in" arguments.
122
For example, this code will list a file named "foo" (if it exists):
124
ls_foo = command('/bin/ls', 'foo')
127
While this invocation will list "foo" and "bar" (assuming they exist):
131
def callable_command(*args):
132
all_args = base_args + args
133
return run(*all_args)
135
return callable_command
139
def environ(**kwargs):
140
"""A context manager to temporarily change environment variables.
142
If an existing environment variable is changed, it is restored during
146
>>> os.environ['MY_VARIABLE'] = 'foo'
147
>>> with environ(MY_VARIABLE='bar'): print os.getenv('MY_VARIABLE')
149
>>> print os.getenv('MY_VARIABLE')
151
>>> del os.environ['MY_VARIABLE']
153
If we are adding environment variables, they are removed during context
157
>>> with environ(MY_VAR1='foo', MY_VAR2='bar'):
158
... print os.getenv('MY_VAR1'), os.getenv('MY_VAR2')
160
>>> os.getenv('MY_VAR1') == os.getenv('MY_VAR2') == None
164
for key, value in kwargs.items():
165
backup[key] = os.getenv(key)
166
os.environ[key] = value
170
for key, value in backup.items():
174
os.environ[key] = value
177
def file_append(filename, line):
178
r"""Append given `line`, if not present, at the end of `filename`.
183
>>> f = tempfile.NamedTemporaryFile('w', delete=False)
184
>>> f.write('line1\n')
186
>>> file_append(f.name, 'new line\n')
187
>>> open(f.name).read()
190
Nothing happens if the file already contains the given `line`::
192
>>> file_append(f.name, 'new line\n')
193
>>> open(f.name).read()
196
A new line is automatically added before the given `line` if it is not
197
present at the end of current file content::
200
>>> f = tempfile.NamedTemporaryFile('w', delete=False)
203
>>> file_append(f.name, 'new line\n')
204
>>> open(f.name).read()
207
The file is created if it does not exist::
210
>>> filename = tempfile.mktemp()
211
>>> file_append(filename, 'line1\n')
212
>>> open(filename).read()
215
if not line.endswith('\n'):
217
with open(filename, 'a+') as f:
218
lines = f.readlines()
219
if line not in lines:
220
if not lines or lines[-1].endswith('\n'):
226
def file_prepend(filename, line):
227
r"""Insert given `line`, if not present, at the beginning of `filename`.
232
>>> f = tempfile.NamedTemporaryFile('w', delete=False)
233
>>> f.write('line1\n')
235
>>> file_prepend(f.name, 'line0\n')
236
>>> open(f.name).read()
239
If the file starts with the given `line`, nothing happens::
241
>>> file_prepend(f.name, 'line0\n')
242
>>> open(f.name).read()
245
If the file contains the given `line`, but not at the beginning,
246
the line is moved on top::
248
>>> file_prepend(f.name, 'line1\n')
249
>>> open(f.name).read()
252
if not line.endswith('\n'):
254
with open(filename, 'r+') as f:
255
lines = f.readlines()
261
lines.insert(0, line)
266
def generate_ssh_keys(path, passphrase=''):
267
"""Generate ssh key pair, saving them inside the given `directory`.
269
>>> generate_ssh_keys('/tmp/id_rsa')
271
>>> open('/tmp/id_rsa').readlines()[0].strip()
272
'-----BEGIN RSA PRIVATE KEY-----'
273
>>> open('/tmp/id_rsa.pub').read().startswith('ssh-rsa')
275
>>> os.remove('/tmp/id_rsa')
276
>>> os.remove('/tmp/id_rsa.pub')
278
If either of the key files already exist, generate_ssh_keys() will
281
Note that ssh-keygen will prompt if the keyfiles already exist, but
282
when we're using it non-interactively it's better to pre-empt that
285
>>> with open('/tmp/id_rsa', 'w') as key_file:
286
... key_file.write("Don't overwrite me, bro!")
287
>>> generate_ssh_keys('/tmp/id_rsa') # doctest: +ELLIPSIS
288
Traceback (most recent call last):
289
Exception: File /tmp/id_rsa already exists...
290
>>> os.remove('/tmp/id_rsa')
292
>>> with open('/tmp/id_rsa.pub', 'w') as key_file:
293
... key_file.write("Don't overwrite me, bro!")
294
>>> generate_ssh_keys('/tmp/id_rsa') # doctest: +ELLIPSIS
295
Traceback (most recent call last):
296
Exception: File /tmp/id_rsa.pub already exists...
297
>>> os.remove('/tmp/id_rsa.pub')
299
if os.path.exists(path):
300
raise Exception("File {} already exists.".format(path))
301
if os.path.exists(path + '.pub'):
302
raise Exception("File {}.pub already exists.".format(path))
303
return subprocess.call([
304
'ssh-keygen', '-q', '-t', 'rsa', '-N', passphrase, '-f', path])
307
def get_su_command(user, args):
308
"""Return a command line as a sequence, prepending "su" if necessary.
310
This can be used together with `run` when the `su` context manager is not
311
enough (e.g. an external program uses uid rather than euid).
313
run(*get_su_command(user, ['bzr', 'whoami']))
315
If the su is requested as current user, the arguments are returned as
319
>>> current_user = getpass.getuser()
321
>>> get_su_command(current_user, ('ls', '-l'))
324
Otherwise, "su" is prepended::
326
>>> get_su_command('nobody', ('ls', '-l', 'my file'))
327
('su', 'nobody', '-c', "ls -l 'my file'")
329
if get_user_ids(user)[0] != os.getuid():
330
args = [i for i in args if i is not None]
331
return ('su', user, '-c', join_command(args))
335
def get_user_home(user):
336
"""Return the home directory of the given `user`.
338
>>> get_user_home('root')
341
If the user does not exist, return a default /home/[username] home::
343
>>> get_user_home('_this_user_does_not_exist_')
344
'/home/_this_user_does_not_exist_'
347
return pwd.getpwnam(user).pw_dir
349
return os.path.join(os.path.sep, 'home', user)
352
def get_user_ids(user):
353
"""Return the uid and gid of given `user`, e.g.::
355
>>> get_user_ids('root')
358
userdata = pwd.getpwnam(user)
359
return userdata.pw_uid, userdata.pw_gid
362
def install_extra_repositories(*repositories):
363
"""Install all of the extra repositories and update apt.
365
Given repositories can contain a "{distribution}" placeholder, that will
366
be replaced by current distribution codename.
368
:raises: subprocess.CalledProcessError
370
distribution = run('lsb_release', '-cs').strip()
371
# Starting from Oneiric, `apt-add-repository` is interactive by
372
# default, and requires a "-y" flag to be set.
373
assume_yes = None if distribution == 'lucid' else '-y'
374
for repo in repositories:
375
repository = repo.format(distribution=distribution)
376
run('apt-add-repository', assume_yes, repository)
377
run('apt-get', 'clean')
378
run('apt-get', 'update')
381
def join_command(args):
382
"""Return a valid Unix command line from `args`.
384
>>> join_command(['ls', '-l'])
387
Arguments containing spaces and empty args are correctly quoted::
389
>>> join_command(['command', 'arg1', 'arg containing spaces', ''])
390
"command arg1 'arg containing spaces' ''"
392
return ' '.join(pipes.quote(arg) for arg in args)
396
"""Create leaf directories (given as `args`) and all intermediate ones.
399
>>> base_dir = tempfile.mktemp(suffix='/')
400
>>> dir1 = tempfile.mktemp(prefix=base_dir)
401
>>> dir2 = tempfile.mktemp(prefix=base_dir)
402
>>> mkdirs(dir1, dir2)
403
>>> os.path.isdir(dir1)
405
>>> os.path.isdir(dir2)
408
If the leaf directory already exists the function returns without errors::
412
An `OSError` is raised if the leaf path exists and it is a file::
414
>>> f = tempfile.NamedTemporaryFile(
415
... 'w', delete=False, prefix=base_dir)
417
>>> mkdirs(f.name) # doctest: +ELLIPSIS
418
Traceback (most recent call last):
421
for directory in args:
423
os.makedirs(directory)
424
except OSError as err:
425
if err.errno != errno.EEXIST or os.path.isfile(directory):
429
def run(*args, **kwargs):
430
"""Run the command with the given arguments.
432
The first argument is the path to the command to run.
433
Subsequent arguments are command-line arguments to be passed.
435
This function accepts all optional keyword arguments accepted by
438
args = [i for i in args if i is not None]
439
pipe = subprocess.PIPE
440
process = subprocess.Popen(
441
args, stdout=kwargs.pop('stdout', pipe),
442
stderr=kwargs.pop('stderr', pipe),
443
close_fds=kwargs.pop('close_fds', True), **kwargs)
444
stdout, stderr = process.communicate()
445
if process.returncode:
446
exception = subprocess.CalledProcessError(
447
process.returncode, repr(args))
448
# The output argument of `CalledProcessError` was introduced in Python
449
# 2.7. Monkey patch the output here to avoid TypeErrors in older
450
# versions of Python, still preserving the output in Python 2.7.
451
exception.output = ''.join(filter(None, [stdout, stderr]))
457
"""Return the name of this script."""
458
return os.path.basename(sys.argv[0])
461
def search_file(regexp, filename):
462
"""Return the first line in `filename` that matches `regexp`."""
463
with open(filename) as f:
465
if re.search(regexp, line):
469
def ssh(location, user=None, key=None, caller=subprocess.call):
470
"""Return a callable that can be used to run ssh shell commands.
472
The ssh `location` and, optionally, `user` must be given.
473
If the user is None then the current user is used for the connection.
475
The callable internally uses the given `caller`::
479
>>> sshcall = ssh('example.com', 'myuser', caller=caller)
480
>>> root_sshcall = ssh('example.com', caller=caller)
481
>>> sshcall('ls -l') # doctest: +ELLIPSIS
482
('ssh', '-t', ..., 'myuser@example.com', '--', 'ls -l')
483
>>> root_sshcall('ls -l') # doctest: +ELLIPSIS
484
('ssh', '-t', ..., 'example.com', '--', 'ls -l')
486
The ssh key path can be optionally provided::
488
>>> root_sshcall = ssh('example.com', key='/tmp/foo', caller=caller)
489
>>> root_sshcall('ls -l') # doctest: +ELLIPSIS
490
('ssh', '-t', ..., '-i', '/tmp/foo', 'example.com', '--', 'ls -l')
492
If the ssh command exits with an error code,
493
a `subprocess.CalledProcessError` is raised::
495
>>> ssh('loc', caller=lambda cmd: 1)('ls -l') # doctest: +ELLIPSIS
496
Traceback (most recent call last):
497
CalledProcessError: ...
499
If ignore_errors is set to True when executing the command, no error
500
will be raised, even if the command itself returns an error code.
502
>>> sshcall = ssh('loc', caller=lambda cmd: 1)
503
>>> sshcall('ls -l', ignore_errors=True)
508
'-t', # Yes, this second -t is deliberate. See `man ssh`.
509
'-o', 'StrictHostKeyChecking=no',
510
'-o', 'UserKnownHostsFile=/dev/null',
513
sshcmd.extend(['-i', key])
515
location = '{}@{}'.format(user, location)
516
sshcmd.extend([location, '--'])
518
def _sshcall(cmd, ignore_errors=False):
519
command = sshcmd + [cmd]
520
retcode = caller(command)
521
if retcode and not ignore_errors:
522
raise subprocess.CalledProcessError(retcode, ' '.join(command))
529
"""A context manager to temporarily run the script as a different user."""
530
uid, gid = get_user_ids(user)
533
home = get_user_home(user)
534
with environ(HOME=home):
536
yield Env(uid, gid, home)
538
os.setegid(os.getgid())
539
os.seteuid(os.getuid())
542
def user_exists(username):
543
"""Return True if given `username` exists, e.g.::
545
>>> user_exists('root')
547
>>> user_exists('_this_user_does_not_exist_')
551
pwd.getpwnam(username)
557
def wait_for_page_contents(url, contents, timeout=120, validate=None):
559
validate = operator.contains
560
start_time = time.time()
563
stream = urllib2.urlopen(url)
564
except (urllib2.HTTPError, urllib2.URLError):
568
if validate(page, contents):
570
if time.time() - start_time >= timeout:
571
raise RuntimeError('timeout waiting for contents of ' + url)
577
Calculate the difference between two dictionaries as:
580
(3) keys same in both but changed values
581
(4) keys same in both and unchanged values
584
# Based on answer by hughdbrown at:
585
# http://stackoverflow.com/questions/1165352
587
def __init__(self, current_dict, past_dict):
588
self.current_dict = current_dict
589
self.past_dict = past_dict
590
self.set_current = set(current_dict)
591
self.set_past = set(past_dict)
592
self.intersect = self.set_current.intersection(self.set_past)
596
return self.set_current - self.intersect
600
return self.set_past - self.intersect
604
return set(key for key in self.intersect
605
if self.past_dict[key] != self.current_dict[key])
608
return set(key for key in self.intersect
609
if self.past_dict[key] == self.current_dict[key])
612
return self.current_dict != self.past_dict
615
def added_or_changed(self):
616
return self.added.union(self.changed)
618
def _changes(self, keys):
622
new[k] = self.current_dict.get(k)
623
old[k] = self.past_dict.get(k)
624
return "%s -> %s" % (old, new)
632
unchanged: %s""") % (
633
self._changes(self.added),
634
self._changes(self.removed),
635
self._changes(self.changed),
636
list(self.unchanged))
643
"""Handle JSON (de)serialization."""
645
def __init__(self, path, default=None, serialize=None, deserialize=None):
647
self.default = default or {}
648
self.serialize = serialize or json.dump
649
self.deserialize = deserialize or json.load
652
return os.path.exists(self.path)
656
with open(self.path) as f:
657
return self.deserialize(f)
661
with open(self.path, 'w') as f:
662
self.serialize(data, f)