wifi-gps-mapper-public/cracklab/server/services/method_definition.py

from __future__ import annotations

import re
from typing import Any


class MethodDefinitionError(ValueError):
    pass


_RESOURCE_RE = re.compile(r"^[a-z][a-z0-9_-]*:[A-Za-z0-9][A-Za-z0-9._-]*$")
_STEP_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
_MASK_CHARSET_SIZES = {
    "?l": 26,
    "?u": 26,
    "?d": 10,
    "?s": 33,
    "?a": 95,
    "?b": 256,
}
_MASK_MIN_LENGTH = 8


def validate_method_definition(
    definition: dict[str, Any],
) -> dict[str, Any]:
    if not isinstance(definition, dict):
        raise MethodDefinitionError(
            "Method definition must be a JSON object"
        )

    if definition.get("schema_version") != 1:
        raise MethodDefinitionError(
            "Unsupported method definition schema version"
        )

    hash_mode = definition.get("hash_mode")

    if not isinstance(hash_mode, int) or isinstance(hash_mode, bool):
        raise MethodDefinitionError(
            "Method definition hash_mode must be a non-negative integer"
        )

    if hash_mode < 0:
        raise MethodDefinitionError(
            "Method definition hash_mode must be a non-negative integer"
        )

    steps = definition.get("steps")

    if not isinstance(steps, list) or not steps:
        raise MethodDefinitionError(
            "Method definition must contain a non-empty steps list"
        )

    seen_step_ids: set[str] = set()

    for step in steps:
        _validate_step(step, seen_step_ids)

    return definition


def _validate_step(
    step: Any,
    seen_step_ids: set[str],
) -> None:
    if not isinstance(step, dict):
        raise MethodDefinitionError(
            "Each method step must be a JSON object"
        )

    step_id = step.get("step_id")

    if not isinstance(step_id, str) or not _STEP_ID_RE.fullmatch(step_id):
        raise MethodDefinitionError(
            "Each method step must have a valid step_id"
        )

    if step_id in seen_step_ids:
        raise MethodDefinitionError(
            f"Duplicate step_id: {step_id}"
        )

    seen_step_ids.add(step_id)

    attack_mode = step.get("attack_mode")

    if not isinstance(attack_mode, int) or isinstance(attack_mode, bool):
        raise MethodDefinitionError(
            f"Invalid attack_mode for step: {step_id}"
        )

    if attack_mode < 0:
        raise MethodDefinitionError(
            f"Invalid attack_mode for step: {step_id}"
        )

    if attack_mode == 0:
        _require_dictionary(step, step_id)

    elif attack_mode == 1:
        _require_dictionaries(step, step_id)

    elif attack_mode == 3:
        _require_mask(
            step,
            step_id,
            minimum_length=_MASK_MIN_LENGTH,
        )

    elif attack_mode == 6:
        _require_dictionary(step, step_id)
        _require_mask(
            step,
            step_id,
            minimum_length=_MASK_MIN_LENGTH,
        )

    elif attack_mode == 7:
        _require_mask(
            step,
            step_id,
            minimum_length=_MASK_MIN_LENGTH,
        )
        _require_dictionary(step, step_id)

    if "rules" in step:
        _validate_resource_object(
            step["rules"],
            "rules",
            step_id,
        )


def _require_dictionary(
    step: dict[str, Any],
    step_id: str,
) -> None:
    if "dictionary" not in step:
        raise MethodDefinitionError(
            f"Step {step_id} requires dictionary"
        )

    _validate_resource_object(
        step["dictionary"],
        "dictionary",
        step_id,
    )


def _require_dictionaries(
    step: dict[str, Any],
    step_id: str,
) -> None:
    dictionaries = step.get("dictionaries")

    if not isinstance(dictionaries, list) or len(dictionaries) != 2:
        raise MethodDefinitionError(
            f"Step {step_id} requires exactly two dictionaries"
        )

    for dictionary in dictionaries:
        _validate_resource_object(
            dictionary,
            "dictionary",
            step_id,
        )


def _require_mask(
    step: dict[str, Any],
    step_id: str,
    minimum_length: int | None = None,
) -> None:
    mask = step.get("mask")

    if not isinstance(mask, str) or not mask:
        raise MethodDefinitionError(
            f"Step {step_id} requires a non-empty mask"
        )

    length = 0
    index = 0

    while index < len(mask):
        if mask[index] != "?":
            length += 1
            index += 1
            continue

        if index + 1 >= len(mask):
            raise MethodDefinitionError(
                f"Step {step_id} mask contains a lone '?' "
                f"at position {index + 1}"
            )

        token = mask[index:index + 2]

        if token == "??":
            length += 1
            index += 2
            continue

        if token not in _MASK_CHARSET_SIZES:
            raise MethodDefinitionError(
                f"Step {step_id} contains unsupported mask token "
                f"{token!r} at position {index + 1}"
            )

        length += 1
        index += 2

    if (
        minimum_length is not None
        and length < minimum_length
    ):
        raise MethodDefinitionError(
            f"Step {step_id} mask must produce at least "
            f"{minimum_length} characters; got {length}"
        )

def _validate_resource_object(
    value: Any,
    resource_type: str,
    step_id: str,
) -> None:
    if not isinstance(value, dict):
        raise MethodDefinitionError(
            f"Step {step_id} {resource_type} must be an object"
        )

    resource = value.get("resource")

    if not isinstance(resource, str) or not _RESOURCE_RE.fullmatch(resource):
        raise MethodDefinitionError(
            f"Step {step_id} {resource_type}.resource must be "
            "a valid logical resource ID"
        )

    expected_prefix = (
        "wordlist:"
        if resource_type == "dictionary"
        else f"{resource_type}:"
    )

    if not resource.startswith(expected_prefix):
        raise MethodDefinitionError(
            f"Step {step_id} {resource_type}.resource must start "
            f"with {expected_prefix!r}"
        )