thttp.py

"""
You may use this code under the UNLICENSE or MIT License.
See README.md for details.

https://github.com/sesh/thttp
"""

import gzip
import json as json_lib
import mimetypes
import secrets
import ssl
from base64 import b64encode
from collections import namedtuple
from http import HTTPStatus
from http.cookiejar import CookieJar
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode
from urllib.request import (
    HTTPCookieProcessor,
    HTTPRedirectHandler,
    HTTPSHandler,
    Request,
    build_opener,
)

Response = namedtuple("Response", "request content json status url headers cookiejar")


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def request(
    url,
    params={},
    json=None,
    data=None,
    headers={},
    method="GET",
    verify=True,
    redirect=True,
    cookiejar=None,
    basic_auth=None,
    timeout=None,
    files={},  # note: experimental
):
    """
    Returns a (named)tuple with the following properties:
        - request
        - content
        - json (dict; or None)
        - headers (dict; all lowercase keys)
            - https://stackoverflow.com/questions/5258977/are-http-headers-case-sensitive
        - status
        - url (final url, after any redirects)
        - cookiejar
    """
    method = method.upper()
    headers = {k.lower(): v for k, v in headers.items()}  # lowercase headers

    if params:
        url += "?" + urlencode(params)  # build URL from query parameters

    if json and data:
        raise Exception("Cannot provide both json and data parameters")

    if method not in ["POST", "PATCH", "PUT"] and (json or data):
        raise Exception("Request method must POST, PATCH or PUT if json or data is provided")

    if files and method != "POST":
        raise Exception("Request method must be POST when uploading files")

    if not timeout:
        timeout = 60

    if json:  # if we have json, dump it to a string and put it in our data variable
        headers["content-type"] = "application/json"
        data = json_lib.dumps(json).encode("utf-8")
    elif data and not isinstance(data, (str, bytes)):
        data = urlencode(data).encode()
    elif isinstance(data, str):
        data = data.encode()
    elif files:
        boundary = secrets.token_hex()

        headers["Content-Type"] = f"multipart/form-data; boundary={boundary}"
        data = b""

        for key, file in files.items():
            file_data = file.read()  # okay, we want this to stay as a byte-string
            if isinstance(file_data, str):
                file_data = file_data.encode("utf-8")
            fn = file.name

            mime, _ = mimetypes.guess_type(fn)
            if not mime:
                print("Using default mimetype")
                mime = "application/octet-stream"

            data += b"--" + boundary.encode() + b"\r\n"
            data += b'Content-Disposition: form-data; name="' + key.encode() + b'"; filename="' + fn.encode() + b'"\r\n'
            data += b"Content-Type: " + mime.encode() + b"\r\n\r\n"
            data += file_data + b"\r\n"
            data += b"--" + boundary.encode() + b"--\r\n"

        data = data
        headers["Content-Length"] = len(data)

    if basic_auth and len(basic_auth) == 2 and "authorization" not in headers:
        username, password = basic_auth
        headers["authorization"] = f'Basic {b64encode(f"{username}:{password}".encode()).decode("ascii")}'

    if not cookiejar:
        cookiejar = CookieJar()

    ctx = ssl.create_default_context()
    if not verify:  # ignore ssl errors
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE

    handlers = []
    handlers.append(HTTPSHandler(context=ctx))
    handlers.append(HTTPCookieProcessor(cookiejar=cookiejar))

    if not redirect:
        no_redirect = NoRedirect()
        handlers.append(no_redirect)

    opener = build_opener(*handlers)
    req = Request(url, data=data, headers=headers, method=method)

    try:
        with opener.open(req, timeout=timeout) as resp:
            status, content, resp_url = (resp.getcode(), resp.read(), resp.geturl())
            headers = {k.lower(): v for k, v in list(resp.info().items())}

            if "gzip" in headers.get("content-encoding", ""):
                content = gzip.decompress(content)

            json = (
                json_lib.loads(content)
                if "application/json" in headers.get("content-type", "").lower() and content
                else None
            )
    except HTTPError as e:
        status, content, resp_url = (e.code, e.read(), e.geturl())
        headers = {k.lower(): v for k, v in list(e.headers.items())}

        if "gzip" in headers.get("content-encoding", ""):
            content = gzip.decompress(content)

        json = (
            json_lib.loads(content)
            if "application/json" in headers.get("content-type", "").lower() and content
            else None
        )

    return Response(req, content, json, status, resp_url, headers, cookiejar)


def pretty(response, headers_only=False):
    RESET = "\033[0m"
    HIGHLIGHT = "\033[34m"
    HTTP_STATUSES = {x.value: x.name for x in HTTPStatus}

    # status code
    print(HIGHLIGHT + str(response.status) + " " + RESET + HTTP_STATUSES.get(response.status, ""))

    # headers
    for k in sorted(response.headers.keys()):
        print(HIGHLIGHT + k + RESET + ": " + response.headers[k])

    if headers_only:
        return

    # blank line
    print()

    # response body
    if response.json:
        print(json_lib.dumps(response.json, indent=2))
    else:
        print(response.content.decode())


import contextlib  # noqa: E402
import os  # noqa: E402
import unittest  # noqa: E402
from io import StringIO  # noqa: E402
from unittest.mock import patch  # noqa: E402


class RequestTestCase(unittest.TestCase):
    def test_cannot_provide_json_and_data(self):
        with self.assertRaises(Exception):
            request(
                "https://httpbingo.org/post",
                json={"name": "Brenton"},
                data="This is some form data",
            )

    def test_should_fail_if_json_or_data_and_not_p_method(self):
        with self.assertRaises(Exception):
            request("https://httpbingo.org/post", json={"name": "Brenton"})

        with self.assertRaises(Exception):
            request("https://httpbingo.org/post", json={"name": "Brenton"}, method="HEAD")

    def test_should_set_content_type_for_json_request(self):
        response = request("https://httpbingo.org/post", json={"name": "Brenton"}, method="POST")
        self.assertEqual(response.request.headers["Content-type"], "application/json")

    def test_should_work(self):
        response = request("https://httpbingo.org/get")
        self.assertEqual(response.status, 200)

    def test_should_create_url_from_params(self):
        response = request(
            "https://httpbingo.org/get",
            params={"name": "brenton", "library": "tiny-request"},
        )
        self.assertEqual(response.url, "https://httpbingo.org/get?name=brenton&library=tiny-request")

    def test_should_return_headers(self):
        response = request("https://httpbingo.org/response-headers", params={"Test-Header": "value"})
        self.assertEqual(response.headers["test-header"], "value")

    def test_should_populate_json(self):
        response = request("https://httpbingo.org/json")
        self.assertTrue("slideshow" in response.json)

    def test_should_return_response_for_404(self):
        response = request("https://httpbingo.org/404")
        self.assertEqual(response.status, 404)
        self.assertTrue("application/json" in response.headers["content-type"])

    def test_should_fail_with_bad_ssl(self):
        with self.assertRaises(URLError):
            request("https://expired.badssl.com/")

    def test_should_load_bad_ssl_with_verify_false(self):
        response = request("https://expired.badssl.com/", verify=False)
        self.assertEqual(response.status, 200)

    def test_should_form_encode_non_json_post_requests(self):
        response = request("https://httpbingo.org/post", data={"name": "test-user"}, method="POST")
        self.assertEqual(response.json["form"]["name"], ["test-user"])

    def test_should_follow_redirect(self):
        response = request(
            "https://httpbingo.org/redirect-to",
            params={"url": "https://example.org/"},
        )
        self.assertEqual(response.url, "https://example.org/")
        self.assertEqual(response.status, 200)

    def test_should_not_follow_redirect_if_redirect_false(self):
        response = request(
            "https://httpbingo.org/redirect-to",
            params={"url": "https://example.org/"},
            redirect=False,
        )
        self.assertEqual(response.status, 302)

    def test_cookies(self):
        response = request(
            "https://httpbingo.org/cookies/set",
            params={"cookie": "test"},
            redirect=False,
        )
        response = request("https://httpbingo.org/cookies", cookiejar=response.cookiejar)
        self.assertEqual(response.json["cookie"], "test")

    def test_basic_auth(self):
        response = request("http://httpbingo.org/basic-auth/user/passwd", basic_auth=("user", "passwd"))
        self.assertEqual(response.json["authorized"], True)

    def test_should_handle_gzip(self):
        response = request("http://httpbingo.org/gzip", headers={"Accept-Encoding": "gzip"})
        self.assertEqual(response.json["gzipped"], True)

    def test_should_handle_gzip_error(self):
        response = request("http://httpbingo.org/status/418", headers={"Accept-Encoding": "gzip"})
        self.assertEqual(response.content, b"I'm a teapot!")

    def test_should_timeout(self):
        import socket

        with self.assertRaises((TimeoutError, socket.timeout)):
            request("http://httpbingo.org/delay/3", timeout=1)

    def test_should_handle_head_requests(self):
        response = request("http://httpbingo.org/head", method="HEAD")
        self.assertTrue(response.content == b"")

    def test_should_post_data_string(self):
        response = request(
            "https://ntfy.sh/thttp-test-ntfy",
            data="The thttp test suite was executed!",
            method="POST",
        )
        self.assertTrue(response.json["topic"] == "thttp-test-ntfy")

    def test_pretty_output(self):
        response = request("https://basehtml.xyz")

        f = StringIO()
        with contextlib.redirect_stdout(f):
            pretty(response)

        f.seek(0)
        output = f.read()

        self.assertTrue("text/html; charset=utf-8" in output)
        self.assertTrue("<h1>base.html</h1>" in output)

    def test_pretty_output_headers_only(self):
        response = request("https://basehtml.xyz")

        f = StringIO()
        with contextlib.redirect_stdout(f):
            pretty(response, headers_only=True)

        f.seek(0)
        output = f.read()

        self.assertTrue("text/html; charset=utf-8" in output)
        self.assertTrue("<h1>base.html</h1>" not in output)

    def test_thttp_with_mocked_response(self):
        mocked_response = Response(None, None, {"response": "mocked"}, 200, None, None, None)

        with patch("thttp.request", side_effect=[mocked_response]):
            response = request("https://example.org")
            self.assertEqual("mocked", response.json["response"])

    def test_upload_single_file(self):
        token = os.environ.get("MEDIAPUB_TOKEN")
        url = os.environ.get("MEDIAPUB_URL")

        if not token or not url:
            self.skipTest("Skipping media upload test because environment variables are not available")

        for fn in ["test-image.png", "LICENSE.md"]:
            with open(fn, "rb" if fn.endswith("png") else "r") as f:
                response = request(
                    url,
                    headers={"Authorization": f"Bearer {token}"},
                    files={"file": f},
                    method="POST",
                )

            self.assertEqual(response.status, 201)
            self.assertTrue("location" in response.headers)