Black_Raven
144d749678
Extractor extracts data from responses: - Supports JSON path extraction - Supports header and status extraction - Saves extracted values to Context
32 lines
1.0 KiB
Python
32 lines
1.0 KiB
Python
from typing import Any
|
|
from .context import Context
|
|
from .itm import ExtractModel
|
|
from .executor import Response
|
|
|
|
|
|
class Extractor:
|
|
def __init__(self, context: Context):
|
|
self.context = context
|
|
|
|
def extract(self, configs: list[ExtractModel], response: Response):
|
|
for config in configs:
|
|
value = self._get_value(response, config)
|
|
self.context.set_extracted(config.name, value)
|
|
|
|
def _get_value(self, response: Response, config: ExtractModel) -> Any:
|
|
if config.source == "status":
|
|
return response.status
|
|
if config.source == "header":
|
|
return response.headers.get(config.json_path)
|
|
return self._json_path(response.body, config.json_path)
|
|
|
|
def _json_path(self, data: dict, path: str) -> Any:
|
|
parts = path.lstrip("$.").split(".")
|
|
current = data
|
|
for part in parts:
|
|
if isinstance(current, dict):
|
|
current = current.get(part)
|
|
else:
|
|
return None
|
|
return current
|