nxdns/website/demo/gen/rewrite_history.py

"""Spread the bulk-generated query log over the previous day with a household
rhythm, then rebuild every projection from the rewritten rows.

Run only while nxdns is stopped. usage: rewrite_history.py DATA_DIR
"""
import datetime
import random
import sqlite3
import sys
import time

DATA = sys.argv[1]
NOW = int(time.time())
START = NOW - 24 * 3600 + 120
END = NOW - 300
CREATED_AT = START - 1200

# Relative activity per local hour of day, per device.
CURVES = {
    "192.168.50.41": [2, 1, 1, 1, 1, 1, 3, 20, 25, 8, 6, 6, 10, 8, 6, 8, 15, 45, 80, 100, 100, 90, 60, 20],  # TV
    "192.168.50.52": [6, 3, 3, 3, 3, 5, 25, 70, 60, 35, 35, 40, 60, 45, 35, 40, 45, 60, 70, 80, 85, 80, 60, 25],  # iPhone
    "192.168.50.53": [5, 3, 3, 3, 3, 4, 15, 60, 70, 40, 40, 45, 65, 50, 40, 45, 50, 65, 75, 80, 75, 70, 50, 20],  # Pixel
    "192.168.50.54": [1, 1, 1, 1, 1, 1, 4, 35, 20, 2, 2, 2, 4, 2, 2, 40, 90, 100, 85, 80, 60, 25, 4, 2],  # Kids' phone
    "192.168.50.20": [2, 1, 1, 1, 1, 1, 1, 4, 10, 15, 15, 15, 20, 15, 15, 20, 30, 50, 80, 100, 100, 85, 50, 15],  # Desktop PC
    "192.168.50.21": [2, 1, 1, 1, 1, 1, 2, 15, 60, 100, 100, 95, 60, 90, 100, 95, 80, 40, 25, 25, 20, 12, 5, 3],  # MacBook Air
}


def weight(ip, ts):
    return CURVES[ip][datetime.datetime.fromtimestamp(ts).hour]


def sample_times(ip, n):
    minutes = list(range(START, END, 60))
    weights = [weight(ip, m) for m in minutes]
    return sorted(m + random.randrange(60) for m in random.choices(minutes, weights, k=n))


def lookups(rows):
    """Group a client's rows into lookups: the A/AAAA/HTTPS burst of one name."""
    groups, last = [], None
    for row_id, ts, domain_id in rows:
        if last and last[0] == domain_id and ts - last[1] <= 1:
            groups[-1].append(row_id)
        else:
            groups.append([row_id])
        last = (domain_id, ts)
    return groups


# The repository's own bucket recompute (src/storage/repositories/queries_repo.zig,
# `recompute_bucket_sql`), verbatim; ?1 is the bucket start, ?2 its exclusive end.
RECOMPUTE_BUCKET_SQL = [
    """INSERT INTO bucket_totals (bucket, queries, blocked, cached, rt_sum, rt_count)
SELECT ?1, count(*), coalesce(sum(blocked <> 0), 0), coalesce(sum(cache_hit = 1), 0),
       coalesce(sum(response_time_us), 0), count(response_time_us)
  FROM query_log
 WHERE timestamp >= ?1 AND timestamp < ?2
HAVING count(*) > 0""",
    """INSERT INTO bucket_clients (bucket, client_ip, queries)
SELECT ?1, client_ip, count(*)
  FROM query_log
 WHERE timestamp >= ?1 AND timestamp < ?2
 GROUP BY client_ip""",
    """INSERT INTO bucket_types (bucket, qtype, count)
SELECT ?1, coalesce(qtype, -1), count(*)
  FROM query_log
 WHERE timestamp >= ?1 AND timestamp < ?2
 GROUP BY coalesce(qtype, -1)""",
    """INSERT INTO bucket_routes (bucket, route_kind, source_present, source_text, count)
SELECT ?1, route_kind, source IS NOT NULL, coalesce(source, ''), count(*)
  FROM (SELECT route_kind,
               CASE route_kind
                    WHEN 'upstream' THEN upstream
                    WHEN 'forward_zone' THEN forward_zone
               END AS source
          FROM query_log
         WHERE timestamp >= ?1 AND timestamp < ?2)
 GROUP BY route_kind, source""",
]
GRAIN = 1800

COLUMNS = (
    "timestamp, domain_id, client_ip, qtype, blocked, response_time_us, cache_hit, upstream, qclass, rcode,"
    " group_id, group_name, policy_action, policy_reason, matched, source_id, source_name, cname_target,"
    " safe_search_target, route_kind, forward_zone"
)

ql = sqlite3.connect(f"{DATA}/querylog.db", isolation_level=None)
ql.execute("BEGIN IMMEDIATE")
new_ts = {}
for (ip,) in ql.execute("SELECT DISTINCT client_ip FROM query_log").fetchall():
    rows = ql.execute(
        "SELECT id, timestamp, domain_id FROM query_log WHERE client_ip = ? ORDER BY timestamp, id", (ip,)
    ).fetchall()
    groups = lookups(rows)
    for group, ts in zip(groups, sample_times(ip, len(groups))):
        for row_id in group:
            new_ts[row_id] = ts
ql.execute("CREATE TEMP TABLE remap (old_id INTEGER PRIMARY KEY, ts INTEGER NOT NULL)")
ql.executemany("INSERT INTO remap VALUES (?, ?)", new_ts.items())
rest = COLUMNS.split(", ", 1)[1]
ql.execute(f"CREATE TEMP TABLE rewritten AS SELECT r.ts AS timestamp, {rest}, q.id AS old_id"
           " FROM query_log q JOIN remap r ON r.old_id = q.id")
before = ql.execute("SELECT count(*) FROM query_log").fetchone()[0]
ql.execute("DELETE FROM query_log")
ql.execute(f"INSERT INTO query_log ({COLUMNS}) SELECT {COLUMNS} FROM rewritten ORDER BY timestamp, old_id")
after = ql.execute("SELECT count(*) FROM query_log").fetchone()[0]
assert before == after == len(new_ts), (before, after, len(new_ts))

for table in ("bucket_totals", "bucket_clients", "bucket_types", "bucket_routes"):
    ql.execute(f"DELETE FROM {table}")
buckets = [b for (b,) in ql.execute(f"SELECT DISTINCT timestamp / {GRAIN} * {GRAIN} FROM query_log")]
for bucket in buckets:
    for sql in RECOMPUTE_BUCKET_SQL:
        ql.execute(sql, (bucket, bucket + GRAIN))
ql.execute("UPDATE querylog_meta SET created_at = ?, available_since = ? WHERE id = 1", (CREATED_AT, CREATED_AT + 1))
ql.execute("COMMIT")
ql.execute("PRAGMA wal_checkpoint(TRUNCATE)")
span = ql.execute("SELECT min(timestamp), max(timestamp), count(*) FROM query_log").fetchone()
print(f"querylog: {span[2]} rows from {span[0]} to {span[1]}, {len(buckets)} buckets rebuilt")

cfg = sqlite3.connect(f"{DATA}/config.db", isolation_level=None)
cfg.execute("ATTACH ? AS ql", (f"{DATA}/querylog.db",))
cfg.execute("BEGIN IMMEDIATE")
cfg.execute("""UPDATE clients SET
    first_seen = coalesce((SELECT min(timestamp) FROM ql.query_log WHERE client_ip = clients.ip), first_seen),
    last_seen  = coalesce((SELECT max(timestamp) FROM ql.query_log WHERE client_ip = clients.ip), last_seen)""")
cfg.execute("UPDATE rules SET created_at = ?", (CREATED_AT,))
cfg.execute("COMMIT")
cfg.execute("PRAGMA wal_checkpoint(TRUNCATE)")
for row in cfg.execute("SELECT ip, name, first_seen, last_seen FROM clients ORDER BY ip"):
    print("client", row)