#!/usr/bin/env python3
"""Ollama shim v2 — obfuscated relay on nightmare (replaces v1, same systemd unit).
- Binds tailnet-only 100.103.34.0:28443 (nonstandard port, nothing fingerprints as Ollama)
- Bearer token required (X-Auth-Token); constant-time compare
- Generic 404 + fake nginx server header on anything else (no Ollama fingerprint)
- Proxies /gen -> localhost:11434/api/generate, returns ONE non-chunked body
"""
import json
import hmac
import http.server, socketserver, urllib.request

LOCAL = "http://127.0.0.1:11434"
BIND = ("100.103.34.0", 28443)
TOKEN = open("/opt/ollama-shim/shared_secret.txt").read().strip()

class H(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.0"
    server_version = "nginx"
    sys_version = ""

    def _authed(self):
        got = self.headers.get("X-Auth-Token", "")
        return hmac.compare_digest(got, TOKEN)

    def _deny(self):
        body = b'{"error":"not found"}'
        self.send_response(404)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if not self._authed():
            return self._deny()
        if self.path == "/tags":
            body = urllib.request.urlopen(LOCAL + "/api/tags", timeout=10).read()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        else:
            self._deny()

    def do_POST(self):
        if not self._authed():
            return self._deny()
        if self.path != "/gen":
            return self._deny()
        n = int(self.headers.get("Content-Length", 0))
        payload = json.loads(self.rfile.read(n) or b"{}")
        payload["stream"] = False
        req = urllib.request.Request(LOCAL + "/api/generate",
                                     data=json.dumps(payload).encode(),
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=600) as r:
            body = r.read()
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *a):
        pass

class TS(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True

if __name__ == "__main__":
    TS(BIND, H).serve_forever()