zonemaster-tools/zmtools/server.py

"""HTTP front for the tools. Read-only GET endpoints returning JSON, mounted by
the gringo host nginx at https://zonemaster.es/tools-api/.

    GET /tools-api/propagation?domain=&type=
    GET /tools-api/resolvers?domain=&type=&resolver=
    GET /tools-api/security?domain=          (also records the scan)
    GET /tools-api/security/history?domain=
    GET /tools-api/infra?domain=
    GET /tools-api/graph?domain=
    GET /tools-api/blocklists?domain=
    GET /tools-api/blacklists?domain=
    GET /tools-api/health
"""

import json
import logging
import os
import sqlite3
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlsplit

from . import blacklists, blocklists, dnsutil, graph, infra, resolvers, security

log = logging.getLogger('zmtools')
HISTORY_DB = os.environ.get('ZMTOOLS_HISTORY_DB', '/var/db/zmtools/history.sqlite')

# Requests per minute per client, by cost. A blacklist check is ~80 queries.
COST = {'propagation': 2, 'resolvers': 2, 'security': 3, 'infra': 2, 'graph': 2,
        'blocklists': 1, 'blacklists': 3, 'history': 0}
BUDGET = 30          # cost units per client per minute
CACHE_TTL = {'propagation': 20, 'resolvers': 20, 'default': 120}
MAX_INFLIGHT = 12    # concurrent heavy requests, all clients together

_buckets, _cache = {}, {}
_state_lock = threading.Lock()
_inflight = threading.BoundedSemaphore(MAX_INFLIGHT)


def _allow(client, cost):
    now = time.time()
    with _state_lock:
        level, last = _buckets.get(client, (0.0, now))
        level = max(0.0, level - (now - last) * BUDGET / 60.0)
        if level + cost > BUDGET:
            _buckets[client] = (level, now)
            return False, int((level + cost - BUDGET) * 60 / BUDGET) + 1
        _buckets[client] = (level + cost, now)
        if len(_buckets) > 50000:
            _buckets.clear()
    return True, 0


def _cached(key, ttl, fn):
    now = time.time()
    with _state_lock:
        hit = _cache.get(key)
        if hit and hit[0] > now:
            return hit[1], True
    val = fn()
    with _state_lock:
        _cache[key] = (now + ttl, val)
        if len(_cache) > 5000:
            for k in [k for k, v in _cache.items() if v[0] <= now]:
                del _cache[k]
    return val, False


def _history_db():
    os.makedirs(os.path.dirname(HISTORY_DB), exist_ok=True)
    con = sqlite3.connect(HISTORY_DB, timeout=5)
    con.execute('CREATE TABLE IF NOT EXISTS scans (domain TEXT, at REAL, score INTEGER, '
                'grade TEXT, statuses TEXT)')
    con.execute('CREATE INDEX IF NOT EXISTS scans_domain ON scans (domain, at)')
    return con


def record_scan(result):
    statuses = {c['id']: c['status'] for c in result['checks']}
    con = _history_db()
    try:
        last = con.execute('SELECT at, statuses FROM scans WHERE domain=? ORDER BY at DESC LIMIT 1',
                           (result['domain'],)).fetchone()
        # One row per change, or per hour when nothing changed.
        if last and json.loads(last[1]) == statuses and time.time() - last[0] < 3600:
            return
        con.execute('INSERT INTO scans VALUES (?,?,?,?,?)',
                    (result['domain'], time.time(), result['score'], result['grade'],
                     json.dumps(statuses)))
        con.commit()
    finally:
        con.close()


def history(domain):
    con = _history_db()
    try:
        rows = con.execute('SELECT at, score, grade, statuses FROM scans WHERE domain=? '
                           'ORDER BY at DESC LIMIT 50', (domain,)).fetchall()
    finally:
        con.close()
    return {'domain': domain, 'scans': [
        {'at': r[0], 'score': r[1], 'grade': r[2], 'statuses': json.loads(r[3])} for r in rows]}


def _security(domain):
    result = security.scan(domain)
    try:
        record_scan(result)
    except sqlite3.Error:
        log.exception('could not record scan')
    return result


ROUTES = {
    'propagation': lambda q: resolvers.propagation(_domain(q, ptr_ok=True), q.get('type')),
    'resolvers': lambda q: resolvers.resolver_test(
        _domain(q), q.get('type'),
        dnsutil.public_ip(q['resolver']) if q.get('resolver') else None),
    'security': lambda q: _security(_domain(q)),
    'infra': lambda q: infra.detect(_domain(q)),
    'graph': lambda q: graph.build(_domain(q)),
    'blocklists': lambda q: blocklists.check(_domain(q)),
    'blacklists': lambda q: blacklists.check(_domain(q)),
}


def _domain(q, ptr_ok=False):
    raw = q.get('domain', '')
    if ptr_ok and dnsutil.is_public_ip(raw.strip()):
        return raw.strip()
    return dnsutil.normalize_domain(raw)


class Handler(BaseHTTPRequestHandler):
    server_version = 'zmtools/1.0'
    sys_version = ''

    def log_message(self, fmt, *args):
        log.info('%s %s', self._client(), fmt % args)

    def _client(self):
        return self.headers.get('X-Real-IP') or self.client_address[0]

    def _send(self, status, body, cache=0, extra=None):
        data = json.dumps(body, separators=(',', ':')).encode()
        self.send_response(status)
        self.send_header('Content-Type', 'application/json; charset=utf-8')
        self.send_header('Content-Length', str(len(data)))
        self.send_header('Cache-Control', f'public, max-age={cache}' if cache else 'no-store')
        self.send_header('X-Content-Type-Options', 'nosniff')
        for k, v in (extra or {}).items():
            self.send_header(k, v)
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        url = urlsplit(self.path)
        path = url.path.rstrip('/')
        if path.startswith('/tools-api'):
            path = path[len('/tools-api'):]
        q = {k: v[0][:300] for k, v in parse_qs(url.query).items()}
        tool = path.strip('/')
        if tool == 'health':
            return self._send(200, {'ok': True, 'blocklists': os.path.exists(blocklists.DB_PATH)})
        if tool == 'security/history':
            try:
                return self._send(200, history(dnsutil.normalize_domain(q.get('domain', ''))))
            except dnsutil.BadInput as e:
                return self._send(400, {'error': 'bad_input', 'detail': str(e)})
        fn = ROUTES.get(tool)
        if fn is None:
            return self._send(404, {'error': 'not_found'})
        ok, wait = _allow(self._client(), COST[tool])
        if not ok:
            return self._send(429, {'error': 'rate_limited', 'retry_after': wait},
                              extra={'Retry-After': str(wait)})
        if not _inflight.acquire(timeout=20):
            return self._send(503, {'error': 'busy'}, extra={'Retry-After': '10'})
        try:
            key = (tool, tuple(sorted(q.items())))
            ttl = CACHE_TTL.get(tool, CACHE_TTL['default'])
            t0 = time.monotonic()
            result, hit = _cached(key, ttl, lambda: fn(q))
            result = dict(result, cached=hit, elapsed_ms=round((time.monotonic() - t0) * 1000))
            self._send(200, result, cache=ttl)
        except dnsutil.BadInput as e:
            self._send(400, {'error': 'bad_input', 'detail': str(e)})
        except Exception:
            log.exception('tool %s failed for %r', tool, q)
            self._send(500, {'error': 'internal'})
        finally:
            _inflight.release()


def _refresher():
    while True:
        try:
            if blocklists.refresh():
                log.info('blocklists refreshed')
        except Exception:
            log.exception('blocklist refresh failed')
        time.sleep(3600)


def main():
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(name)s %(message)s')
    logging.getLogger('httpx').setLevel(logging.WARNING)  # one line per DoH probe otherwise
    host = os.environ.get('ZMTOOLS_HOST', '127.0.0.1')
    port = int(os.environ.get('ZMTOOLS_PORT', '5100'))
    threading.Thread(target=_refresher, daemon=True, name='refresh').start()
    srv = ThreadingHTTPServer((host, port), Handler)
    srv.daemon_threads = True
    log.info('listening on %s:%d', host, port)
    srv.serve_forever()


if __name__ == '__main__':
    main()