8
\$\begingroup\$

It's an extremely simply HTTP server in Python using the socket library, and a few others to get the MIME type, etc...

I've also avoided the ../../ vulnerability, although some of the code in the send_file function seems a bit weak.

It should also be PEP8 compliant, aside from maybe some trailing whitespace in comments.

import filetype
import socket
import _thread


class ServerSocket():
    ''' Recieves connections, parses the request and ships it off to a handler
    '''
    def __init__(self, address, handler=None, *args, **kwargs):
        ''' Creates a server socket and defines a handler for the server
        '''
        self.socket = socket.socket()
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.bind(address)

        self.handler_args = args
        self.handler_kwargs = kwargs

        if handler:
            self.handler = handler  # The custom handler
        else:
            self.handler = Handler  # The default handler

    def initialise(self, open_connections=5):
        ''' Initilises the server socket and has it listen for connections
        '''
        self.socket.listen(open_connections)
        self.listen()

    def parse(self, data):
        ''' Splits a packet into
                the request,
                the headers (which includes the request),
                and contents
        '''
        stringed = str(data, 'utf-8')
        split = stringed.split('\r\n\r\n')
        headers = split[0]
        if len(split) > 1:
            content = split[1]
        else:
            content = []
        request = headers.split(' ')[0]

        return request, headers, content

    def handle(self, client, address):
        ''' Parses the data and handles the request. It then closes the connection
        '''
        try:
            data = client.recv(1024)
        except ConnectionResetError:
            if self.handler_kwargs["logging"] is True:
                print(f'{address[0]} unexpectedly quit')
            client.close()
            return
        parsed = self.parse(data)
        handler = self.handler(self.handler_args, self.handler_kwargs)
        handler.handle(client, parsed, address)
        client.close()

    def listen(self):
        ''' Listens until a keyboard interrupt and handles each connection in a
                new thread
        '''
        try:
            while True:
                client_data = self.socket.accept()
                if self.handler_kwargs['logging'] is True:
                    print(f'Connection from {client_data[1][0]}')
                _thread.start_new_thread(self.handle, client_data)
        except KeyboardInterrupt:
            self.socket.close()


class Handler():
    ''' Handles requests from the Server Socket
    '''
    def __init__(self, args, kwargs):
        self.args = args
        self.kwargs = kwargs

    def set_status(self, code, message):
        ''' Used to add a status line:
                - 'HTTP/1.1 200 OK' or 'HTTP/1.1 404 Not Found', etc...
        '''
        self.reply_headers = [f'HTTP/1.0 {code} {message}']

    def set_header(self, header, content):
        ''' Defines a custom header and adds it to the response
        '''
        self.reply_headers += [f'{header}: {content}']

    def response(self, content):
        ''' Adds to the content of the response
        '''
        if type(content) == str:
            self.reply_content += content.split('\n')
        else:
            self.reply_content += [content]

    def calculate_content_length(self):
        ''' Calculates the content length and adds it to the header
        '''
        length = len(self.reply_content) * 2
        lengths = [len(line) for line in self.reply_content]
        length += sum(lengths)
        self.set_header('Content-Length', length)

    def get_type(self, file_name):
        return filetype.guess('./public/'+file_name)

    def extract_file_name(self, file_name=None):
        if file_name:
            f_name = file_name[1:]
        else:
            f_name = self.request_status.split(' ')[1][1:]
        return f_name

    def send_file(self, file_name=None):
        if file_name is None:
            file_name = self.extract_file_name()

        if file_name == '':
            file_name = 'index.html'
        elif file_name[0] in './':
            self.set_status(403, "Forbidden")
            self.set_header('Content-Type', 'text/html')
            self.reply_content = ['<p>Error 403: Forbidden</p>']
            return

        try:
            with open('./public/'+file_name, 'rb') as file:
                file_contents = file.read()
        except FileNotFoundError:
            self.set_status(404, 'Not Found')
            self.set_header('Content-Type', 'text/html')
            self.reply_content = ['<p>Error 404: File not found</p>']
            return
        file_type = self.get_type(file_name)
        if file_type is not None:
            self.set_header('Content-Type', file_type.MIME)
        elif file_name.split('.')[-1] == 'html':
            self.set_header('Content-Type', 'text/html')
        else:
            self.set_header('Content-Type', 'text/txt')
        self.response(file_contents)

    def get_request_address(self):
        return self.address

    def parse_headers(self, headers):
        t = {}
        for header in headers[1:]:
            t[header.split(': ')[0]] = header.split(': ')[1]
        return t

    def reply(self):
        ''' Assembles the response and sends it to the client
        '''
        if self.reply_headers[0][0:4] != "HTTP":
            self.set_status(200, 'OK')
            self.set_header('Content-Type', 'text/html')
            self.reply_content = ['<p>Response Status unspecified</p>']

        self.calculate_content_length()

        message = '\r\n'.join(self.reply_headers)
        message += '\r\n\r\n'
        try:
            message += '\r\n'.join(self.reply_content)
            message += '\r\n'
        except TypeError:
            message = bytes(message, 'utf-8')
            message += b'\r\n'.join(self.reply_content)
            message += b'\r\n'

        try:
            if type(message) == str:
                self.client.send(bytes(message, 'utf-8'))
            else:
                self.client.send(message)
        except:
            pass

    def handle(self, client, parsed_data, address):
        ''' Initialises variables and case-switches the request type to
                determine the handler function
        '''
        self.client = client
        self.address = address
        self.reply_headers = []
        self.reply_content = []
        self.headers = True
        self.request_status = parsed_data[1].split('\r\n')[0]
        request = parsed_data[0]
        headers = self.parse_headers(parsed_data[1].split('\r\n'))
        contents = parsed_data[2]

        if request == "GET":
            func = self.get
        elif request == "POST":
            func = self.post
        elif request == "HEAD":
            func = self.head
        elif request == "PUT":
            func = self.put
        elif request == "DELETE":
            func = self.delete
        elif request == "CONNECT":
            func = self.connect
        elif request == "OPTIONS":
            func = self.options
        elif request == "TRACE":
            func = self.trace
        elif request == "PATCH":
            func = self.patch
        else:
            func = self.default

        func(headers, contents)
        self.reply()

    def default(self, headers, contents):
        ''' If the request is not known, defaults to this
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Unknown Request Type</p>''')

    def get(self, headers, contents):
        ''' Overwrite to customly handle GET requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a GET Request</p>''')

    def post(self, headers, contents):
        ''' Overwrite to customly handle POST requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a POST Request</p>''')

    def head(self, headers, contents):
        ''' Overwrite to customly handle HEAD requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a HEAD Request</p>''')

    def put(self, headers, contents):
        ''' Overwrite to customly handle PUT requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a PUT Request</p>''')

    def delete(self, headers, contents):
        ''' Overwrite to customly handle DELETE requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a DELETE Request</p>''')

    def connect(self, headers, contents):
        ''' Overwrite to customly handle CONNECT requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a CONNECT Request</p>''')

    def options(self, headers, contents):
        ''' Overwrite to customly handle OPTIONS requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got an OPTIONS Request</p>''')

    def trace(self, headers, contents):
        ''' Overwrite to customly handle TRACE requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a TRACE Request</p>''')

    def patch(self, headers, contents):
        ''' Overwrite to customly handle PATCH requests
        '''
        self.set_status(200, 'OK')
        self.set_header('Content-Type', 'text/html')
        self.response('''<p>Successfully got a PATCH Request</p>''')


# =======================================================================

if __name__ == "__main__":
    import sys

    class CustomHandler(Handler):
        def get(self, headers, contents):
            self.set_status(200, 'OK')
            request_address = self.get_request_address()[0]
            file_name = self.extract_file_name()
            print(f'{request_address} -> {headers["Host"]}/{file_name}')
            self.send_file()

    def run():
        if len(sys.argv) == 2:
            port = int(sys.argv[1])
        else:
            port = 80

        try:
            print('Initialising...', end='')
            http_server = ServerSocket(
                ('0.0.0.0', port),
                CustomHandler,
                logging=True
            )
            print('Done')
            http_server.initialise()
        except Exception as e:
            print(f'{e}')

    run()

\$\endgroup\$

1 Answer 1

9
\$\begingroup\$

ServerSocket

self.handler

The handler evaluation in __init__ can be accomplished with an or ternary operation. It's clearer to the reader as to what's going on. Also, the name could be changed to HandlerClass, since it represents a class rather than an instance:

    self.HandlerClass = handler or Handler

self.handler logging

The check against handler_kwargs makes the code a bit difficult to follow, since the SocketServer is now in charge of something that arguably the handler should be doing. If the server is who should be doing the logging, then I would leave the logging check out of the handler_kwargs entirely. Store self.logging as a boolean setting on the server instance, and just check against that:

# for traceback handling
import traceback 

class SocketServer:
    def __init__(self, address, handler=None, *args, **kwargs):
        ~skipping some code~

        # this is so you don't unpack logging from kwargs since
        # it looks like you need it in your handler
        self.logging = kwargs.get('logging', False)

        ...

    def handle(self):
        try:
            data = client.recv(1024)
        except ConnectionResetError as e:
            # use the more pythonic `if bool` check
            # here, rather than comparing against a singleton
            if self.logging:
                print(f'{address[0]} unexpectedly quit: {e}', file=sys.stderr)
                traceback.print_exception(*sys.exc_info(), file=sys.stderr)     

I've added a sys.stderr stream to your print statement, and added the exception to your print statement. I've also added a traceback print which points to stderr.

handle Refactor

The client.close(); return statement can also be refactored using try/except's else feature and a finally block, since you always want the client to close

        try:
            data = client.recv(1024)
        except ConnectionResetError as e:
            if self.logging:
                print(f'{address[0]} unexpectedly quit: {e}', file=sys.stderr)
                traceback.print_exception(*sys.exc_info(), file=sys.stderr) 
        # this is if no exception fired
        else:
            parsed = self.parse(data)
            handler = self.HandlerClass()
            handler.handle(client, parsed, address)
        # this will always execute
        finally:
            client.close()

Convert bytes to str

parse

First, when converting bytes to str, use bytes.decode() rather than str(byte_obj, encoding). The headers and content evaluation can be handled using argument unpacking, and from there you can use an or expression with any on contents to either take the first result or create an empty list:

    def parse(self, data):
        ''' Splits a packet into
                the request,
                the headers (which includes the request),
                and contents
        '''
        stringed = data.decode().split('\r\n\r\n')

        # The headers and content blocks can be handled also with argument
        # unpacking and a ternary operator on content:

        headers, *content = data.decode().split('\r\n\r\n')
        content = content[0] if content else []

        request = headers.split(' ')[0]

        return request, headers, content

Handler

Adding to lists

There are lots of cases where you do this:

self.some_list += [some_value]

Just use .append(some_value), it's faster since you don't have to create a list just to add the value to an existing list. For example:

    def set_header(self, header, content):
        ''' Defines a custom header and adds it to the response
        '''
        # this is much quicker, do this instead
        self.reply_headers.append(f'{header}: {content}')

type-checking

Use isinstance rather than type(object) == some_type:

# This
    def response(self, content):
        ''' Adds to the content of the response
        '''
        if type(content) == str:
            self.reply_content += content.split('\n')
        else:
            self.reply_content += [content]

# should be this
    def response(self, content):
        if isinstance(content, str):
            self.reply_content.extend(content.split('\n'))
        else:
            self.reply_content.append(content)

Note that I'm also switching the addition of lists to appropriate calls to append and extend. Though, looking through the rest of your code, you only use this against str or bytes types, so I'd refactor to the following:

    def response(self, content):
        # force everything to bytes-type for convenience, that way
        # you never have to worry about TypeErrors later
        content = content if isinstance(content, bytes) else content.encode()

        self.reply_content.append(content)

content-length calculation

Here, you can drop the creation of the lengths list, and just use sum on map:

    def calculate_content_length(self):
        ''' Calculates the content length and adds it to the header
        '''
        length = len(self.reply_content) * 2

        # sum will take any iterable, even generators and maps
        # len in map is the function to be applied to each element
        length += sum(map(len, self.reply_content))
        self.set_header('Content-Length', length)

It's faster and doesn't build as many objects inside the function

Magic Numbers

In your extract_file_name function, you are slicing the file name from the first element, though it's not completely clear why:

    def extract_file_name(self, file_name=None):
        if file_name:
            f_name = file_name[1:]
        else:
            f_name = self.request_status.split(' ')[1][1:]
        return f_name

This is usually a code smell and you should include a docstring and/or comments to explain why you slice that way. Otherwise, the index is a "magic number" and can be difficult for you or others to maintain later. You also never use the file_name argument at any point when calling this function, so I might just leave it out.

Sending Files

Most of the improvements here follow the ones that have been suggested above:

Ternary Operation or Use or for file_name

file_name = file_name or self.extract_file_name()

## or 

file_name = file_name if file_name else self.extract_file_name()

The latter mirrors your if statement a bit more

Checking an empty string

Use if some_string not if some_string == ''.

Checking string start and end values

The str.startswith and str.endswith methods will help here, and they avoid indexing or slicing which can impair readability:

if file_name.startswith('./'):
    # do something


if file_name.endswith('html'):
    # do something

The logic in send_file could be cleaned up a bit. First, I think a handle_error method would do nicely to clean up some of the repeated code where you handle exceptions:

    def handle_error(self, code, short_reason, reason):
        """
        code: integer error code
        short_reason: string denoting the short error reason
        reason: string denoting the full error reason

        example:
            self.handle_error(404, 'Not Found', 'File Not Found')
        """
        self.set_status(code, short_reason)
        self.set_header('Content-Type', 'text/html')
        self.reply_content = [f'<p>Error {code}: {reason} </p>']  

Next, I think the filename checking can be refactored to make a bit more sense. First, it seems a bit counterintuitive that a path that might start with '.' is forbidden. What about ./public/index.html? If you're trying to avoid folder traversal such as paths with . and .. in them, it might not be that bad to use a regex. For example, what if I tried to give you a path like `../../root_file.txt'? It would pass your test, even though it will traverse back directories. I would do something like the following:

import re
~skipping lots of code~

    @staticmethod
    def valid_path(path):
        """
        Will take any unix-path and check for any '.' and '..' directories
        in it. Example:

            import re
            re_path = re.compile('^\.+$')

            some_path = '/root/path/to/../folder/./file.txt'
            public_path = 'public/folder/../../../file.txt'

            next(filter(re_path.match, some_path.split('/')))
            # '..'

            '/'.join(filter(re_path.match, some_path.split('/')))
            # '..'
            valid_path = '/path/to/file.txt'
            next(filter(re_path.match, valid_path.split('/')))
            StopIteration
        """
        path = path.lstrip('.').lstrip('/')
        re_path = re.compile('^\.+$')
        try:
            match = next(filter(re_path.match, path.split('/')))
        except StopIteration:
            return True
        return False

This way you can check if there are backout paths without colliding with something inocuous like ./public/file.txt, and you can handle the following FileNotFoundError. Otherwise, you'll get a None on return, and then you can return an AccessDenied:

    def send_file(self, file_name=None):
        ~snip~
        # this is more pythonic than doing
        # if file_name == ''
        file_name = file_name or 'index.html'

        # will either be True or False
        if not self.valid_path(file_name):
            # I've added the keywords here for clarity
            self.handle_error(403, short_reason='Forbidden', reason='Forbidden')
            return  

        try:
            # use and f-string here rather than string
            # concatenation
            with open(f'./public/{file_name}', 'rb') as fh:
                file_contents = file.read()
        except FileNotFoundError:
            self.handle_error(404, short_reason='Not Found', reason='File Not Found')
            return

        file_type = self.get_type(file_name)
        if file_type is not None:
            self.set_header('Content-Type', file_type.MIME)
        # use str.endswith here
        elif file_name.endswith('html'):
            self.set_header('Content-Type', 'text/html')
        else:
            self.set_header('Content-Type', 'text/txt')

        self.response(file_contents)

Redundant Methods

The method get_request_address in my opinion doesn't need to be there. Just access self.address wherever you call it.

parse_headers

If it were up to me, I'd make the headers a dictionary from the get-go, but this can be re-factored like so:

    # this can be static since you never access any self attributes
    @staticmethod
    def parse_headers(headers):
        t = {}
        for header in headers[1:]:
            # unpack k and v from one call to split
            k, v = header.split(': ')
            t[k] = v
        return t

Or, even more succinctly

    @staticmethod
    def parse_headers(headers):
        return dict(header.split(': ') for header in headers[1:])

reply

Again, I'd force everything to bytes. You try to append to a str (message) with mixed types, which will fail:

    def reply(self):
        ''' Assembles the response and sends it to the client
        '''
        if not self.reply_headers[0].startswith("HTTP"):
            self.set_status(200, 'OK')
            self.set_header('Content-Type', 'text/html')
            self.reply_content = ['<p>Response Status unspecified</p>']

        self.calculate_content_length()

        # here's how to coerce to bytes on reply_headers
        message = b'\r\n'.join(map(str.encode, self.reply_headers))
        message += b'\r\n\r\n'


        message += b'\r\n'.join(x.encode() if isinstance(x, str) else x for x in self.reply_content))
        message += b'\r\n'

        # Now you don't have to type-check
        try:
            self.client.send(message)
        except:
            pass

Avoiding the type-checking entirely makes your code more streamlined and maintainable.

Last, instead of using pass on an unexpected Exception, I'd raise some sort of 500:

     try:
         self.client.send(message)
     except Exception as e:
         traceback.print_exception(*sys.exc_info(), file=sys.stderr)
         self.handle_error(500, 'Server Error', 'Internal Server Error')

Then at the very least there's some visibility either by the client or by you that something bad happened.

handle

This is a classic case of don't repeat yourself. Looking at the big block of if statements, your execution times will be slightly worse for a PATCH than a GET. To refactor, I would use a dictionary that binds the names of functions to operations, yielding a constant time lookup:

    def __init__(self, ...):
        ~snip~
        self.functions = {
            'GET': self.get,
            'POST': self.post,
            'HEAD': self.head,
            'PUT': self.put,
            'DELETE': self.delete,
            'CONNECT': self.connect,
            'OPTIONS': self.options,
            'TRACE': self.trace,
            'PATCH': self.patch
        }


    def call_method(self, contents, request_type=None, headers=None):
        headers = headers or {}

        headers['Content-Type'] = headers.get('Content-Type', 'text/html')

        self.set_status(200, 'OK')

        for k, v in headers.items():
            self.set_header(k, v)

        func = self.functions.get(request_type.upper() if request_type else '', self.default)

        func(headers, contents)
        self.reply()
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.