#!/usr/bin/env python

"""
File:    check_cache_headers.py
Version: Python 3.6.5
Author:  Nicholas Russo (http://njrusmc.net)
Purpose: Quickly check to ensure the HTTP cache testing files all have the
         expected Cache-Control headers. These should not be overwritten
         during new CI pushes, but this script guarantees it.
"""

import sys
import requests


def main():
    """
    Execution begins here.
    """

    # Define base URL for all test files
    base_url = "http://njrusmc.net/cache"

    # Enumerate the files and the expected Cache-Control header string
    file_dict = {
        "zero128k_none.test": None,
        "zero128k_public.test": "public",
        "zero128k_public60.test": "public, max-age=60",
        "zero128k_public3600.test": "public, max-age=3600",
        "zero128k_private.test": "private",
        "zero128k_private60.test": "private, max-age=60",
        "zero128k_private3600.test": "private, max-age=3600",
        "zero128k_nocache.test": "no-cache",
        "zero128k_nocache60.test": "no-cache, max-age=60",
        "zero128k_nocache3600.test": "no-cache, max-age=3600",
        "zero128k_nostore.test": "no-store",
        "rand128k_none.test": None,
        "rand128k_public.test": "public",
        "rand128k_public60.test": "public, max-age=60",
        "rand128k_public3600.test": "public, max-age=3600",
        "rand128k_private.test": "private",
        "rand128k_private60.test": "private, max-age=60",
        "rand128k_private3600.test": "private, max-age=3600",
        "rand128k_nocache.test": "no-cache",
        "rand128k_nocache60.test": "no-cache, max-age=60",
        "rand128k_nocache3600.test": "no-cache, max-age=3600",
        "rand128k_nostore.test": "no-store",
    }

    # Create a long lived session and iterate over dict items
    sess = requests.session()
    failures = 0
    for filename, cc_header in file_dict.items():

        # Reveal test case and issue HEAD request
        print(f"Testing {filename} for {cc_header} ... ", end="")
        resp = sess.head(f"{base_url}/{filename}")

        # If a Cache-Control header was specified, test for equality
        if cc_header:
            if resp.headers["Cache-Control"] == cc_header:
                print("OK!")
            else:
                print("FAIL!")
                failures += 1

        # Else a Cache-Control header was not specified, ensure absence
        else:
            if "Cache-Control" not in resp.headers:
                print("OK!")
            else:
                print("FAIL!")
                failures += 1

    # Exit with rc equal to failure count (0 means success)
    sys.exit(failures)


if __name__ == "__main__":
    main()