import uuid
import stix2
import os
import shutil

from uuid import UUID

from stix2 import Vulnerability
from stix2extensions._extensions import vulnerability_opencti_ExtensionDefinitionSMO

# create the directories

tmp_directories = [
    "tmp_object_store",
]

for directory in tmp_directories:
    if not os.path.exists(directory):
        os.makedirs(directory)

# define UUID for generating UUIDv5s -- use stix2extensions namespace as this is an SDO

namespace=UUID("1abb62b9-e513-5f55-8e73-8f6d7b55c237")

# define values that are recycled between objects

### dogesec-demo

created_by_ref="identity--" + str(uuid.uuid5(namespace, f"dogesec-demo"))
created="2020-01-01T00:00:00.000Z"
modified="2020-01-01T00:00:00.000Z"

# Create Vulnerability SDO object

### vulnerability--132b38d2-c0d2-5c11-967f-9959a7099a2f

example_opencti_VulnerabilitySDO = Vulnerability(
                        id="vulnerability--"+ str(uuid.uuid5(namespace, f"Another demo Vulnerability")),
                        created_by_ref=created_by_ref,
                        created=created,
                        modified=modified,
                        name="CVE-XXX-XXXX",
                        external_references=[
                            {
                                "source_name": "cve",
                                "url": "https://nvd.nist.gov/vuln/detail/CVE-XXX-XXXX",
                                "external_id": "CVE-XXX-XXXX"
                            }
                        ],
                        object_marking_refs=[
                            "marking-definition--94868c89-83c2-464b-929b-a1a8aa3c8487", # this is TLP:CLEAR
                            "marking-definition--" + str(uuid.uuid5(namespace, f"stix2extensions")) # marking-definition--97ba4e8b-04f6-57e8-8f6e-3a0f0a7dc0fb
                        ],
                        x_opencti_cvss_v2_base_score=4.8,
                        x_opencti_cvss_v2_vector_string="AV:N/AC:L/Au:N/C:P/I:P/A:P",
                        x_opencti_cvss_base_score=4.8,
                        x_opencti_cvss_vector_string="CVSS:3.1/AV:N/AC:L/PR:H/UI:R/S:C/C:L/I:L/A:N",
                        extensions={
                            vulnerability_opencti_ExtensionDefinitionSMO.id: {
                                    "extension_type": "toplevel-property-extension"
                            }
                        }
                    )

# Write the objects to the filestore
## https://stix2.readthedocs.io/en/latest/guide/filesystem.html#FileSystemSource

### Creating FileSystemStore and adding MarkingDefinitionSMO for each directory

fs_directories = {
    "tmp_object_store": example_opencti_VulnerabilitySDO
}

for directory, vulnerability_sdo in fs_directories.items():
    fs_store = stix2.FileSystemStore(directory)
    fs_store.add([vulnerability_sdo])

# Now move those files into the standardised locations for easy download

final_directories = [
    "example_objects/sdos"
]

for directory in final_directories:
    if not os.path.exists(directory):
        os.makedirs(directory)

shutil.move("tmp_object_store/vulnerability/vulnerability--" + str(uuid.uuid5(namespace, f"Another demo Vulnerability")) + "/20200101000000000.json", "example_objects/properties/vulnerability--" + str(uuid.uuid5(namespace, f"Another demo Vulnerability")) + ".json")

shutil.rmtree("tmp_object_store")