1import uctypes
2
3# http://www.gnu.org/software/tar/manual/html_node/Standard.html
4TAR_HEADER = {
5    "name": (uctypes.ARRAY | 0, uctypes.UINT8 | 100),
6    "size": (uctypes.ARRAY | 124, uctypes.UINT8 | 11),
7}
8
9DIRTYPE = "dir"
10REGTYPE = "file"
11
12
13def roundup(val, align):
14    return (val + align - 1) & ~(align - 1)
15
16
17class FileSection:
18    def __init__(self, f, content_len, aligned_len):
19        self.f = f
20        self.content_len = content_len
21        self.align = aligned_len - content_len
22
23    def read(self, sz=65536):
24        if self.content_len == 0:
25            return b""
26        if sz > self.content_len:
27            sz = self.content_len
28        data = self.f.read(sz)
29        sz = len(data)
30        self.content_len -= sz
31        return data
32
33    def readinto(self, buf):
34        if self.content_len == 0:
35            return 0
36        if len(buf) > self.content_len:
37            buf = memoryview(buf)[: self.content_len]
38        sz = self.f.readinto(buf)
39        self.content_len -= sz
40        return sz
41
42    def skip(self):
43        sz = self.content_len + self.align
44        if sz:
45            buf = bytearray(16)
46            while sz:
47                s = min(sz, 16)
48                self.f.readinto(buf, s)
49                sz -= s
50
51
52class TarInfo:
53    def __str__(self):
54        return "TarInfo(%r, %s, %d)" % (self.name, self.type, self.size)
55
56
57class TarFile:
58    def __init__(self, name=None, fileobj=None):
59        if fileobj:
60            self.f = fileobj
61        else:
62            self.f = open(name, "rb")
63        self.subf = None
64
65    def next(self):
66        if self.subf:
67            self.subf.skip()
68        buf = self.f.read(512)
69        if not buf:
70            return None
71
72        h = uctypes.struct(uctypes.addressof(buf), TAR_HEADER, uctypes.LITTLE_ENDIAN)
73
74        # Empty block means end of archive
75        if h.name[0] == 0:
76            return None
77
78        d = TarInfo()
79        d.name = str(h.name, "utf-8").rstrip("\0")
80        d.size = int(bytes(h.size), 8)
81        d.type = [REGTYPE, DIRTYPE][d.name[-1] == "/"]
82        self.subf = d.subf = FileSection(self.f, d.size, roundup(d.size, 512))
83        return d
84
85    def __iter__(self):
86        return self
87
88    def __next__(self):
89        v = self.next()
90        if v is None:
91            raise StopIteration
92        return v
93
94    def extractfile(self, tarinfo):
95        return tarinfo.subf
96