features.py

"""AdVig feature extractor.

~30 features, all computable with cheap integer ops + tiny lexicon lookups,
designed to port 1:1 to C on an ESP8266-class MCU. No allocations beyond
small fixed buffers needed at inference time.
"""
import math
import re

# multi-part public suffixes (subset covering top traffic)
MULTIPART_SUFFIXES = {
    "co.uk", "org.uk", "ac.uk", "gov.uk", "co.jp", "ne.jp", "or.jp", "ac.jp",
    "com.au", "net.au", "org.au", "edu.au", "gov.au", "co.nz", "net.nz",
    "org.nz", "co.in", "net.in", "org.in", "gov.in", "co.za", "org.za",
    "web.za", "com.br", "net.br", "org.br", "gov.br", "com.mx", "org.mx",
    "com.cn", "net.cn", "org.cn", "gov.cn", "com.hk", "org.hk", "com.sg",
    "com.my", "org.my", "com.tr", "org.tr", "co.kr", "or.kr", "com.ar",
    "com.co", "org.co", "com.pl", "org.pl", "com.ru", "org.ru", "net.ru",
    "com.ua", "org.ua", "net.ua", "com.tw", "org.tw", "com.vn", "com.ph",
    "co.id", "or.id", "com.pk", "org.pk", "com.bd", "co.th", "or.th",
    "com.sa", "com.eg", "com.ng", "com.gh", "co.ke", "com.il", "org.il",
    "co.at", "or.at", "ac.at", "co.hu", "org.hu", "co.ro", "org.ro",
    "com.es", "org.es", "com.pt", "org.pt", "co.it", "org.it", "com.gr",
    "co.be", "org.be", "co.nl", "org.nl", "co.dk", "co.no", "co.se",
    "com.fi", "com.ie", "co.cz", "co.cl", "com.pe", "com.uy", "com.ec",
    "com.do", "com.gt", "com.sv", "com.ni", "com.pa", "com.ve",
}

TRUSTED_TLDS = {"com", "org", "net", "edu", "gov", "mil", "int"}

ADHEAVY_TLDS = {
    "xyz", "top", "click", "link", "online", "site", "club", "icu", "buzz",
    "stream", "cfd", "sbs", "shop", "store", "fun", "space", "website",
    "live", "rest", "monster", "quest", "cyou", "cam", "bar", "gdn", "mom",
    "lol", "bond", "autos", "boat", "review", "country", "kim", "work",
    "bid", "trade", "webcam", "dating", "adult", "porn", "sex", "casino",
}

VENDOR_TOKENS = {
    "doubleclick", "googlesyndication", "googleadservices", "adnxs",
    "criteo", "criteo", "taboola", "outbrain", "scorecardresearch",
    "quantserve", "quantcast", "moatads", "pubmatic", "rubiconproject",
    "openx", "smartadserver", "zedo", "yieldmo", "adform", "adroll",
    "adcolony", "chartboost", "applovin", "inmobi", "mopub", "admob",
    "adservice", "adsystem", "adsense", "adsrvr", "amplitude",
}

BIGTECH_TOKENS = {
    "google", "googleapis", "facebook", "fbcdn", "instagram", "amazon",
    "awsstatic", "cloudfront", "microsoft", "apple", "icloud", "netflix",
    "twitter", "tiktok", "bing", "linkedin", "reddit", "wikipedia",
    "cloudflare", "akamai", "fastly", "youtube", "yahoo", "ebay",
}

FEATURE_NAMES = [
    "length", "label_count", "max_label_len", "digit_count", "max_digit_run",
    "hyphen_count", "entropy", "vowel_ratio", "starts_with_www",
    "has_punycode", "subdomain_depth", "tld_trusted", "tld_adheavy",
    "tld_is_cctld", "tld_length",
    "tok_ad", "tok_advert", "tok_banner", "tok_promo", "tok_sponsor",
    "tok_track", "tok_analytics", "tok_metrics", "tok_telemetry",
    "tok_beacon", "tok_pixel", "tok_tag", "tok_click", "tok_impression",
    "tok_affiliate", "tok_syndication", "tok_vendor",
    "tok_bigtech", "bigtech_and_adtoken",
]

_VOWELS = frozenset("aeiou")


def _registrable(domain: str) -> str:
    parts = domain.split(".")
    if len(parts) >= 3 and ".".join(parts[-2:]) in MULTIPART_SUFFIXES:
        return ".".join(parts[-3:])
    if len(parts) >= 2:
        return ".".join(parts[-2:])
    return domain


def extract_features(domain: str):
    d = domain.lower()
    n = len(d)

    labels = d.split(".")
    label_count = len(labels)
    max_label_len = max(len(x) for x in labels)
    digit_count = sum(ch.isdigit() for ch in d)
    hyphen_count = d.count("-")

    max_digit_run = run = 0
    for ch in d:
        if ch.isdigit():
            run += 1
            if run > max_digit_run:
                max_digit_run = run
        else:
            run = 0

    freq = {}
    for ch in d:
        freq[ch] = freq.get(ch, 0) + 1
    entropy = -sum((c / n) * math.log2(c / n) for c in freq.values()) if n else 0.0
    vowels = sum(1 for ch in d if ch in _VOWELS)
    vowel_ratio = vowels / n if n else 0.0

    starts_with_www = float(d.startswith("www.") or ".www." in ("." + d))
    has_punycode = float("xn--" in d)

    reg = _registrable(d)
    reg_parts = reg.split(".")
    subdomain_depth = label_count - len(reg_parts)

    tld = labels[-1]
    tld_trusted = float(tld in TRUSTED_TLDS)
    tld_adheavy = float(tld in ADHEAVY_TLDS)
    tld_is_cctld = float(len(tld) == 2 and not tld_trusted)
    tld_length = len(tld)

    segs = re.split(r"[._-]", d)
    segset = set(segs)

    def any_seg(pred):
        return float(any(pred(t) for t in segs))

    tok_ad = any_seg(lambda t: t in ("ad", "ads"))
    tok_advert = any_seg(lambda t: t.startswith("advert") or t.startswith("adver"))
    tok_banner = any_seg(lambda t: "banner" in t)
    tok_promo = any_seg(lambda t: t.startswith("promo"))
    tok_sponsor = any_seg(lambda t: "sponsor" in t)
    tok_track = any_seg(lambda t: t.startswith("track"))
    tok_analytics = any_seg(lambda t: "analytic" in t)
    tok_metrics = any_seg(lambda t: t.startswith("metric") or t.startswith("stat") or t == "st")
    tok_telemetry = any_seg(lambda t: t.startswith("telemetr"))
    tok_beacon = any_seg(lambda t: "beacon" in t)
    tok_pixel = any_seg(lambda t: "pixel" in t)
    tok_tag = any_seg(lambda t: t.startswith("tag"))
    tok_click = any_seg(lambda t: t.startswith("click") or t in ("clk", "clks"))
    tok_impression = any_seg(lambda t: t.startswith("impr"))
    tok_affiliate = any_seg(lambda t: t.startswith("affil") or t == "aff")
    tok_syndication = any_seg(lambda t: t.startswith("syndic"))
    tok_vendor = float(bool(segset & VENDOR_TOKENS))
    tok_bigtech = float(bool(segset & BIGTECH_TOKENS))
    ad_any = max(tok_ad, tok_advert, tok_track, tok_analytics, tok_vendor,
                 tok_banner, tok_beacon, tok_pixel)
    bigtech_and_adtoken = ad_any * tok_bigtech

    return [
        float(n), float(label_count), float(max_label_len),
        float(digit_count), float(max_digit_run), float(hyphen_count),
        entropy, vowel_ratio, starts_with_www, has_punycode,
        float(subdomain_depth), tld_trusted, tld_adheavy, tld_is_cctld,
        float(tld_length),
        tok_ad, tok_advert, tok_banner, tok_promo, tok_sponsor, tok_track,
        tok_analytics, tok_metrics, tok_telemetry, tok_beacon, tok_pixel,
        tok_tag, tok_click, tok_impression, tok_affiliate,
        tok_syndication, tok_vendor, tok_bigtech, bigtech_and_adtoken,
    ]


def extract_batch(domains):
    import numpy as np
    return np.array([extract_features(d) for d in domains], dtype=np.float32)


if __name__ == "__main__":
    tests = ["ads.doubleclick.net", "www.google.com", "stats.g.analytics.example.com",
             "mail.yahoo.co.uk", "pixel-tracking.adnxs.com"]
    for t in tests:
        feats = dict(zip(FEATURE_NAMES, extract_features(t)))
        keep = ["length", "subdomain_depth", "tok_ad", "tok_track", "tok_analytics",
                "tok_vendor", "tok_bigtech", "bigtech_and_adtoken", "tld_trusted"]
        print(f"{t:45s}", {k: feats[k] for k in keep})