sysreptor/api/src/reportcreator_api/pentests/customfields/utils.py

import dataclasses
import enum
import random
from lorem_text import lorem
from typing import Any, Iterable, Optional, Union, OrderedDict
from django.utils import timezone

from reportcreator_api.pentests.customfields.types import FieldDataType, FieldDefinition, FieldOrigin
from reportcreator_api.utils.utils import is_date_string, is_uuid
from reportcreator_api.utils.error_messages import format_path


def contains(a, b):
    """
    Checks if dict a contains dict b recursively
    """
    if not b:
        return True
        
    if type(a) != type(b):
        return False

    for k, v in b.items():
        if k not in a:
            return False
        if isinstance(v, dict):
            if not contains(a[k], v):
                return False
        elif isinstance(v, (list, tuple)):
            raise ValueError('Cannot diff lists')
        elif v != b[k]:
            return False
    return True





def has_field_structure_changed(old: dict[str, FieldDefinition], new: dict[str, FieldDefinition]):
    if set(old.keys()) != set(new.keys()):
        return True
    
    for k in old.keys():
        field_type = old[k].type
        if field_type != new[k].type:
            return True
        elif field_type == FieldDataType.OBJECT and has_field_structure_changed(old[k].properties, new[k].properties):
            return True
        elif field_type == FieldDataType.LIST and has_field_structure_changed({'items': old[k].items}, {'items': new[k].items}):
            return True
        elif field_type == FieldDataType.ENUM and set(map(lambda c: c.value, old[k].choices)) - set(map(lambda c: c.value, new[k].choices)):
            # Existing enum choice was removed
            return True
    
    return False


class HandleUndefinedFieldsOptions(enum.Enum):
    FILL_NONE = 'fill_none'
    FILL_DEFAULT = 'fill_default'
    FILL_DEMO_DATA = 'fill_demo_data'


def _default_or_demo_data(definition: FieldDefinition, demo_data: Any, handle_undefined: HandleUndefinedFieldsOptions):
    if handle_undefined == HandleUndefinedFieldsOptions.FILL_NONE:
        return None
    elif handle_undefined == HandleUndefinedFieldsOptions.FILL_DEFAULT:
        return definition.default
    elif handle_undefined == HandleUndefinedFieldsOptions.FILL_DEMO_DATA:
        return definition.default or demo_data


def ensure_defined_structure(value, definition: Union[dict[str, FieldDefinition], FieldDefinition], handle_undefined: HandleUndefinedFieldsOptions = HandleUndefinedFieldsOptions.FILL_DEFAULT, include_undefined=False):
    """
    Ensure that the returned data is valid for the given field definition.
    Recursively check for undefined fields and set a value.
    Returns only data of defined fields, if value contains undefined fields, this data is not returned.
    """
    if isinstance(definition, dict):
        out = value if include_undefined else {}
        for k, d in definition.items():
            out[k] = ensure_defined_structure(value=(value if isinstance(value, dict) else {}).get(k), definition=d, handle_undefined=handle_undefined)
        return out
    else:
        if definition.type == FieldDataType.OBJECT:
            return ensure_defined_structure(value, definition.properties, handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.LIST:
            if isinstance(value, list):
                return [ensure_defined_structure(value=e, definition=definition.items, handle_undefined=handle_undefined) for e in value]
            else:
                if handle_undefined == HandleUndefinedFieldsOptions.FILL_DEMO_DATA and definition.items.type != FieldDataType.USER:
                    return [ensure_defined_structure(value=None, definition=definition.items, handle_undefined=handle_undefined) for _ in range(2)]
                else:
                    return []
        elif definition.type == FieldDataType.MARKDOWN and not isinstance(value, str):
            return _default_or_demo_data(definition, lorem.paragraphs(3), handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.STRING and not isinstance(value, str):
            return _default_or_demo_data(definition, lorem.words(2), handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.CVSS and not isinstance(value, str):
            return _default_or_demo_data(definition, 'n/a', handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.ENUM and not (isinstance(value, str) and value in {c.value for c in definition.choices}):
            return _default_or_demo_data(definition, next(iter(map(lambda c: c.value, definition.choices)), None), handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.COMBOBOX and not isinstance(value, str):
            return _default_or_demo_data(definition, next(iter(definition.suggestions), None), handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.DATE and not (isinstance(value, str) and is_date_string(value)):
            return _default_or_demo_data(definition, timezone.now().date().isoformat(), handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.NUMBER and not isinstance(value, (int, float)):
            return _default_or_demo_data(definition, random.randint(1, 10), handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.BOOLEAN and not isinstance(value, bool):
            return _default_or_demo_data(definition, random.choice([True, False]), handle_undefined=handle_undefined)
        elif definition.type == FieldDataType.USER and not (isinstance(value, str) or is_uuid(value)):
            return None
        else:
            return value


def check_definitions_compatible(a: Union[dict[str, FieldDefinition], FieldDefinition], b: Union[dict[str, FieldDefinition], FieldDefinition], path: Optional[tuple[str]] = None) -> tuple[bool, list[str]]:
    """
    Check if definitions are compatible and values can be converted without data loss.
    """
    path = path or tuple()
    valid = True
    errors = []
    if isinstance(a, dict) and isinstance(b, dict):
        for k in set(a.keys()).intersection(b.keys()):
            res_valid, res_errors = check_definitions_compatible(a[k], b[k], path=path + tuple([k]))
            valid = valid and res_valid
            errors.extend(res_errors)
    elif isinstance(a, FieldDefinition) and isinstance(b, FieldDefinition):
        if a.type != b.type:
            valid = False
            errors.append(f'Field "{format_path(path)}" has different types: "{a.type.value}" vs. "{b.type.value}"')
        elif a.type == FieldDataType.LIST:
            res_valid, res_errors = check_definitions_compatible(a.items, b.items, path=path + tuple(['[]']))
            valid = valid and res_valid
            errors.extend(res_errors) 
        elif a.type == FieldDataType.ENUM:
            missing_choices = {c.value for c in a.choices} - {c.value for c in b.choices}
            if missing_choices:
                valid = False
                missing_choices_str = ', '.join(map(lambda c: f'"{c}"', missing_choices))
                errors.append(f'Field "{format_path(path)}" has missing enum choices: {missing_choices_str}')
    return valid, errors


def set_field_origin(definition: Union[dict[str, FieldDefinition], FieldDefinition], predefined_fields: Union[dict, FieldDefinition, None]):
    """
    Sets definition.origin recursively
    """
    if isinstance(definition, dict):
        out = {}
        for k, d in definition.items():
            out[k] = set_field_origin(d, predefined_fields=predefined_fields.get(k) if predefined_fields else None)
        return out
    else:
        out = dataclasses.replace(definition, origin=getattr(predefined_fields, 'origin', FieldOrigin.CUSTOM))
            
        if out.type == FieldDataType.OBJECT:
            out.properties = set_field_origin(out.properties, predefined_fields=getattr(predefined_fields, 'properties', None))
        elif out.type == FieldDataType.LIST:
            out.items = set_field_origin(out.items, predefined_fields=getattr(predefined_fields, 'items', None))
        return out


def iterate_fields(value: Union[dict, Any], definition: Union[dict[str, FieldDefinition], FieldDefinition], path: Optional[tuple[str]] = None) -> Iterable[tuple[tuple[str], Any, FieldDefinition]]:
    """
    Recursively iterate over all defined fields
    """
    if not definition:
        return

    path = path or tuple()
    if isinstance(definition, dict):
        for k, d in definition.items():
            yield from iterate_fields(value=(value if isinstance(value, dict) else {}).get(k), definition=d, path=path + tuple([k]))
    else:
        # Current field
        yield path, value, definition

        # Nested structures
        if definition.type == FieldDataType.OBJECT:
            yield from iterate_fields(value=value or {}, definition=definition.properties, path=path)
        elif definition.type == FieldDataType.LIST:
            for idx, v in enumerate(value if isinstance(value, list) else []):
                yield from iterate_fields(value=v, definition=definition.items, path=path + tuple(['[' + str(idx) + ']']))