| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel, DistilBertModel |
| from transformers.modeling_outputs import SequenceClassifierOutput |
| from .configuration_fourthwall import FourthWallConfig |
|
|
| LABEL_COLS = [ |
| '000 - Normal', |
| '272 - Protocol Manipulation', |
| '242 - Code Injection', |
| '88 - OS Command Injection', |
| '126 - Path Traversal', |
| '66 - SQL Injection', |
| '16 - Dictionary-based Password Attack', |
| '310 - Scanning for Vulnerable Software', |
| '153 - Input Data Manipulation', |
| '274 - HTTP Verb Tampering', |
| '194 - Fake the Source of Data', |
| '34 - HTTP Response Splitting', |
| '33 - HTTP Request Smuggling' |
| ] |
|
|
| BLOCK_CLASSES = { |
| '242 - Code Injection', '88 - OS Command Injection', '126 - Path Traversal', |
| '66 - SQL Injection', '16 - Dictionary-based Password Attack', |
| '153 - Input Data Manipulation', '274 - HTTP Verb Tampering', |
| '194 - Fake the Source of Data', '34 - HTTP Response Splitting', |
| '33 - HTTP Request Smuggling' |
| } |
|
|
| DEFAULT_UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" |
|
|
| class FourthWallForSequenceClassification(PreTrainedModel): |
| config_class = FourthWallConfig |
| base_model_prefix = "bert" |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.num_labels = config.num_labels |
| self.config = config |
| |
| |
| self.bert = DistilBertModel(config) |
| self.classifier = nn.Sequential( |
| nn.Dropout(config.seq_classif_dropout), |
| nn.Linear(config.dim, config.num_labels) |
| ) |
| self.post_init() |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| head_mask=None, |
| inputs_embeds=None, |
| labels=None, |
| output_attentions=None, |
| output_hidden_states=None, |
| return_dict=None, |
| ): |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
| outputs = self.bert( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| head_mask=head_mask, |
| inputs_embeds=inputs_embeds, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=return_dict, |
| ) |
|
|
| hidden_state = outputs[0] |
| cls_output = hidden_state[:, 0] |
| logits = self.classifier(cls_output) |
|
|
| loss = None |
| if labels is not None: |
| loss_fct = nn.BCEWithLogitsLoss() |
| loss = loss_fct(logits, labels) |
|
|
| if not return_dict: |
| output = (logits,) + outputs[1:] |
| return ((loss,) + output) if loss is not None else output |
|
|
| return SequenceClassifierOutput( |
| loss=loss, |
| logits=logits, |
| hidden_states=outputs.hidden_states if output_hidden_states else None, |
| attentions=outputs.attentions if output_attentions else None, |
| ) |
|
|
| @classmethod |
| def serialize_request(cls, method="GET", path="/", ua="", ct="", body="", max_body_chars=500): |
| """Serialize raw HTTP components into model input format.""" |
| parts = [] |
| m = str(method).strip().upper() if method and str(method).lower() not in ('nan', 'none', '') else "GET" |
| parts.append(f"[METHOD] {m}") |
| |
| p = str(path).strip() if path and str(path).lower() not in ('nan', 'none', '') else "/" |
| if not p.startswith("/"): |
| p = "/" + p |
| parts.append(f"[PATH] {p}") |
| |
| ua_str = str(ua).strip() if ua and str(ua).lower() not in ('nan', 'none', '') else DEFAULT_UA |
| parts.append(f"[UA] {ua_str}") |
| |
| if ct and str(ct).lower() not in ('nan', 'none', ''): |
| parts.append(f"[CT] {str(ct).strip()}") |
| |
| if body and str(body).lower() not in ('nan', 'none', ''): |
| b_str = str(body).strip()[:max_body_chars] |
| parts.append(f"[BODY] {b_str}") |
| |
| return " ".join(parts) |
|
|
