Proxy Server using python socket library

mail@pastecode.io avatar
unknown
python
2 years ago
5.0 kB
1
Indexable
Never
import os
import sys
import socket
import _thread

def main():
    try:
        # Create a new proxy server that will sit between browser and remote server
        serverSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        serverSocket.bind(('', 5000))
        serverSocket.listen(50)


    except socket.error as e:
        print(e)
        sys.exit()

    while True:
        try:
            # create connection between proxy server and browser client
            (client_socket, address) = serverSocket.accept()

            # start a new thread for each client
            _thread.start_new_thread(client_connection, (client_socket, address))
        except Exception:
            serverSocket.close()
            sys.exit()



# returns headers based on cache expiration that will be sent to remote server
def getHeaders(host, if_modify=False, modify_date=None):
    if (if_modify):
        return f'GET / HTTP/1.1\r\nAccept: */*\r\nCache-Control: no-cache\r\nAccept-Encoding: utf8\r\nHost: {host}\r\nIf-Modified-Since: {modify_date}\r\nConnection: keep-alive\r\n\r\n'.encode()
    else:
        return f'GET / HTTP/1.1\r\nAccept: */*\r\nCache-Control: no-cache\r\nAccept-Encoding: utf8\r\nHost: {host}\r\nConnection: keep-alive\r\n\r\n'.encode()

# returns a name (str) for a url's cache file
def filename_for_url(url):
    ret = url.replace(".", "_")
    ret = ret.replace("/", "-")
    ret += '.txt'
    return ret

# saves cache
def saveCache(filename, content=b''):
    with open(filename, 'w') as file:
        file.write(content.decode())

# checks if cache exists. if exists, returns the cache
def checkCache(filename):
    if (os.path.isfile(filename)):
        with open(filename, 'r') as file:
            return True, file.read()
    return False, None

# This handles a single client
def client_connection(client_socket: socket.socket, address):
    # receive request from browser
    req = client_socket.recv(8192)

    # Get the destination address from request
    destination_address = req.decode().split("\n")[0].split(" ")[1]

    # Shorten the destination address i.e., no http://
    t = destination_address.find("://")
    short_url = destination_address
    if (t != -1):
        short_url = destination_address[t+3:]

    # Separate port if exists
    port = 80
    if (short_url.find(":") != -1):
        short_url, port = short_url.split(":")
    if (short_url.find("/") != -1):
        short_url = short_url.split("/")[0]

    # Choose a filename based on url
    filename = filename_for_url(short_url)
    cache_exists, cache = checkCache(filename)

    try:
        ss = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        ss.connect((short_url, int(port)))

        if (cache_exists):
            print("cache exists")

            # Check the last modified date of cached file
            data = cache.split("\n")
            last_modify_date = None
            for i in range(len(data)):
                if (data[i].find("Last-Modified:") != -1):
                    last_modify_date = data[i][15:]
            headers = getHeaders(short_url, True, last_modify_date)

            # send request to remote server to check if the cache is still valid
            ss.send(headers)
            res = ss.recv(8192)
            res_dec = res.decode()

            # Check the response status
            # If the cache is still valid send the cached files to browser
            if (int(res_dec.split('/r/n')[0].split(" ")[1]) == 304):
                print("Not modified sending cache")
                client_socket.sendall(cache.encode())

            # cache is not valid, so send new file from remote server
            else:
                print("Sending Modified version")
                client_socket.send(res)

                # if the body is not completed, below code will wait for rest of the body to come
                response = res
                while True:
                    res = ss.recv(8192)
                    if (len(res) > 0):
                        response += res
                        client_socket.send(res)
                    else:
                        break

                # update the cache
                saveCache(filename, response)

        # If no cache exists then fetch file from remote server
        else:
            print('No cache found')
            headers = getHeaders(short_url)
            ss.send(headers)

            response = b''
            while True:
                res = ss.recv(8192)
                if (len(res) > 0):
                    response += res
                    client_socket.send(res)
                else:
                    break

            # create new cache for this site
            saveCache(filename, response)

    except socket.error as e:
        print(e)
    finally:
        if (client_socket):
            client_socket.close()
        if (ss):
            ss.close()

if __name__ == '__main__':
    main()