ceios/builder/tests/test_model.py

import json
import os
from pathlib import Path
from tempfile import NamedTemporaryFile

import pytest
import stix2

import attack_flow.model


def test_load_attack_flow_bundle():
    bundle_json = {
        "type": "bundle",
        "id": "bundle--3b210ed6-4aac-4620-9e75-79a9b7ae99c5",
        "objects": [
            {
                "type": "attack-flow",
                "spec_version": "2.1",
                "id": "attack-flow--77694729-3848-4261-a294-889837d58460",
                "name": "Test Flow",
                "extensions": {
                    "extension-definition--fb9c968a-745b-4ade-9b25-c324172197f4": {
                        "extension_type": "new-sdo"
                    }
                },
            }
        ],
    }
    with NamedTemporaryFile("w") as f:
        json.dump(bundle_json, f)
        f.seek(0, os.SEEK_SET)
        flow_bundle = attack_flow.model.load_attack_flow_bundle(Path(f.name))
        assert flow_bundle.id == "bundle--3b210ed6-4aac-4620-9e75-79a9b7ae99c5"


def test_confidence_label_to_num():
    assert attack_flow.model.confidence_label_to_num("Speculation") == 0
    assert attack_flow.model.confidence_label_to_num("Even Odds") == 50


def test_confidence_label_to_num_with_invalid_label():
    with pytest.raises(ValueError):
        attack_flow.model.confidence_label_to_num("Foobar")


def test_confidence_label_num_to_label():
    assert attack_flow.model.confidence_num_to_label(0) == "Speculation"
    assert attack_flow.model.confidence_num_to_label(45) == "Even Odds"


def test_confidence_label_num_to_label_with_invalid_num():
    with pytest.raises(ValueError):
        assert attack_flow.model.confidence_num_to_label(-10)

    with pytest.raises(ValueError):
        assert attack_flow.model.confidence_num_to_label(110)


def test_get_flow_object():
    bundle = stix2.Bundle(
        attack_flow.model.AttackFlow(
            id="attack-flow--0c545a6f-3da2-4fa8-9789-68fd98257d10",
            name="Test Flow",
        ),
        id="bundle--b704f0ad-5df9-4386-b0f4-3317859fd4e0",
    )
    flow = attack_flow.model.get_flow_object(bundle)
    assert flow.id == "attack-flow--0c545a6f-3da2-4fa8-9789-68fd98257d10"