mitmproxy/netlib/tutils.py

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

from netlib import tcp, utils, 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, six.string_types):
        if expected.lower() not in str(actual).lower():
            six.reraise(AssertionError, AssertionError(
                "Expected %s, but caught %s" % (
                    repr(str(expected)), actual
                )
            ), exc_tb)
    else:
        if not isinstance(actual, expected):
            six.reraise(AssertionError, AssertionError(
                "Expected %s, but caught %s %s" % (
                    expected.__name__, actual.__class__.__name__, str(actual)
                )
            ), exc_tb)


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__)


def treq(content="content", scheme="http", host="address", port=22):
    """
    @return: libmproxy.protocol.http.HTTPRequest
    """
    headers = http.Headers()
    headers["header"] = "qvalue"
    req = http.Request(
        "relative",
        "GET",
        scheme,
        host,
        port,
        "/path",
        (1, 1),
        headers,
        content,
        None,
        None,
    )
    return req


def treq_absolute(content="content"):
    """
    @return: libmproxy.protocol.http.HTTPRequest
    """
    r = treq(content)
    r.form_in = r.form_out = "absolute"
    r.host = "address"
    r.port = 22
    r.scheme = "http"
    return r


def tresp(content="message"):
    """
    @return: libmproxy.protocol.http.HTTPResponse
    """

    headers = http.Headers()
    headers["header_response"] = "svalue"

    resp = http.semantics.Response(
        (1, 1),
        200,
        "OK",
        headers,
        content,
        timestamp_start=time.time(),
        timestamp_end=time.time(),
    )
    return resp