zipfs.py

   1 #! /usr/bin/env python3
   2
   3 from io import SEEK_CUR, SEEK_END
   4 import fuse
   5 from contextlib import ExitStack
   6 import zipstream
   7 from errno import ENOENT, ENOSYS, ESPIPE
   8 from shorthand import ceildiv
   9 from zipfile import ZipInfo, ZipFile, ZIP_STORED, ZIP_DEFLATED
  10 import zlib
  11 from os import fsdecode, fsencode
  12
  13 def main():
  14     parser = fuse.ArgumentParser()
  15     parser.add_argument('--scan-files', action='store_true')
  16     [file, mnt, args] = parser.parse_args()
  17     fs = Filesystem(mnt, args)
  18     with ExitStack() as cleanup:
  19         fs.zip = cleanup.enter_context(open(file, 'rb'))
  20         fs.nodes = {fuse.ROOT_ID: empty_dir()}
  21         fs.namelen = 0
  22         if args.scan_files:
  23             for r in zipstream.iter_files(fs.zip):
  24                 [flags, method, crc, comp, uncomp, name, extra] = r
  25                 member = ZipInfo()
  26
  27                 METHODS = {
  28                     zipstream.METHOD_STORE: ZIP_STORED,
  29                     8: ZIP_DEFLATED,
  30                 }
  31                 member.compress_type = METHODS[method]
  32
  33                 member.compress_size = comp
  34                 member.file_size = uncomp
  35                 member.filename = name
  36                 fs.add_member(member, fs.zip.tell() + extra)
  37                 fs.zip.seek(+extra + comp, SEEK_CUR)
  38             fs.size = fs.zip.tell()
  39         else:
  40             fs.size = fs.zip.seek(0, SEEK_END)
  41             for member in ZipFile(fs.zip, 'r').infolist():
  42                 fs.add_member(member)
  43         fs.decomp_node = None
  44
  45         fs.make_dir()
  46         cleanup.callback(fs.close)
  47         fs.mount("zipfs", file)
  48         while True:
  49             fs.handle_request()
  50
  51 def empty_dir():
  52     return {'type': fuse.Filesystem.DT_DIR, 'links': 2, 'dir': dict()}
  53
  54 class Filesystem(fuse.Filesystem):
  55     def add_member(self, member, offset=None):
  56         name = member.filename.split('/')
  57
  58         dir = self.nodes[fuse.ROOT_ID]
  59         for dirname in name[:-1]:
  60             new_node = fuse.ROOT_ID + len(self.nodes)
  61             node = dir['dir'].setdefault(dirname, new_node)
  62             if node is new_node:
  63                 dir['links'] += 1
  64                 self.nodes[new_node] = empty_dir()
  65                 self.namelen = max(self.namelen, len(dirname))
  66             dir = self.nodes[node]
  67
  68         if name[-1] > '':
  69             new_node = fuse.ROOT_ID + len(self.nodes)
  70             node = dir['dir'].setdefault(name[-1], new_node)
  71             assert node is new_node
  72             self.nodes[new_node] = {
  73                 'type': Filesystem.DT_REG,
  74                 'size': member.file_size,
  75                 'offset': offset,
  76                 'member': member,
  77                 'blocks': ceildiv(member.compress_size, 512),
  78             }
  79             self.namelen = max(self.namelen, len(name[-1]))
  80
  81     def statfs(self):
  82         return dict(
  83             blocks=self.size,
  84             files=len(self.nodes),
  85             namelen=self.namelen,
  86         )
  87
  88     def getattr(self, node):
  89         return self.nodes[node]
  90
  91     def lookup(self, node, name):
  92         node = self.nodes[node]['dir'].get(fsdecode(name))
  93         if node is None:
  94             raise OSError(ENOENT, None)
  95         return node
  96
  97     def readdir(self, node, offset):
  98         dir = self.nodes[node]['dir']
  99         for [i, [name, node]] in enumerate(dir.items()):
 100             if i < offset:
 101                 continue
 102             type = self.nodes[node]['type']
 103             yield (node, fsencode(name), type, i + 1)
 104
 105     def read(self, node, offset, size):
 106         node = self.nodes[node]
 107         comp = node['member'].compress_size
 108         if node['member'].compress_type == ZIP_STORED:
 109             if offset >= comp:
 110                 return b''
 111             return self.read_file(node, offset, min(size, comp - offset))
 112         if node['member'].compress_type == ZIP_DEFLATED:
 113             if offset == 0:
 114                 self.decomp_node = node
 115                 self.decomp_obj = zlib.decompressobj(-zlib.MAX_WBITS)
 116                 self.decomp_offset = 0
 117                 self.comp_offset = 0
 118                 self.comp_left = comp
 119             elif node is not self.decomp_node \
 120                     or offset != self.decomp_offset:
 121                 raise OSError(ESPIPE, None)
 122
 123             result = b''
 124             while len(result) < size and not self.decomp_obj.eof:
 125                 data = max(size, 0x10000) - len(self.decomp_obj.unconsumed_tail)
 126                 data = min(data, self.comp_left)
 127                 data = self.read_file(node, self.comp_offset, data)
 128                 self.comp_left -= len(data)
 129                 self.comp_offset += len(data)
 130
 131                 data = self.decomp_obj.unconsumed_tail + data
 132                 decomp = self.decomp_obj.decompress(data, size - len(result))
 133                 if not data and not decomp:
 134                     raise EOFError()
 135                 result += decomp
 136             self.decomp_offset += len(result)
 137             return result
 138         raise OSError(ENOSYS, None)
 139
 140     def read_file(self, node, offset, size):
 141         if node['offset'] is None:
 142             self.zip.seek(node['member'].header_offset)
 143             r = self.zip.read(4)
 144             assert r == zipstream.FILE_SIG
 145             r = zipstream.read_file_header(self.zip)
 146             [flags, method, crc, comp, uncomp, name, extra] = r
 147             node['offset'] = self.zip.tell() + extra
 148         self.zip.seek(node['offset'] + offset)
 149         return self.zip.read(size)
 150
 151 if __name__ == '__main__':
 152     with fuse.handle_termination():
 153         main()