1#!/usr/bin/env python
2
3import contextlib
4import hashlib
5import os
6import os.path
7import shutil
8import stat
9import subprocess
10import sys
11import xattr
12
13class TestCaseSkipError(Exception):
14	pass
15
16def skip_if_no_compression_support(type):
17	"""
18	Raises TestCaseSkipError if the type is "lzma" and the test is running on
19	darwin (OS X). In the future, we should add a hidden debugging flag to xar
20	to determine valid compression types. This will skip incorrectly if a
21	custom xar is used on OS X, or if a custom xar on another platform is
22	built without bzip2 or lzma.
23
24	"""
25	if sys.platform == "darwin" and type == "lzma":
26		raise TestCaseSkipError("{t} support not compiled in".format(t=type))
27
28@contextlib.contextmanager
29def directory_created(directory_path):
30	"""
31	Creates the named directory and provides the path to the directory to the
32	calling code. Automatically removes the directory when finished.
33
34	Usage:
35	with directory_created("foobar") as path:
36		do_stuff_with_path
37
38	"""
39	os.mkdir(directory_path)
40	try:
41		yield os.path.realpath(directory_path)
42	finally:
43		if os.path.exists(directory_path):
44			shutil.rmtree(directory_path)
45
46@contextlib.contextmanager
47def archive_created(archive_path, content_path, *extra_args, **extra_kwargs):
48	"""
49	Creates a named xar archive of the specified content path, returning the
50	path to the archive. Automatically removes the archive when finished.
51
52	Usage:
53	with archive_created("/bin", "bin.xar") as path:
54		do_stuff_with(path)
55
56	"""
57	cmd = ["xar", "-c", "-f", archive_path, content_path]
58	if extra_args:
59		cmd += list(extra_args)
60	try:
61		subprocess.check_call(cmd, **extra_kwargs)
62		assert os.path.exists(archive_path), "failed to create archive \"{p}\" but xar did not report an error".format(p=archive_path)
63		yield os.path.realpath(archive_path)
64	finally:
65		if os.path.exists(archive_path):
66			os.unlink(archive_path)
67
68HASH_CHUNK_SIZE = 32768
69
70def _md5_path(path):
71	with open(path, "r") as f:
72		h = hashlib.md5()
73		while True:
74			last = f.read(HASH_CHUNK_SIZE)
75			if not last:
76				break
77			h.update(last)
78		return h.digest()
79
80def assert_identical_directories(path1, path2):
81	"""
82	Verifies two directories have identical contents. Checks file type (via
83	the high byte of the mode), size, atime, and mtime, but does not check
84	other attributes like uid and gid, since they can be expected to change.
85
86	"""
87	seen = set([])
88	for file1 in os.listdir(path1):
89		seen.add(file1)
90		entry1 = os.path.join(path1, file1)
91		entry2 = os.path.join(path2, file1)
92		assert os.path.exists(entry2), "\"{f1}\" exists in \"{p1}\" but not \"{p2}\"".format(f1=file1, p1=path1, p2=path2)
93
94		# Extended attributes
95		xattr1 = xattr.xattr(entry1)
96		xattr2 = xattr.xattr(entry2)
97		assert set(xattr1.list()) == set(xattr2.list()), "list of extended attributes on \"{f1}\" ({l1}) differs from \"{f2}\" ({l2})".format(f1=entry1, l1=xattr1.list(), f2=entry2, l2=xattr2.list())
98		for attribute in xattr1.list():
99			assert xattr1.get(attribute) == xattr2.get(attribute), "extended attribute \"{a1}\" on \"{f1}\" doesn't match value from \"{f2}\"".format(a1=attribute, f1=entry1, f2=entry2)
100
101		# Why do it this way? We want to lstat() instead of stat(), so we can't use os.path.isdir() and friends
102		stat1 = os.lstat(entry1)
103		stat2 = os.lstat(entry2)
104
105		# Modes
106		mode1 = stat1.st_mode
107		mode2 = stat2.st_mode
108		if stat.S_ISREG(mode1):
109			assert stat.S_ISREG(mode2)
110		if stat.S_ISDIR(mode1):
111			assert stat.S_ISDIR(mode2)
112		if stat.S_ISLNK(mode1):
113			assert stat.S_ISLNK(mode2)
114		if stat.S_ISCHR(mode1):
115			assert stat.S_ISCHR(mode2)
116		if stat.S_ISBLK(mode1):
117			assert stat.S_ISBLK(mode2)
118		if stat.S_ISFIFO(mode1):
119			assert stat.S_ISFIFO(mode2)
120		if stat.S_ISSOCK(mode1):
121			assert stat.S_ISSOCK(mode2)
122
123		# Sizes and the like
124		assert stat1.st_size == stat2.st_size, "size mismatch for \"{e1}\" ({s1}) and \"{e2}\" ({s2})".format(e1=entry1, s1=stat1.st_size, e2=entry2, s2=stat2.st_size)
125		assert stat1.st_mtime == stat2.st_mtime, "mtime mismatch for \"{e1}\" and \"{e2}\"".format(e1=entry1, e2=entry2)
126		assert _md5_path(entry1) == _md5_path(entry2), "md5 hash mismatch for \"{e1}\" and \"{e2}\"".format(e1=entry1, e2=entry2)
127		if os.path.isdir(entry1):
128			assert_identical_directories(entry1, entry2)
129	for file2 in os.listdir(path2):
130		assert file2 in seen, "\"{f2}\" exists in \"{p2}\" but not \"{p1}\"".format(f2=file2, p1=path1, p2=path2)
131
132def touch(path):
133	if not os.path.exists(path):
134		with open(path, "w"):
135			pass
136	os.utime(path, None)
137
138@contextlib.contextmanager
139def chdir(*args, **kwargs):
140	cwd = os.getcwd()
141	os.chdir(*args, **kwargs)
142	try:
143		yield os.getcwd()
144	finally:
145		os.chdir(cwd)
146