#!/usr/bin/python

import os
import sys
import subprocess

# 2020-09-17    Somehow broke gzip compression
#               Can't get subprocess.Popen to work anymore

# 2022-**-**    Increased max load average from 3.5 to 6
# 2022-05-21    and to 8
# 2022-08-08    to 12

'''
compressout
===========

Simple CGI output compression module.

Not all webservers support compressing the output stream of CGI scripts.
This script determines whether or not a client accept gzip compression and
compresses the output stream of the script.

NOTICE: The `cgitb` module will write to stdout if the script crashes,
you should use a browser that does not accept gzip, when you are
testing your scripts.

NOTICE: This module uses two global variable: `use_gzip` and `body`.
It's supposed to be used like an object, but rather than using a class,
this is imported "pre-created" or so to say.

NOTICE: There are two constants:
  `max_load_avg_1min`:
    This is the maximum allowed load average under one minute.
    If the one minute load average exceeds this value, compressout will
    return a 503 response to the client and abort the process.
  `http503_body`:
    A plain text error message to send if the server is overloaded.
You should modify these constants to fit your own needs.


Example / TL;DR
===============

import compressout
compressout.init()
compressout.write_head('Content-Type: text/plain\r\n')
compressout.write_head('Foo: test\r\nBar: test\r\n')
# Blank line required for terminating the HEAD section
compressout.write_head('\r\n')
compressout.write_body('Hello world!\n')
compressout.write_body('Bye.\n')
compressout.done()


Functions
=========

init(write_headers=True)
------------------------

    Initialize the module.  This function will detect if the client
    supports gzip.
    If `write_headers` is True, the function writes a
    'Vary: Content-Encoding' header and (if gzip is used) a
    'Content-Encdoing: gzip' header.


write_head(s) and write_h(s)
----------------------------
    
    This function is used to print all HTTP headers **and the blank line
    separating the head from the body**.
    Write `s` to standard output, will never go through gzip.


write_body(s) and write_b(s)
----------------------------

    Write part of body.
    NOTICE: You need to have printed the blank line after the headers
    with the `write_h` (or `write_head`) fuction.
    
    If gzip is supported by the client
    ----------------------------------
    
        `s` will be appended to a local buffer which the `done` function
        will compress and print.
    
    If gzip is not supported
    ------------------------
    
        `s` will go straight to stdout. The `done` function won't do
        anything.


done()
------

    Done writing output.
    This function will invoke gzip.


Dos and don'ts
==============

    * Try to call `init` and `done` at convenient locations like on the
      "outside" of a main function, i.e. don't repeat yourself by calling
      these two functions everywhere in your code.
    * Never call `write_head` after any call to `write_body`.
    * Always call `done` when your done.
    * Use only compressout to write output, otherwise you'll have a mess.
    * NOTICE: The `cgitb` module will write to stdout if the script
      crashes, you should use a browser that does not accept gzip,
      when you are testing your scripts.  Eg, lwp-request.
      `GET http://example.com/ | less` is excellent for debuggin.
'''

### GLOBALS ###

use_gzip = None     # Whether or not to compress the body
body = ''           # The body is stored here if it is to be compressed

### END GLOBALS ###


### CONSTANTS -- Configure for your own needs ###

# If the load average of the last one minute exceeds the hard coded value,
# this script will return a 503 response and abort the process.
max_load_avg_1min = 12      # 3.5

http503_body = '''
Service temporarily unavailable!
Wait at least two minutes before trying again.
Re-attempting prematurely may result in banning your IP address.

-- END --

'''

#               #############################################

if sys.version_info[0] > 2:
    def output(s):
        if isinstance(s, str):
            sys.stdout.buffer.write(s.encode('utf-8'))
        elif isinstance(s, bytes):
            sys.stdout.buffer.write(s)
        else:
            raise TypeError("Unsupported datatype")
    flush = sys.stdout.buffer.flush
else:
    output = sys.stdout.write
    flush = sys.stdout.flush


def overload_test(too_late=False):
    '''
    '''
    if os.getloadavg()[0] > max_load_avg_1min:
        if not too_late:
            output('Status: 503\n')
            output('Content-Type: text/plain\n')
            output('Retry-After: 90\n')
            output(http503_body)
            flush()
        #os.abort()
        sys.exit(1)


def init(write_headers=True):
    '''
    Initialize the module.  This function will detect if the client
    support gzip.  This will also set the global variable `debug_cookie`.
    If `write_headers`, write a 'Vary' and (if used)
    'Content-Encoding' header.
    '''
    
    global use_gzip
    global body
    global debug_cookie
    
    # This is the only place where sending a 503 message will work.
    # write_h:
    #   - Message body may need to be compressed.
    #   - Possibility of conflicting Status headers.
    # write_b:
    #   - Message body may need to be compressed.
    #   - Message body may be application/xhtml+xml
    # done:
    #   - Message body needs to be compressed if `use_gzip`.
    #   - Body has already been written if not `use_gzip`.
    overload_test(too_late=False)

    use_gzip = 'gzip' in os.getenv('HTTP_ACCEPT_ENCODING', '')
    body = ''

    use_gzip = False

    if write_headers:
        output('Vary: Accept-Encoding\n')
        if use_gzip:
            output('Content-Encoding: gzip\n')

    debug_cookie = 'debug=on' in os.getenv('HTTP_COOKIE', '')
    if 'debug=on' in os.getenv('QUERY_STRING', ''):
        output('Set-Cookie: debug=on\n')
        debug_cookie = True
    if 'debug=off' in os.getenv('QUERY_STRING', ''):
        output('Set-Cookie: debug=off\n')
        debug_cookie = False


def write_head(s):
    write_h(s)

def write_h(s):
    '''
    Write part of header.
    Write `s` to standard output, will never go through gzip.
    '''

    overload_test(too_late=True)
    output(s)


def write_body(s):
    write_b(s)

def write_b(s):
    '''
    Write part of body.
    
    gzip is supported by the client
    -------------------------------
    
        `s` will be appended to a local buffer
        which `done` will compress and print.
    
    gzip is not supported
    ---------------------
    
        `s` will go straight to stdout.
    '''
    
    global body
    overload_test(too_late=True)
    
    if use_gzip:
        body += s
    else:
        output(s)


def done():
    '''
    Done writing output.
    This function will invoke gzip.
    '''
    
    overload_test(too_late=True)
    if use_gzip:
        gzip = subprocess.Popen(
            ['gzip'],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )
        if sys.version_info[0] > 2:
            body = body.encode('utf-8')
            sys.stderr.write('Body encoded\n')
        sys.stderr.write('Just before communicate\n')
        gzip_stdout = gzip.communicate(body)[0]
        sys.stderr.write('Just after communicate\n')
        #gzip_stdout = things[0]
        #sys.stderr.write('After extracting data\n')
        #sys.stderr.write(gzip_stderr)
        output(gzip_stdout)
        sys.stderr.write('done() complete\n')