File size: 2,754 Bytes
f60cdbe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | 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()
|