pgs/pgscript.py

import socket
import os
from datetime import datetime
import pathlib
import html
server_start = datetime.now()
import pgs
print(pgs.add(1,2))

def get_server_uptime():
    seconds = (datetime.now() - server_start).total_seconds()
    if seconds > 60*60: # hour 
        return f"{int(seconds/60/60)} hours"
    elif seconds > 60: # minute
        return f"{int(seconds/60)} minutes"
    return f"{seconds} seconds"

# Hostname, only resolved at startup.
hostname = socket.gethostname()

# Specify environment based on hostname.
env = "dev" if hostname.startswith("retoor") else "prod"

# This dict will contain the connections in this format: {downstream_fd:upstream_fd]
streams = {}

# This is debug variable. It holds the number of connections total made.
counter = 0

def is_ssh(header_bytes):
    return b'SSH' in header_bytes

def is_http(header_bytes):
    """
    Check if the header is an HTTP request.
    """
    return b'HTTP/1.1' in header_bytes or b'HTTP/1.0' in header_bytes or b'HTTP/2.0' in header_bytes or b'HTTP/3.0' in header_bytes or b'Connection: ' in header_bytes

def is_https(header_bytes):
    return not any([is_ssh(header_bytes), is_http(header_bytes)])

def on_connect(downstream):
    """
    This is a connection router which will be called by the server every 
    time a client connects. This function will be used to determine
    the upstream. The downstream and upstream are file descriptors.
    The upstream is not connected yet, it only holds a file descriptor.
    The connection will be made in this function. The connection will be 
    set non blocking after this function by pgs.
    This way of routing is so dynamic that you can:
     - Create [your-site].localhost redirects without configuring DNS.
     - Run multiple services on the same port. It is possible to run ssh 
       on the same port as your HTTPS server. This is a good idea in sense
       of security. Very unique, who does / expects that?
     - Rewrite urls in general.
     - Make clients always connect to the same upstream. Servers only have
       to manage their own session instead of having to communicate with
       redis.
     - You can inject headers in the request.
     - You can add HTTP Basic Authentication to protect all your services
       in a very early stage.
     - This server is quick, it can act as ddos protection.
     - You can make your server act as a load balancer.
     - You can make your server act as a reverse proxy.
     - You can apply rate limits.
     - You can cache responses.
     - You can implement a complete custom protocol here. Complete own 
       design. This feature will probably moved in the future.
     - You can do static file serving.
     - You can protect sensitive data not leaving the network by
       intercepting it.
     - You can call AI to make modications.
     - You can call databases.
     - You can save statistics.
    """

    global streams
    global counter

    counter += 1

    print("Connection nr.", counter)

    #u = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    #print("FD:",u.fileno())
    peek = pgs.read(downstream, 4096)

    redirect_to = []
    
    if pgs.is_ssh(peek):
        print("Forwarding to ssh molodetz")
        redirect_to = "molodetz.nl",22
        u = socket.fromfd(pgs.connect(*redirect_to), socket.AF_INET, socket.SOCK_STREAM)
    elif pgs.is_http(peek):
        
        if b'/random' in peek or b'random.' in peek:
            
            print("Forwarding to 127.0.0.1:3028.")
            peek = peek.replace(b'/random', b'/')
            peek = peek.replace(b'random.', b'')
            redirect_to = "127.0.0.1",3028
            u = socket.fromfd(pgs.connect(*redirect_to), socket.AF_INET, socket.SOCK_STREAM) 
        elif b'molodetz.local' in peek:
            print("Forwarding to 127.0.0.1:8082.")
            peek = peek.replace(b'molodetz.local', b'localhost')
            redirect_to = "127.0.0.1",8082
            u = socket.fromfd(pgs.connect(*redirect_to), socket.AF_INET, socket.SOCK_STREAM) 
        elif b'bench.local' in peek:
            print("Responding with bench page.")
            body = f"""<html>\n<head>\n<title>Benchmark page.</title>\n</head>\n<body>\n<h1>Bench</h1>\n<p>{counter}</p>\n</body>\n</html>\n""".encode()
            
            s = socket.fromfd(downstream, socket.AF_INET, socket.SOCK_STREAM)
            s.sendall(
                b'HTTP/1.1 200 Pretty Good Server.\r\n'
                + b'Content-Length: ' + str(len(body)).encode() + b'\r\n'
                + b'Content-Type: text/html\r\n\r\n'
                + body
            )
            u = None
            s.shutdown(socket.SHUT_RDWR)
        else:
            # 404
            if env == "prod":
                pgs.write(downstream,
                    b'HTTP/1.1 403 Authorization Required.\r\n\r\n')
            else:
                pgscript_source = html.escape(pathlib.Path(__file__).read_text())
                content = f"""<pre>
Server: Pretty Good Server
Environment: {env}
Total connections: {counter}
Local hostname: {hostname}
Downstream FD: {downstream}
Current time server: {datetime.now()}
Server started on: {server_start}
Server uptime: {get_server_uptime()}
</pre>
<h3>Source code of pgscript.py</h3>
<i>Location: {pathlib.Path(__file__).resolve()}</i>
<pre style="color:blue;">
{pgscript_source}
</pre>
                """
                body = f"""<html>\n<head>\n<title>Debug page.</title>\n</head>\n<body>\n<h1>Pretty Good Server</h1>\n<h3>Debugging information</h3>\n<p>{content}</p>\n</body>\n</html>\n"""
                headers = ["HTTP/1.1 200 Pretty Good Server.",
                    "Content-Length: {len(body)}",
                    "Content-Type: text/html",
                    ""
                ]
                headers = "\r\n".join(headers)
                #response = f"{headers}{body}" 
                response = f"{headers}\r\n\r\n{body}"
                pgs.write(downstream,response)
            
            # Unset socket so the server will close it.
            # Do not disconnect in python!
            # Instead of a 404, we also could've displayed a custom page.
            # Maybe some server statistics?
            u = None

    elif is_https(peek) and env == "prod":
        print("Forwarding to dev.to")
        redirect_to = "devrant.com", 443
        u = socket.fromfd(pgs.connect(*redirect_to), socket.AF_INET, socket.SOCK_STREAM) 
    else:
        # Error.
        print("Could not find upstream for header content.")
        print(f"Closing connection. Your current environment: {env}")
        # Don't have to close socket, pgs will do that himself.
        # Socket handling is done at one place to avoid race conditions.
        u = None

    if not u:
        return -1
    
    
    #os.write(upstream,peek)
    
    # Keep track of connections. Not sure if this is needed.
    upstream = u.fileno()
    
    # Remove reference to the socket so it doesn't get garbage collected.
    # This could break the connection. This way, it stays open.
    u = None 
    streams[downstream] = dict(upstream=upstream,upstream_host=redirect_to[0],upstream_port=redirect_to[1])
    streams[upstream] = dict(dowstream=downstream, upstream_host=redirect_to[0],upstream_port=redirect_to[1])

    # Return exact same value as what is given as parameter.
    return upstream

def on_headers(downstream, headers):
    stream = streams[downstream]
    if stream['upstream_host'] == b'devrant.com':
        headers = headers.replace(b'localhost', b'devrant.com')
        headers = headers.replace(b'molodetz.nl', b'devrant.com')
    if stream['upstream_host'] == b'molodetz.nl':
        headers = headers.replace(b'localhost', b'molodetz.nl')
    return headers