diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py new file mode 100644 index 0000000..0a39d8d --- /dev/null +++ b/tests/test_data_generator.py @@ -0,0 +1,158 @@ +from cyhy_db.models import CVE, RequestDoc +from cyhy_db.models.request_doc import Agency, Contact, Location, Window +from cyhy_db.models.enum import ( + AgencyType, + DayOfWeek, + CVSSVersion, + PocType, + ReportPeriod, + ReportType, + ScanType, + Scheduler, + Stage, +) +from cyhy_db.utils import utcnow +from mimesis_factory import MimesisField +from mimesis.locales import DEFAULT_LOCALE +from mimesis.providers.base import BaseProvider +from pytest_factoryboy import register +import factory +from mimesis import Generic +import ipaddress +import random +from datetime import datetime + + +class CyHyProvider(BaseProvider): + class Meta: + name = "cyhy_provider" + + def cve_id(self, year=None): + # If year is None, generate a random year between 1999 and the current year + if year is None: + year = self.random.randint(1999, datetime.now().year) + + # Generate a random number for the CVE, ensuring it has a leading zero if necessary + number = self.random.randint(1, 99999) + + return f"CVE-{year}-{number:05d}" + + def network_ipv4(self): + # Generate a base IP address + base_ip = generic.internet.ip_v4() + # Choose a random CIDR between 24-30 to ensure a smaller network size and avoid host bits set error + cidr = random.randint(24, 30) + # Create the network address + network = ipaddress.IPv4Network(f"{base_ip}/{cidr}", strict=False) + return network + + +generic = Generic(locale=DEFAULT_LOCALE) +generic.add_provider(CyHyProvider) + + +@register +class CVEFactory(factory.Factory): + class Meta(object): + model = CVE + + id = factory.LazyFunction(lambda: generic.cyhy_provider.cve_id()) + cvss_score = factory.LazyFunction(lambda: round(random.uniform(0, 10), 1)) + cvss_version = factory.LazyFunction(lambda: random.choice(list(CVSSVersion))) + severity = factory.LazyFunction(lambda: random.randint(1, 4)) + + +class AgencyFactory(factory.Factory): + class Meta(object): + model = Agency + + name = factory.Faker("company") + # Generate an acronym from the name + acronym = factory.LazyAttribute( + lambda o: "".join(word[0].upper() for word in o.name.split()) + ) + type = factory.LazyFunction(lambda: random.choice(list(AgencyType))) + contacts = factory.LazyFunction( + lambda: [ContactFactory() for _ in range(random.randint(1, 5))] + ) + location = factory.LazyFunction(lambda: LocationFactory()) + + +class ContactFactory(factory.Factory): + class Meta(object): + model = Contact + + email = factory.Faker("email") + name = factory.Faker("name") + phone = factory.Faker("phone_number") + type = factory.LazyFunction(lambda: random.choice(list(PocType))) + + +class LocationFactory(factory.Factory): + class Meta(object): + model = Location + + country_name = factory.Faker("country") + country = factory.Faker("country_code") + county_fips = factory.Faker("numerify", text="##") + county = factory.Faker("city") + gnis_id = factory.Faker("numerify", text="#######") + name = factory.Faker("city") + state_fips = factory.Faker("numerify", text="##") + state_name = factory.Faker("state") + state = factory.Faker("state_abbr") + + +class WindowFactory(factory.Factory): + class Meta(object): + model = Window + + day = factory.LazyFunction(lambda: random.choice(list(DayOfWeek))) + duration = factory.LazyFunction(lambda: random.randint(0, 168)) + start = factory.Faker("time", pattern="%H:%M:%S") + + +class RequestDocFactory(factory.Factory): + class Meta(object): + model = RequestDoc + + id = factory.LazyAttribute( + lambda o: o.agency.acronym + "-" + str(random.randint(1, 1000)) + ) + agency = factory.SubFactory(AgencyFactory) + enrolled = factory.LazyFunction(utcnow) + init_stage = factory.LazyFunction(lambda: random.choice(list(Stage))) + key = factory.Faker("password") + period_start = factory.LazyFunction(utcnow) + report_period = factory.LazyFunction(lambda: random.choice(list(ReportPeriod))) + retired = factory.LazyFunction(lambda: random.choice([True, False])) + scheduler = factory.LazyFunction(lambda: random.choice(list(Scheduler))) + stakeholder = factory.LazyFunction(lambda: random.choice([True, False])) + windows = factory.LazyFunction( + lambda: [WindowFactory() for _ in range(random.randint(1, 5))] + ) + networks = factory.LazyFunction( + lambda: [ + generic.cyhy_provider.network_ipv4() for _ in range(random.randint(1, 5)) + ] + ) + # create a set of 1 to 3 random scan types from the ScanType enum + scan_types = factory.LazyFunction( + lambda: set( + [random.choice(list(ScanType)) for _ in range(random.randint(1, 3))] + ) + ) + + +async def test_create_cves(): + for _ in range(100): + cve = CVEFactory() + print(cve) + await cve.save() + + +async def test_create_request_docs(): + for _ in range(100): + request_doc = RequestDocFactory() + print(request_doc) + await request_doc.save()