devicehub-teal/ereuse_devicehub/parser/schemas.py

from datetime import datetime

from flask import current_app as app
from marshmallow import Schema as MarshmallowSchema
from marshmallow import ValidationError, pre_load, validates_schema
from marshmallow.fields import DateTime, Dict, Integer, List, Nested, String
from marshmallow_enum import EnumField

from ereuse_devicehub.resources.enums import Severity, SnapshotSoftware
from ereuse_devicehub.resources.schemas import Thing

# from marshmallow_enum import EnumField


class Snapshot_lite_data(MarshmallowSchema):
    hwmd_version = String(required=True)
    lshw = Dict(required=True)
    dmidecode = String(required=True)
    lspci = String(required=True)
    hwinfo = String(required=True)
    smart = List(Dict(), required=False)


class Test(MarshmallowSchema):
    type = String(required=True)


class Steps(MarshmallowSchema):
    num = Integer(data_key='step', required=True)
    start_time = DateTime(data_key='date_init', required=False)
    end_time = DateTime(data_key='date_end', required=False)
    severity = EnumField(Severity)

    @pre_load
    def preload_datas(self, data: dict):
        # import pdb; pdb.set_trace()
        data['severity'] = Severity.Info.name
        data.pop('duration', None)
        data.pop('commands', None)

        if not data.pop('success', False):
            data['severity'] = Severity.Error.name

        if data.get('date_init'):
            data['date_init'] = datetime.fromtimestamp(data['date_init']).isoformat()

        if data.get('date_end'):
            data['date_end'] = datetime.fromtimestamp(data['date_end']).isoformat()
        else:
            data['date_end'] = data['date_init']


class Sanitize(MarshmallowSchema):
    steps = Nested(Steps, many=True, required=True, data_key='erasure_steps')
    validation = Dict()
    device_info = Dict()
    method = Dict(required=True)
    sanitize_version = String()
    severity = EnumField(Severity, required=True)

    @pre_load
    def preload_datas(self, data: dict):
        data['severity'] = Severity.Info.name

        if not data.pop('result', False):
            data['severity'] = Severity.Error.name


class Snapshot_lite(Thing):
    uuid = String(required=True)
    version = String(required=True)
    schema_api = String(required=True)
    software = EnumField(
        SnapshotSoftware,
        required=True,
        description='The software that generated this Snapshot.',
    )
    sid = String(required=True)
    timestamp = String(required=True)
    settings_version = String(required=False)
    hwmd = Nested(Snapshot_lite_data, required=True)
    tests = Nested(Test, many=True, collection_class=list, required=False)
    sanitize = Nested(Sanitize, many=True, collection_class=list, required=False)

    @validates_schema
    def validate_workbench_version(self, data: dict):
        if data['schema_api'] not in app.config['SCHEMA_HWMD']:
            raise ValidationError(
                'Min. supported Workbench version is '
                '{} but yours is {}.'.format(
                    app.config['SCHEMA_HWMD'][0], data['version']
                ),
                field_names=['version'],
            )