diff --git a/libmproxy/encoding.py b/libmproxy/encoding.py index b56c03a69..aab4141b8 100644 --- a/libmproxy/encoding.py +++ b/libmproxy/encoding.py @@ -10,13 +10,21 @@ ENCODINGS = set(["identity", "gzip", "deflate"]) def decode(encoding, content): encoding_map = { - "identity": decode_identity, + "identity": identity, "gzip": decode_gzip, "deflate": decode_deflate, } - return encoding_map.get(encoding, decode_identity)(content) + return encoding_map.get(encoding, identity)(content) -def decode_identity(content): +def encode(encoding, content): + encoding_map = { + "identity": identity, + "gzip": encode_gzip, + "deflate": encode_deflate, + } + return encoding_map.get(encoding, identity)(content) + +def identity(content): """ Returns content unchanged. Identity is the default value of Accept-Encoding headers. @@ -30,9 +38,16 @@ def decode_gzip(content): except IOError: return None +def encode_gzip(content): + s = cStringIO.StringIO() + gf = gzip.GzipFile(fileobj=s, mode='wb') + gf.write(content) + gf.close() + return s.getvalue() + def decode_deflate(content): """ - Returns decompress data for DEFLATE. Some servers may respond with + Returns decompressed data for DEFLATE. Some servers may respond with compressed data without a zlib header or checksum. An undocumented feature of zlib permits the lenient decompression of data missing both values. @@ -46,3 +61,9 @@ def decode_deflate(content): return zlib.decompress(content, -15) except zlib.error: return None + +def encode_deflate(content): + """ + Returns compressed content, always including zlib header and checksum. + """ + return zlib.compress(content) diff --git a/test/test_encoding.py b/test/test_encoding.py index 00632c0e8..dff146ee7 100644 --- a/test/test_encoding.py +++ b/test/test_encoding.py @@ -4,30 +4,28 @@ import libpry import cStringIO import gzip, zlib -class udecode_identity(libpry.AutoTree): - def test_decode(self): - assert 'string' == encoding.decode('identity', 'string') +class uidentity(libpry.AutoTree): + def test_simple(self): + assert "string" == encoding.decode("identity", "string") + assert "string" == encoding.encode("identity", "string") def test_fallthrough(self): - assert 'string' == encoding.decode('nonexistent encoding', 'string') + assert "string" == encoding.decode("nonexistent encoding", "string") + assert "string" == encoding.encode("nonexistent encoding", "string") -class udecode_gzip(libpry.AutoTree): +class ugzip(libpry.AutoTree): def test_simple(self): - s = cStringIO.StringIO() - gf = gzip.GzipFile(fileobj=s, mode='wb') - gf.write('string') - gf.close() - assert 'string' == encoding.decode('gzip', s.getvalue()) + assert "string" == encoding.decode("gzip", encoding.encode("gzip", "string")) assert None == encoding.decode("gzip", "bogus") -class udecode_deflate(libpry.AutoTree): +class udeflate(libpry.AutoTree): def test_simple(self): - assert 'string' == encoding.decode('deflate', zlib.compress('string')) - assert 'string' == encoding.decode('deflate', zlib.compress('string')[2:-4]) + assert "string" == encoding.decode("deflate", encoding.encode("deflate", "string")) + assert "string" == encoding.decode("deflate", encoding.encode("deflate", "string")[2:-4]) assert None == encoding.decode("deflate", "bogus") tests = [ - udecode_identity(), - udecode_gzip(), - udecode_deflate() + uidentity(), + ugzip(), + udeflate() ]