tor-browser

app.py (5644B)

---

#!/usr/bin/env python

import asyncio
import http
import http.cookies
import pathlib
import signal
import urllib.parse
import uuid

import websockets
from websockets.frames import CloseCode


# User accounts database

USERS = {}


def create_token(user, lifetime=1):
   """Create token for user and delete it once its lifetime is over."""
   token = uuid.uuid4().hex
   USERS[token] = user
   asyncio.get_running_loop().call_later(lifetime, USERS.pop, token)
   return token


def get_user(token):
   """Find user authenticated by token or return None."""
   return USERS.get(token)


# Utilities


def get_cookie(raw, key):
   cookie = http.cookies.SimpleCookie(raw)
   morsel = cookie.get(key)
   if morsel is not None:
       return morsel.value


def get_query_param(path, key):
   query = urllib.parse.urlparse(path).query
   params = urllib.parse.parse_qs(query)
   values = params.get(key, [])
   if len(values) == 1:
       return values[0]


# Main HTTP server

CONTENT_TYPES = {
   ".css": "text/css",
   ".html": "text/html; charset=utf-8",
   ".ico": "image/x-icon",
   ".js": "text/javascript",
}


async def serve_html(path, request_headers):
   user = get_query_param(path, "user")
   path = urllib.parse.urlparse(path).path
   if path == "/":
       if user is None:
           page = "index.html"
       else:
           page = "test.html"
   else:
       page = path[1:]

   try:
       template = pathlib.Path(__file__).with_name(page)
   except ValueError:
       pass
   else:
       if template.is_file():
           headers = {"Content-Type": CONTENT_TYPES[template.suffix]}
           body = template.read_bytes()
           if user is not None:
               token = create_token(user)
               body = body.replace(b"TOKEN", token.encode())
           return http.HTTPStatus.OK, headers, body

   return http.HTTPStatus.NOT_FOUND, {}, b"Not found\n"


async def noop_handler(websocket):
   pass


# Send credentials as the first message in the WebSocket connection


async def first_message_handler(websocket):
   token = await websocket.recv()
   user = get_user(token)
   if user is None:
       await websocket.close(CloseCode.INTERNAL_ERROR, "authentication failed")
       return

   await websocket.send(f"Hello {user}!")
   message = await websocket.recv()
   assert message == f"Goodbye {user}."


# Add credentials to the WebSocket URI in a query parameter


class QueryParamProtocol(websockets.WebSocketServerProtocol):
   async def process_request(self, path, headers):
       token = get_query_param(path, "token")
       if token is None:
           return http.HTTPStatus.UNAUTHORIZED, [], b"Missing token\n"

       user = get_user(token)
       if user is None:
           return http.HTTPStatus.UNAUTHORIZED, [], b"Invalid token\n"

       self.user = user


async def query_param_handler(websocket):
   user = websocket.user

   await websocket.send(f"Hello {user}!")
   message = await websocket.recv()
   assert message == f"Goodbye {user}."


# Set a cookie on the domain of the WebSocket URI


class CookieProtocol(websockets.WebSocketServerProtocol):
   async def process_request(self, path, headers):
       if "Upgrade" not in headers:
           template = pathlib.Path(__file__).with_name(path[1:])
           headers = {"Content-Type": CONTENT_TYPES[template.suffix]}
           body = template.read_bytes()
           return http.HTTPStatus.OK, headers, body

       token = get_cookie(headers.get("Cookie", ""), "token")
       if token is None:
           return http.HTTPStatus.UNAUTHORIZED, [], b"Missing token\n"

       user = get_user(token)
       if user is None:
           return http.HTTPStatus.UNAUTHORIZED, [], b"Invalid token\n"

       self.user = user


async def cookie_handler(websocket):
   user = websocket.user

   await websocket.send(f"Hello {user}!")
   message = await websocket.recv()
   assert message == f"Goodbye {user}."


# Adding credentials to the WebSocket URI in user information


class UserInfoProtocol(websockets.BasicAuthWebSocketServerProtocol):
   async def check_credentials(self, username, password):
       if username != "token":
           return False

       user = get_user(password)
       if user is None:
           return False

       self.user = user
       return True


async def user_info_handler(websocket):
   user = websocket.user

   await websocket.send(f"Hello {user}!")
   message = await websocket.recv()
   assert message == f"Goodbye {user}."


# Start all five servers


async def main():
   # Set the stop condition when receiving SIGINT or SIGTERM.
   loop = asyncio.get_running_loop()
   stop = loop.create_future()
   loop.add_signal_handler(signal.SIGINT, stop.set_result, None)
   loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)

   async with websockets.serve(
       noop_handler,
       host="",
       port=8000,
       process_request=serve_html,
   ), websockets.serve(
       first_message_handler,
       host="",
       port=8001,
   ), websockets.serve(
       query_param_handler,
       host="",
       port=8002,
       create_protocol=QueryParamProtocol,
   ), websockets.serve(
       cookie_handler,
       host="",
       port=8003,
       create_protocol=CookieProtocol,
   ), websockets.serve(
       user_info_handler,
       host="",
       port=8004,
       create_protocol=UserInfoProtocol,
   ):
       print("Running on http://localhost:8000/")
       await stop
       print("\rExiting")


if __name__ == "__main__":
   asyncio.run(main())