tarwriter.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  1. #
  2. # The Qubes OS Project, http://www.qubes-os.org
  3. #
  4. # Copyright (C) 2016 Marek Marczykowski-Górecki
  5. # <marmarek@invisiblethingslab.com>
  6. #
  7. # This library is free software; you can redistribute it and/or
  8. # modify it under the terms of the GNU Lesser General Public
  9. # License as published by the Free Software Foundation; either
  10. # version 2.1 of the License, or (at your option) any later version.
  11. #
  12. # This library is distributed in the hope that it will be useful,
  13. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  14. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  15. # Lesser General Public License for more details.
  16. #
  17. # You should have received a copy of the GNU Lesser General Public
  18. # License along with this library; if not, see <https://www.gnu.org/licenses/>.
  19. import argparse
  20. import functools
  21. import os
  22. import subprocess
  23. import tarfile
  24. import io
  25. BUF_SIZE = 409600
  26. class TarSparseInfo(tarfile.TarInfo):
  27. def __init__(self, name="", sparsemap=None):
  28. super(TarSparseInfo, self).__init__(name)
  29. if sparsemap is not None:
  30. self.type = tarfile.REGTYPE
  31. self.sparsemap = sparsemap
  32. self.sparsemap_buf = self.format_sparse_map()
  33. # compact size
  34. self.size = functools.reduce(lambda x, y: x+y[1], sparsemap,
  35. 0) + len(self.sparsemap_buf)
  36. self.pax_headers['GNU.sparse.major'] = '1'
  37. self.pax_headers['GNU.sparse.minor'] = '0'
  38. self.pax_headers['GNU.sparse.name'] = name
  39. self.pax_headers['GNU.sparse.realsize'] = str(self.realsize)
  40. self.name = '{}/GNUSparseFile.{}/{}'.format(
  41. os.path.dirname(name), os.getpid(), os.path.basename(name))
  42. else:
  43. self.sparsemap = []
  44. self.sparsemap_buf = b''
  45. @property
  46. def realsize(self):
  47. if self.sparsemap:
  48. return self.sparsemap[-1][0] + self.sparsemap[-1][1]
  49. return self.size
  50. def format_sparse_map(self):
  51. sparsemap_txt = (str(len(self.sparsemap)) + '\n' +
  52. ''.join('{}\n{}\n'.format(*entry) for entry in self.sparsemap))
  53. sparsemap_txt_len = len(sparsemap_txt)
  54. if sparsemap_txt_len % tarfile.BLOCKSIZE:
  55. padding = '\0' * (tarfile.BLOCKSIZE -
  56. sparsemap_txt_len % tarfile.BLOCKSIZE)
  57. else:
  58. padding = ''
  59. return (sparsemap_txt + padding).encode()
  60. def tobuf(self, format=tarfile.PAX_FORMAT, encoding=tarfile.ENCODING,
  61. errors="strict"):
  62. # pylint: disable=redefined-builtin
  63. header_buf = super(TarSparseInfo, self).tobuf(format, encoding, errors)
  64. return header_buf + self.sparsemap_buf
  65. def get_sparse_map(input_file):
  66. '''
  67. Return map of the file where actual data is present, ignoring zero-ed
  68. blocks. Last entry of the map spans to the end of file, even if that part is
  69. zero-size (when file ends with zeros).
  70. This function is performance critical.
  71. :param input_file: io.File object
  72. :return: iterable of (offset, size)
  73. '''
  74. zero_block = bytearray(tarfile.BLOCKSIZE)
  75. buf = bytearray(BUF_SIZE)
  76. in_data_block = False
  77. data_block_start = 0
  78. buf_start_offset = 0
  79. while True:
  80. buf_len = input_file.readinto(buf)
  81. if not buf_len:
  82. break
  83. for offset in range(0, buf_len, tarfile.BLOCKSIZE):
  84. if buf[offset:offset+tarfile.BLOCKSIZE] == zero_block:
  85. if in_data_block:
  86. in_data_block = False
  87. yield (data_block_start,
  88. buf_start_offset+offset-data_block_start)
  89. else:
  90. if not in_data_block:
  91. in_data_block = True
  92. data_block_start = buf_start_offset+offset
  93. buf_start_offset += buf_len
  94. if in_data_block:
  95. yield (data_block_start, buf_start_offset-data_block_start)
  96. else:
  97. # always emit last slice to the input end - otherwise extracted file
  98. # will be truncated
  99. yield (buf_start_offset, 0)
  100. def copy_sparse_data(input_stream, output_stream, sparse_map):
  101. '''Copy data blocks from input to output according to sparse_map
  102. :param input_stream: io.IOBase input instance
  103. :param output_stream: io.IOBase output instance
  104. :param sparse_map: iterable of (offset, size)
  105. '''
  106. buf = bytearray(BUF_SIZE)
  107. for chunk in sparse_map:
  108. input_stream.seek(chunk[0])
  109. left = chunk[1]
  110. while left:
  111. if left > BUF_SIZE:
  112. read = input_stream.readinto(buf)
  113. output_stream.write(buf[:read])
  114. else:
  115. buf_trailer = input_stream.read(left)
  116. read = len(buf_trailer)
  117. output_stream.write(buf_trailer)
  118. left -= read
  119. if not read:
  120. raise Exception('premature EOF')
  121. def finalize(output):
  122. '''Write EOF blocks'''
  123. output.write(b'\0' * 512)
  124. output.write(b'\0' * 512)
  125. def main(args=None):
  126. parser = argparse.ArgumentParser()
  127. parser.add_argument('--override-name', action='store', dest='override_name',
  128. help='use this name in tar header')
  129. parser.add_argument('--use-compress-program', default=None,
  130. metavar='COMMAND', action='store', dest='use_compress_program',
  131. help='Filter data through COMMAND.')
  132. parser.add_argument('input_file',
  133. help='input file name')
  134. parser.add_argument('output_file', default='-', nargs='?',
  135. help='output file name')
  136. args = parser.parse_args(args)
  137. input_file = io.open(args.input_file, 'rb')
  138. sparse_map = list(get_sparse_map(input_file))
  139. header_name = args.input_file
  140. if args.override_name:
  141. header_name = args.override_name
  142. tar_info = TarSparseInfo(header_name, sparse_map)
  143. if args.output_file == '-':
  144. output = io.open('/dev/stdout', 'wb')
  145. else:
  146. output = io.open(args.output_file, 'wb')
  147. if args.use_compress_program:
  148. compress = subprocess.Popen([args.use_compress_program],
  149. stdin=subprocess.PIPE, stdout=output)
  150. output = compress.stdin
  151. else:
  152. compress = None
  153. output.write(tar_info.tobuf(tarfile.PAX_FORMAT))
  154. copy_sparse_data(input_file, output, sparse_map)
  155. finalize(output)
  156. input_file.close()
  157. output.close()
  158. if compress is not None:
  159. compress.wait()
  160. return compress.returncode
  161. return 0
  162. if __name__ == '__main__':
  163. main()