#!/usr/bin/env python3
"""Bound one already-buffered gzip member. Offline teaching fixture, not HTTP middleware."""
import gzip
import sys
import unittest
import zlib

INPUT_LIMIT = 64 * 1024
OUTPUT_LIMIT = 256 * 1024


class Rejected(ValueError):
    pass


def decode_one_gzip(blob, input_limit=INPUT_LIMIT, output_limit=OUTPUT_LIMIT):
    """Return complete bytes or reject; never return a partial decoded prefix.

    Caller must bound transport buffering separately. The output limit is not
    a whole-process memory or CPU bound. Concatenated members are rejected.
    """
    for limit in (input_limit, output_limit):
        if type(limit) is not int or limit <= 0:
            raise ValueError('limits must be positive integers')
    if type(blob) is not bytes:
        raise TypeError('bytes required')
    if len(blob) > input_limit:
        raise Rejected('compressed input limit exceeded')
    decoder = zlib.decompressobj(16 + zlib.MAX_WBITS)
    try:
        # Nonzero max_length caps returned output. One extra byte detects excess.
        body = decoder.decompress(blob, output_limit + 1)
    except zlib.error as exc:
        raise Rejected('invalid gzip stream') from exc
    if len(body) > output_limit:
        raise Rejected('decoded output limit exceeded')
    if not decoder.eof:
        raise Rejected('incomplete gzip stream')
    if decoder.unused_data:
        raise Rejected('trailing bytes or additional gzip member')
    return body


class Tests(unittest.TestCase):
    def test_normal(self):
        self.assertEqual(decode_one_gzip(gzip.compress(b'approved report')), b'approved report')

    def test_empty_member(self):
        self.assertEqual(decode_one_gzip(gzip.compress(b'')), b'')

    def test_exact_output_limit(self):
        self.assertEqual(decode_one_gzip(gzip.compress(b'a' * OUTPUT_LIMIT)), b'a' * OUTPUT_LIMIT)

    def test_one_byte_over(self):
        with self.assertRaisesRegex(Rejected, 'decoded output'):
            decode_one_gzip(gzip.compress(b'a' * (OUTPUT_LIMIT + 1)))

    def test_expansion(self):
        blob = gzip.compress(b'A' * (2 * 1024 * 1024), mtime=0)
        self.assertLess(len(blob), INPUT_LIMIT)
        with self.assertRaisesRegex(Rejected, 'decoded output'):
            decode_one_gzip(blob)

    def test_compressed_limit(self):
        blob = gzip.compress(b'abc')
        with self.assertRaisesRegex(Rejected, 'compressed input'):
            decode_one_gzip(blob, input_limit=len(blob) - 1)
        self.assertEqual(decode_one_gzip(blob, input_limit=len(blob)), b'abc')

    def test_truncated_trailer(self):
        with self.assertRaisesRegex(Rejected, 'incomplete'):
            decode_one_gzip(gzip.compress(b'report')[:-1])

    def test_corrupt_checksum(self):
        blob = bytearray(gzip.compress(b'report'))
        blob[-8] ^= 1
        with self.assertRaisesRegex(Rejected, 'invalid'):
            decode_one_gzip(bytes(blob))

    def test_non_gzip(self):
        for blob in (b'', b'not gzip', zlib.compress(b'report')):
            with self.subTest(blob=blob), self.assertRaises(Rejected):
                decode_one_gzip(blob)

    def test_concatenated_members(self):
        with self.assertRaisesRegex(Rejected, 'additional gzip member'):
            decode_one_gzip(gzip.compress(b'first') + gzip.compress(b'second'))

    def test_trailing_garbage(self):
        with self.assertRaisesRegex(Rejected, 'trailing bytes'):
            decode_one_gzip(gzip.compress(b'report') + b'junk')

    def test_invalid_limits(self):
        for bad in (0, -1, True, 2.5, float('nan'), float('inf')):
            for field in ('input_limit', 'output_limit'):
                with self.subTest(field=field, bad=bad), self.assertRaises(ValueError):
                    decode_one_gzip(gzip.compress(b'x'), **{field: bad})

    def test_input_type(self):
        with self.assertRaises(TypeError):
            decode_one_gzip(bytearray(gzip.compress(b'x')))

    def test_bufsize_is_not_a_limit(self):
        blob = gzip.compress(b'A' * (2 * 1024 * 1024), mtime=0)
        result = zlib.decompress(blob, wbits=31, bufsize=OUTPUT_LIMIT)
        self.assertEqual(len(result), 2 * 1024 * 1024)


def demo():
    blob = gzip.compress(b'A' * (2 * 1024 * 1024), mtime=0)
    print('compressed bytes:', len(blob))
    print('compressed check:', 'PASS' if len(blob) <= INPUT_LIMIT else 'FAIL')
    naive = zlib.decompress(blob, wbits=31, bufsize=OUTPUT_LIMIT)
    print('naive decoded bytes:', len(naive))
    try:
        decode_one_gzip(blob)
    except Rejected as exc:
        print('bounded decoder:', str(exc))
    else:
        raise RuntimeError('expansion was unexpectedly accepted')


if __name__ == '__main__':
    if sys.argv[1:] == ['--test']:
        unittest.main(argv=[sys.argv[0]], verbosity=2)
    elif sys.argv[1:]:
        raise SystemExit('usage: agent-gzip-output-limit.py [--test]')
    else:
        demo()