# Written by retoor@molodetz.nl
# This script is a network proxy that intercepts HTTP communication, replaces specific content in headers and data, and passes communication between clients and an upstream server. It supports chunked transfer encoding, keep-alive connections, and forking for client connections.
# Import summary: Uses socket, asyncio, pathlib, os, signal, and json modules from Python standard library.
# MIT License
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF, OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
import socket
import asyncio
import pathlib
import os
import signal
import json
from concurrent.futures import ThreadPoolExecutor as Executor
def set_header_value(headers, key, value):
value = str(value)
headers = headers.decode("utf-8").rstrip("\r\n\r\n")
parts = headers.split(key + ": ")
if len(parts) > 1:
headers = headers.replace(parts[1].split("\r\n")[0], value)
else:
headers = headers + key + ": " + value + "\r\n"
return (headers + "\r\n").encode()
def get_header_value(headers, key):
headers = headers.decode("utf-8")
try:
parts = headers.split(key + ": ")
return parts[1].split("\r\n")[0]
except:
return None
def get_content_length(headers):
try:
return int(get_header_value(headers, "Content-Length"))
except:
return 0
def send_all(sock, data):
while data:
sent = sock.send(data)
data = data[sent:]
class HTTPDocument:
def __init__(self, headers, data):
self.headers = headers
self.data = data
self.original_content_length = get_content_length(headers)
self.content_length = self.original_content_length
def replace(self, old, new):
print("Replaced")
self.data = self.data.replace(old.encode(), new.encode())
if self.original_content_length != len(self.data):
self.headers = set_header_value(self.headers, "Content-Length", len(self.data))
return True
return False
@property
def content(self):
return self.data
class Protocol:
def __init__(self, document, downstream, upstream):
self.upstream = upstream
self.downstream = downstream
self.document = document
self.bytes_sent = 0
def stream(self):
while self.bytes_sent != self.document.content_length:
chunk = self.downstream.recv(1)
self.bytes_sent += 1
self.upstream.sendall(chunk)
def read_until(sock, delim):
data = b""
try:
while True:
d = sock.recv(1)
data += d
if data.endswith(delim):
return data
except Exception as ex:
print(ex)
return None
def communicate(sock, config):
upstream = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
upstream.connect(("127.0.0.1", 8082))
while True:
headers = read_until(sock, b"\r\n\r\n")
print(headers)
if not headers:
sock.close()
break
upstream.sendall(headers)
content_length = get_content_length(headers)
if content_length:
doc = HTTPDocument(headers, b'')
protocol = Protocol(doc, sock, upstream)
protocol.stream()
headers = read_until(upstream, b"\r\n\r\n")
for key, value in config.get("upstream", {}).items():
headers = headers.replace(key.encode(), value.encode())
print(headers)
sock.sendall(headers)
content_length = get_content_length(headers)
handled = content_length > 0
if content_length:
doc = HTTPDocument(headers, b'')
protocol = Protocol(doc, upstream, sock)
protocol.stream()
if get_header_value(headers, "Connection") == "close":
sock.close()
return
if get_header_value(headers, "Upgrade") == "websocket":
sock.close()
return
if get_header_value(headers, "Transfer-Encoding") == "chunked":
while True:
headers = read_until(upstream, b"\r\n")
content_length = int(headers[:-2], 16)
if not content_length:
sock.sendall(b"\r\n")
break
data = b''
while len(data) < content_length:
chunk = upstream.recv(1)
data += chunk
for key, value in config.get("downstream", {}).items():
data = data.replace(key.encode(), value.encode())
content_length = len(data)
data = hex(content_length)[2:].encode() + b"\r\n" + data
data += upstream.recv(2)
sock.sendall(data)
if not data:
break
print(data)
upstream.close()
sock.close()
break
if not handled:
upstream.close()
sock.close()
break
if get_header_value(headers, "Connection") == "keep-alive":
continue
break
def reap_zombie_processes():
while True:
try:
pid, _ = os.waitpid(-1, os.WNOHANG)
if pid == 0:
break
except ChildProcessError:
break
async def serve(host, port, config):
config = json.loads(pathlib.Path(config).read_text())
signal.signal(signal.SIGCHLD, lambda signum, frame: reap_zombie_processes())
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
executor = Executor(100)
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind((host, port))
server.listen()
clients = []
while True:
client, address = server.accept()
clients.append(client)
print(f"Connection from {address}")
pid = os.fork()
if pid == 0:
server.close()
communicate(client, config)
os._exit(0)
else:
client.close()
asyncio.run(serve("0.0.0.0", 3046, ".replace.json"))
Reference in New Issue View Git Blame Copy Permalink