import json
import os
import sys

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

MAX_TOKENS = 128

DECODE_PAIRS = (
    ("%25", "%"),
    ("%2e", "."),
    ("%2E", "."),
    ("%2f", "/"),
    ("%2F", "/"),
    ("%5c", "\\\\"),
    ("%5C", "\\\\"),
    ("%27", "'"),
    ("%20", " "),
    ("%3c", "<"),
    ("%3C", "<"),
    ("%3e", ">"),
    ("%3E", ">"),
    ("%00", ""),
)


def decode_3x(value):
    for _ in range(3):
        for needle, replacement in DECODE_PAIRS:
            value = value.replace(needle, replacement)
    return value


def build_request_text(method, path, query, user_agent):
    return f"{method} {decode_3x(path or '')}?{decode_3x(query or '')} [SEP] {user_agent or ''}"


def load(model_dir):
    tokenizer = Tokenizer.from_file(os.path.join(model_dir, "tokenizer.json"))
    tokenizer.enable_truncation(max_length=MAX_TOKENS)
    tokenizer.enable_padding(pad_id=tokenizer.token_to_id("[PAD]"), pad_token="[PAD]")
    session = ort.InferenceSession(os.path.join(model_dir, "model.onnx"), providers=["CPUExecutionProvider"])
    with open(os.path.join(model_dir, "thresholds.json"), encoding="utf-8") as handle:
        thresholds = json.load(handle)
    return tokenizer, session, thresholds


def score(tokenizer, session, texts):
    encodings = tokenizer.encode_batch(texts)
    feeds = {
        "input_ids": np.array([e.ids for e in encodings], dtype=np.int64),
        "attention_mask": np.array([e.attention_mask for e in encodings], dtype=np.int64),
    }
    logits = session.run(None, feeds)[0].astype(np.float64)
    logits -= logits.max(axis=1, keepdims=True)
    probs = np.exp(logits)
    return (probs[:, 1] / probs.sum(axis=1)).tolist()


def band(value, thresholds):
    if value >= thresholds["dangerous_threshold"]:
        return "dangerous"
    if value >= thresholds["risky_threshold"]:
        return "risky"
    return "none"


REQUESTS = [
    ("GET", "/", "", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.7390.65 Safari/537.36"),
    ("GET", "/download", "file=..%252f..%252f..%252f..%252fetc%252fpasswd", "curl/8.5.0"),
    ("GET", "/wp-content/plugins/hello/wso.php", "", "Mozilla/5.0"),
]


def main():
    model_dir = sys.argv[1] if len(sys.argv) > 1 else os.path.dirname(os.path.abspath(__file__))
    tokenizer, session, thresholds = load(model_dir)
    texts = [build_request_text(*request) for request in REQUESTS]
    print(f"risky >= {thresholds['risky_threshold']:.4f}, dangerous >= {thresholds['dangerous_threshold']:.4f}")
    for text, value in zip(texts, score(tokenizer, session, texts)):
        print(f"{value:.4f}  {band(value, thresholds):9}  {text[:90]}")


if __name__ == "__main__":
    main()