120 lines
3.4 KiB
Python
Executable File
120 lines
3.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
User-space allowlist TCP proxy for the 111 Ollama fallback.
|
|
|
|
Purpose:
|
|
- Keep the real Ollama server bound to 127.0.0.1:11434.
|
|
- Expose 192.168.0.111:11434 only to approved momo-pro clients.
|
|
- Reject noisy LAN / VM clients without requiring sudo/pfctl.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import ipaddress
|
|
import logging
|
|
import os
|
|
import signal
|
|
from typing import Iterable
|
|
|
|
|
|
LISTEN_HOST = os.getenv("OLLAMA111_PROXY_LISTEN_HOST", "192.168.0.111")
|
|
LISTEN_PORT = int(os.getenv("OLLAMA111_PROXY_LISTEN_PORT", "11434"))
|
|
TARGET_HOST = os.getenv("OLLAMA111_PROXY_TARGET_HOST", "127.0.0.1")
|
|
TARGET_PORT = int(os.getenv("OLLAMA111_PROXY_TARGET_PORT", "11434"))
|
|
ALLOWED_CIDRS = tuple(
|
|
item.strip()
|
|
for item in os.getenv(
|
|
"OLLAMA111_PROXY_ALLOWED_CIDRS",
|
|
"127.0.0.1/32,192.168.0.80/32,192.168.0.111/32,192.168.0.188/32",
|
|
).split(",")
|
|
if item.strip()
|
|
)
|
|
|
|
|
|
def _allowed_networks() -> tuple[ipaddress._BaseNetwork, ...]:
|
|
return tuple(ipaddress.ip_network(item, strict=False) for item in ALLOWED_CIDRS)
|
|
|
|
|
|
ALLOWED_NETWORKS = _allowed_networks()
|
|
|
|
|
|
def _is_allowed(peer_ip: str) -> bool:
|
|
try:
|
|
ip = ipaddress.ip_address(peer_ip)
|
|
except ValueError:
|
|
return False
|
|
return any(ip in network for network in ALLOWED_NETWORKS)
|
|
|
|
|
|
async def _pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
while True:
|
|
data = await reader.read(65536)
|
|
if not data:
|
|
break
|
|
writer.write(data)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
try:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def _handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
peer = writer.get_extra_info("peername")
|
|
peer_ip = peer[0] if peer else "unknown"
|
|
if not _is_allowed(peer_ip):
|
|
logging.warning("reject peer=%s allowed=%s", peer_ip, ",".join(ALLOWED_CIDRS))
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
return
|
|
|
|
try:
|
|
target_reader, target_writer = await asyncio.open_connection(TARGET_HOST, TARGET_PORT)
|
|
except Exception as exc:
|
|
logging.error("target connect failed peer=%s target=%s:%s err=%r", peer_ip, TARGET_HOST, TARGET_PORT, exc)
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
return
|
|
|
|
logging.info("proxy peer=%s -> %s:%s", peer_ip, TARGET_HOST, TARGET_PORT)
|
|
await asyncio.gather(
|
|
_pipe(reader, target_writer),
|
|
_pipe(target_reader, writer),
|
|
)
|
|
|
|
|
|
async def _main() -> None:
|
|
logging.basicConfig(
|
|
level=os.getenv("OLLAMA111_PROXY_LOG_LEVEL", "INFO"),
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
)
|
|
server = await asyncio.start_server(_handle_client, LISTEN_HOST, LISTEN_PORT)
|
|
sockets = ", ".join(str(sock.getsockname()) for sock in (server.sockets or []))
|
|
logging.info(
|
|
"ollama111 allow proxy listening=%s target=%s:%s allowed=%s",
|
|
sockets,
|
|
TARGET_HOST,
|
|
TARGET_PORT,
|
|
",".join(ALLOWED_CIDRS),
|
|
)
|
|
|
|
stop = asyncio.Event()
|
|
loop = asyncio.get_running_loop()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
loop.add_signal_handler(sig, stop.set)
|
|
|
|
async with server:
|
|
await stop.wait()
|
|
server.close()
|
|
await server.wait_closed()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(_main())
|