importpkg: refactor commit handling out of process_package*
[~helmut/debian-dedup.git] / importpkg.py
1 #!/usr/bin/python
2 """This tool reads a Debian package from stdin and emits a yaml stream on
3 stdout.  It does not access a database. Therefore it can be run in parallel and
4 on multiple machines. The generated yaml contains multiple documents. The first
5 document contains package metadata. Then a document is emitted for each file.
6 And finally a document consisting of the string "commit" is emitted."""
7
8 import hashlib
9 import itertools
10 import optparse
11 import sys
12 import tarfile
13 import zlib
14
15 import lzma
16 import yaml
17
18 from dedup.arreader import ArReader
19 from dedup.debpkg import process_control, get_tar_hashes
20 from dedup.hashing import DecompressedHash, SuppressingHash, HashedStream, \
21         HashBlacklistContent
22 from dedup.compression import GzipDecompressor, DecompressedStream
23 from dedup.image import GIFHash, PNGHash
24
25 boring_content = set(("", "\n"))
26
27 def sha512_nontrivial():
28     return HashBlacklistContent(hashlib.sha512(), boring_content)
29
30 def gziphash():
31     hashobj = DecompressedHash(GzipDecompressor(), hashlib.sha512())
32     hashobj = SuppressingHash(hashobj, (ValueError, zlib.error))
33     hashobj.name = "gzip_sha512"
34     return HashBlacklistContent(hashobj, boring_content)
35
36 def pnghash():
37     hashobj = PNGHash(hashlib.sha512())
38     hashobj = SuppressingHash(hashobj, (ValueError,))
39     hashobj.name = "png_sha512"
40     return hashobj
41
42 def gifhash():
43     hashobj = GIFHash(hashlib.sha512())
44     hashobj = SuppressingHash(hashobj, (ValueError,))
45     hashobj.name = "gif_sha512"
46     return hashobj
47
48 def decompress_tar(filelike, extension):
49     if extension in (b".lzma", b".xz"):
50         filelike = DecompressedStream(filelike, lzma.LZMADecompressor())
51         extension = b""
52     if extension not in (b"", b".gz", b".bz2"):
53         raise ValueError("unknown compression format with extension %r" %
54                          extension)
55     return tarfile.open(fileobj=filelike,
56                         mode="r|" + extension[1:].decode("ascii"))
57
58 def process_package(filelike, hash_functions):
59     af = ArReader(filelike)
60     af.read_magic()
61     state = "start"
62     while True:
63         try:
64             name = af.read_entry()
65         except EOFError:
66             raise ValueError("data.tar not found")
67         if name.startswith(b"control.tar"):
68             if state != "start":
69                 raise ValueError("unexpected control.tar")
70             state = "control"
71             tf = decompress_tar(af, name[11:])
72             for elem in tf:
73                 if elem.name not in ("./control", "control"):
74                     continue
75                 if state != "control":
76                     raise ValueError("duplicate control file")
77                 state = "control_file"
78                 yield process_control(tf.extractfile(elem).read())
79                 break
80             continue
81         elif name.startswith(b"data.tar"):
82             if state != "control_file":
83                 raise ValueError("missing control file")
84             state = "data"
85             tf = decompress_tar(af, name[8:])
86             for name, size, hashes in get_tar_hashes(tf, hash_functions):
87                 try:
88                     name = name.decode("utf8")
89                 except UnicodeDecodeError:
90                     print("warning: skipping filename with encoding error")
91                     continue # skip files with non-utf8 encoding for now
92                 yield dict(name=name, size=size, hashes=hashes)
93             break
94
95 def hashed_stream_check(hstream, hashvalue):
96     if False: # pylint: disable=using-constant-test
97         yield # defer checking until being iterated
98     while hstream.read(4096):
99         pass
100     if hstream.hexdigest() != hashvalue:
101         raise ValueError("hash sum mismatch")
102
103 def main():
104     parser = optparse.OptionParser()
105     parser.add_option("-H", "--hash", action="store",
106                       help="verify that stdin hash given sha256 hash")
107     options, args = parser.parse_args()
108     hash_functions = [sha512_nontrivial, gziphash, pnghash, gifhash]
109     try:
110         stdin = sys.stdin.buffer
111     except AttributeError: # python2
112         stdin = sys.stdin
113     iters = [("commit",)]
114     if options.hash:
115         stdin = HashedStream(stdin, hashlib.sha256())
116         iters.insert(0, hashed_stream_check(stdin, options.hash))
117     iters.insert(0, process_package(stdin, hash_functions))
118     yaml.safe_dump_all(itertools.chain(*iters), sys.stdout)
119
120 if __name__ == "__main__":
121     main()