mitmproxy/netlib/tutils.py

from io import BytesIO
import tempfile
import os
import time
import shutil
from contextlib import contextmanager
import sys

from netlib import utils, tcp, http


def treader(bytes):
    """
        Construct a tcp.Read object from bytes.
    """
    fp = BytesIO(bytes)
    return tcp.Reader(fp)


@contextmanager
def tmpdir(*args, **kwargs):
    orig_workdir = os.getcwd()
    temp_workdir = tempfile.mkdtemp(*args, **kwargs)
    os.chdir(temp_workdir)

    yield temp_workdir

    os.chdir(orig_workdir)
    shutil.rmtree(temp_workdir)


def _check_exception(expected, actual, exc_tb):
    if isinstance(expected, str):
        if expected.lower() not in str(actual).lower():
            raise AssertionError(
                "Expected %s, but caught %s" % (
                    repr(expected), repr(actual)
                )
            )
    else:
        if not isinstance(actual, expected):
            raise AssertionError(
                "Expected %s, but caught %s %s" % (
                    expected.__name__, actual.__class__.__name__, repr(actual)
                )
            )


def raises(expected_exception, obj=None, *args, **kwargs):
    """
        Assert that a callable raises a specified exception.

        :exc An exception class or a string. If a class, assert that an
        exception of this type is raised. If a string, assert that the string
        occurs in the string representation of the exception, based on a
        case-insenstivie match.

        :obj A callable object.

        :args Arguments to be passsed to the callable.

        :kwargs Arguments to be passed to the callable.
    """
    if obj is None:
        return RaisesContext(expected_exception)
    else:
        try:
            ret = obj(*args, **kwargs)
        except Exception as actual:
            _check_exception(expected_exception, actual, sys.exc_info()[2])
        else:
            raise AssertionError("No exception raised. Return value: {}".format(ret))


class RaisesContext(object):
    def __init__(self, expected_exception):
        self.expected_exception = expected_exception

    def __enter__(self):
        return

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not exc_type:
            raise AssertionError("No exception raised.")
        else:
            _check_exception(self.expected_exception, exc_val, exc_tb)
        return True


test_data = utils.Data(__name__)
# FIXME: Temporary workaround during repo merge.
test_data.dirname = os.path.join(test_data.dirname, "..", "test", "netlib")


def treq(**kwargs):
    """
    Returns:
        netlib.http.Request
    """
    default = dict(
        first_line_format="relative",
        method=b"GET",
        scheme=b"http",
        host=b"address",
        port=22,
        path=b"/path",
        http_version=b"HTTP/1.1",
        headers=http.Headers(((b"header", b"qvalue"), (b"content-length", b"7"))),
        content=b"content"
    )
    default.update(kwargs)
    return http.Request(**default)


def tresp(**kwargs):
    """
    Returns:
        netlib.http.Response
    """
    default = dict(
        http_version=b"HTTP/1.1",
        status_code=200,
        reason=b"OK",
        headers=http.Headers(((b"header-response", b"svalue"), (b"content-length", b"7"))),
        content=b"message",
        timestamp_start=time.time(),
        timestamp_end=time.time(),
    )
    default.update(kwargs)
    return http.Response(**default)