Home | History | Annotate | Download | only in webpagereplay 
      1 #!/usr/bin/env python
      2 # Copyright 2011 Google Inc. All Rights Reserved.
      3 #
      4 # Licensed under the Apache License, Version 2.0 (the "License");
      5 # you may not use this file except in compliance with the License.
      6 # You may obtain a copy of the License at
      7 #
      8 #      http://www.apache.org/licenses/LICENSE-2.0
      9 #
     10 # Unless required by applicable law or agreed to in writing, software
     11 # distributed under the License is distributed on an "AS IS" BASIS,
     12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     13 # See the License for the specific language governing permissions and
     14 # limitations under the License.
     15 
     16 """Apply gzip/deflate to separate chunks of data."""
     17 
     18 import struct
     19 import zlib
     20 
     21 GZIP_HEADER = (
     22     '\037\213'             # magic header
     23     '\010'                 # compression method
     24     '\000'                 # flags (none)
     25     '\000\000\000\000'     # packed time (use zero)
     26     '\002'
     27     '\377')
     28 
     29 
     30 def compress_chunks(uncompressed_chunks, use_gzip):
     31   """Compress a list of data with gzip or deflate.
     32 
     33   The returned chunks may be used with HTTP chunked encoding.
     34 
     35   Args:
     36     uncompressed_chunks: a list of strings
     37        (e.g. ["this is the first chunk", "and the second"])
     38     use_gzip: if True, compress with gzip. Otherwise, use deflate.
     39 
     40   Returns:
     41     [compressed_chunk_1, compressed_chunk_2, ...]
     42   """
     43   if use_gzip:
     44     size = 0
     45     crc = zlib.crc32("") & 0xffffffffL
     46     compressor = zlib.compressobj(
     47         6, zlib.DEFLATED, -zlib.MAX_WBITS, zlib.DEF_MEM_LEVEL, 0)
     48   else:
     49     compressor = zlib.compressobj()
     50   compressed_chunks = []
     51   last_index = len(uncompressed_chunks) - 1
     52   for index, data in enumerate(uncompressed_chunks):
     53     chunk = ''
     54     if use_gzip:
     55       size += len(data)
     56       crc = zlib.crc32(data, crc) & 0xffffffffL
     57       if index == 0:
     58         chunk += GZIP_HEADER
     59     chunk += compressor.compress(data)
     60     if index < last_index:
     61       chunk += compressor.flush(zlib.Z_SYNC_FLUSH)
     62     else:
     63       chunk += (compressor.flush(zlib.Z_FULL_FLUSH) +
     64                 compressor.flush())
     65       if use_gzip:
     66         chunk += (struct.pack("<L", long(crc)) +
     67                   struct.pack("<L", long(size)))
     68     compressed_chunks.append(chunk)
     69   return compressed_chunks
     70 
     71 
     72 def uncompress_chunks(compressed_chunks, use_gzip):
     73   """Uncompress a list of data compressed with gzip or deflate.
     74 
     75   Args:
     76     compressed_chunks: a list of compressed data
     77     use_gzip: if True, uncompress with gzip. Otherwise, use deflate.
     78 
     79   Returns:
     80     [uncompressed_chunk_1, uncompressed_chunk_2, ...]
     81   """
     82   if use_gzip:
     83     decompress = zlib.decompressobj(16 + zlib.MAX_WBITS).decompress
     84   else:
     85     decompress = zlib.decompressobj(-zlib.MAX_WBITS).decompress
     86   return [decompress(c) for c in compressed_chunks]
     87