Python Simple HTTP Client

SimpleHttpClient.py

#!/usr/bin/env python
'''
SimpleHttpClient: Python Simple HTTP Client

Author: Moises P. Sena <moisespsena AT gmail.com>

CONNECT AND CLOSE:

>>> host, port = 'localhost', 80
>>> c = SimpleHttpClient((host, port), host)
>>> c.connect()
>>> c.close()

or

>>> with SimpleHttpClient((host, port), host) as c:
...     # ...
...     pass


COMMANDS:

>>> print(c.request('GET', '/'))
>>> print(c.request('HEAD', '/'))
>>> print(c.request('POST', '/', headers={'Content-Type':'text/plain'}, data='file content'))
>>> def rflines():
...     for l in open("file.txt"):
...         yield l
>>> print(c.request('POST', '/', headers={'Content-Type':'text/plain'}, data=rflines()))

TEST URI:

>>> print(test_uri(c, '/'))

multiples uris:

>>> print(test_uris(c, ['/', '/not-found']))

cli:

`echo -e "/\n/not-found" | python SimpleHttpClient.py localhost 80`
'''

import socket
from cStringIO import StringIO
import sys
import time
import types


CRLF = "\r\n"
# A SIMPLE NOT FOUND URL FOR CLOSE CONNECTION REQUEST
NOT_FOUND_URI = '/-----not-found-%s-uri----' % time.clock()


def read_until(s, match):
    tmp = " " * len(match)
    b = StringIO()

    while not tmp == match:
        c = s.recv(1)
        tmp = "%s%s" % (tmp[1:], c)
        b.write(c)

    return b.getvalue()


def read_response(s, method):
    headers = read_until(s, "\r\n\r\n").strip().split("\r\n")
    protocol, status_code, status_msg = headers[0].split(" ", 2)
    status_code = int(status_code)
    headers = dict([_.split(": ", 1) for _ in headers[1:]])

    if method != 'HEAD' and 'Content-Length' in headers:
        data = s.recv(int(headers['Content-Length']))
    else:
        data = None

    return dict(protocol=protocol, status=(status_code, status_msg),
                headers=headers), data


class SimpleHttpClient(object):
    def __init__(self, contuple, host):
        self.contuple = contuple
        self.host = host
        self.sock = None
        self.closed = False

    def _init_sock(self):
        return socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    def _connect(self, sock):
        sock.connect(self.contuple)

    def connect(self):
        self.sock = self._init_sock()
        self._connect(self.sock)
        self.closed = False

    def _mk_line(self, msg):
        return "%s%s" % (msg, CRLF)

    def _header_msg(self, h, v):
        return "%s: %s%s" % (h, v, CRLF)

    def send(self, *args):
        self.sock.send(*args)

    def send_header(self, h, v):
        self.send(self._header_msg(h, v))

    def send_headers(self, it):
        b = StringIO()
        map(lambda h: b.write(self._header_msg(h[0], h[1])), it)
        b.write(CRLF)
        b.write(CRLF)
        self.send(b.getvalue())

    def send_first_line(self, hmethod, uri, protocol="HTTP/1.1"):
        self.send("%s %s %s%s" % (hmethod, uri, protocol, CRLF))

    def request(self, method, uri, protocol="HTTP/1.1", headers={}, data=None,
                close=False):
        h = {"Host": self.host}
        h.update(headers)

        if close:
            h['Connection'] = 'close'

        self.send_first_line(method, uri, protocol)
        self.send_headers(h.iteritems())

        if not data is None:
            if isinstance(data, types.GeneratorType):
                for d in data:
                    self.send(d)
            else:
                self.send(data)

        r = read_response(self.sock, method)

        r[0].update({'method': method, 'uri': uri, 'req_headers': h})

        if close:
            self.sock.close()
            self.closed = True
        return r

    def close(self):
        '''Send a host request for a not found URI
        with 'Connection: close' Header'''
        r = self.request('HEAD', NOT_FOUND_URI, close=True)
        return r

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


def test_uri(c, uri, **kwargs):
    return c.request('HEAD', uri, **kwargs)


def test_uris(c, uris, **kwargs):
    close = kwargs.pop('close', False)

    for uri in uris:
        yield test_uri(c, uri)

    if close:
        c.close()


def run_test_uris(host, port, uris):
    with SimpleHttpClient((host, int(port)), host) as c:
        r = test_uris(c, uris)
        for _ in r:
            yield _

def uris_from_stdin():
    for l in sys.stdin:
        l = l.strip()
        if l:
            yield l

def main(uris=None):
    if uris is None:
        uris = uris_from_stdin()

    args = sys.argv[1:]
    args.append(uris)

    for i, d in run_test_uris(*args):
        print("%s|%s" % (i['status'][0], i['uri']))


if __name__ == '__main__':
    main()