101 lines
3.0 KiB
Python
Executable File
101 lines
3.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import os
|
|
import selectors
|
|
import signal
|
|
import socket
|
|
import sys
|
|
|
|
|
|
LISTEN_HOST = os.getenv("LISTEN_HOST", "127.0.0.1")
|
|
LISTEN_PORT = int(os.getenv("LISTEN_PORT", "6380"))
|
|
REDIS_HOST = os.getenv("REDIS_HOST", "100.90.90.91")
|
|
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
|
|
REDIS_USER = os.getenv("REDIS_USER", "lisglosips")
|
|
CONNECT_TIMEOUT = float(os.getenv("CONNECT_TIMEOUT", "2.0"))
|
|
|
|
|
|
def credential(name: str) -> str:
|
|
cred_dir = os.getenv("CREDENTIALS_DIRECTORY")
|
|
if not cred_dir:
|
|
raise RuntimeError("CREDENTIALS_DIRECTORY is not set")
|
|
path = os.path.join(cred_dir, name)
|
|
with open(path, "r", encoding="utf-8") as handle:
|
|
return handle.read().strip()
|
|
|
|
|
|
def resp_array(parts: list[str]) -> bytes:
|
|
output = f"*{len(parts)}\r\n".encode("utf-8")
|
|
for part in parts:
|
|
data = part.encode("utf-8")
|
|
output += f"${len(data)}\r\n".encode("utf-8") + data + b"\r\n"
|
|
return output
|
|
|
|
|
|
def read_line(sock: socket.socket) -> bytes:
|
|
data = b""
|
|
while not data.endswith(b"\r\n"):
|
|
chunk = sock.recv(1)
|
|
if not chunk:
|
|
break
|
|
data += chunk
|
|
if len(data) > 4096:
|
|
break
|
|
return data
|
|
|
|
|
|
def connect_upstream() -> socket.socket:
|
|
password = credential("redis_password")
|
|
upstream = socket.create_connection((REDIS_HOST, REDIS_PORT), timeout=CONNECT_TIMEOUT)
|
|
upstream.sendall(resp_array(["AUTH", REDIS_USER, password]))
|
|
reply = read_line(upstream)
|
|
if not reply.startswith(b"+OK"):
|
|
upstream.close()
|
|
raise RuntimeError("Redis AUTH failed")
|
|
upstream.setblocking(False)
|
|
return upstream
|
|
|
|
|
|
def bridge(client: socket.socket) -> None:
|
|
upstream = connect_upstream()
|
|
client.setblocking(False)
|
|
selector = selectors.DefaultSelector()
|
|
selector.register(client, selectors.EVENT_READ, upstream)
|
|
selector.register(upstream, selectors.EVENT_READ, client)
|
|
try:
|
|
while True:
|
|
for key, _ in selector.select(timeout=30):
|
|
src = key.fileobj
|
|
dst = key.data
|
|
data = src.recv(65536)
|
|
if not data:
|
|
return
|
|
dst.sendall(data)
|
|
finally:
|
|
selector.close()
|
|
client.close()
|
|
upstream.close()
|
|
|
|
|
|
def main() -> int:
|
|
signal.signal(signal.SIGCHLD, signal.SIG_IGN)
|
|
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
server.bind((LISTEN_HOST, LISTEN_PORT))
|
|
server.listen(128)
|
|
print(f"lisglosips redis auth proxy listening on {LISTEN_HOST}:{LISTEN_PORT}", flush=True)
|
|
while True:
|
|
client, _ = server.accept()
|
|
pid = os.fork()
|
|
if pid == 0:
|
|
server.close()
|
|
try:
|
|
bridge(client)
|
|
except Exception as exc:
|
|
print(f"redis proxy client failed: {exc}", file=sys.stderr, flush=True)
|
|
return 0
|
|
client.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|