from datetime import date from factory import Sequence, post_generation, SubFactory from factory.alchemy import SQLAlchemyModelFactory from factory.fuzzy import FuzzyChoice, FuzzyText, FuzzyDate, FuzzyInteger from lemur.database import db from lemur.authorities.models import Authority from lemur.certificates.models import Certificate from lemur.destinations.models import Destination from lemur.sources.models import Source from lemur.notifications.models import Notification from lemur.users.models import User from lemur.roles.models import Role from lemur.endpoints.models import Policy, Endpoint from lemur.policies.models import RotationPolicy from .vectors import INTERNAL_VALID_SAN_STR, PRIVATE_KEY_STR class BaseFactory(SQLAlchemyModelFactory): """Base factory.""" class Meta: """Factory configuration.""" abstract = True sqlalchemy_session = db.session class CertificateFactory(BaseFactory): """Certificate factory.""" name = Sequence(lambda n: 'certificate{0}'.format(n)) chain = INTERNAL_VALID_SAN_STR body = INTERNAL_VALID_SAN_STR private_key = PRIVATE_KEY_STR owner = 'joe@example.com' status = FuzzyChoice(['valid', 'revoked', 'unknown']) deleted = False description = FuzzyText(length=128) active = True date_created = FuzzyDate(date(2016, 1, 1), date(2020, 1, 1)) class Meta: """Factory Configuration.""" model = Certificate @post_generation def user(self, create, extracted, **kwargs): if not create: return if extracted: self.user_id = extracted.id @post_generation def authority(self, create, extracted, **kwargs): if not create: return if extracted: self.authority_id = extracted.id @post_generation def notifications(self, create, extracted, **kwargs): if not create: return if extracted: for notification in extracted: self.notifications.append(notification) @post_generation def destinations(self, create, extracted, **kwargs): if not create: return if extracted: for destination in extracted: self.destintations.append(destination) @post_generation def replaces(self, create, extracted, **kwargs): if not create: return if extracted: for replace in extracted: self.replaces.append(replace) @post_generation def sources(self, create, extracted, **kwargs): if not create: return if extracted: for source in extracted: self.sources.append(source) @post_generation def domains(self, create, extracted, **kwargs): if not create: return if extracted: for domain in extracted: self.domains.append(domain) @post_generation def roles(self, create, extracted, **kwargs): if not create: return if extracted: for domain in extracted: self.roles.append(domain) class AuthorityFactory(BaseFactory): """Authority factory.""" name = Sequence(lambda n: 'authority{0}'.format(n)) owner = 'joe@example.com' plugin = {'slug': 'test-issuer'} authority_certificate = SubFactory(CertificateFactory) class Meta: """Factory configuration.""" model = Authority @post_generation def roles(self, create, extracted, **kwargs): if not create: return if extracted: for role in extracted: self.roles.append(role) class RotationPolicyFactory(BaseFactory): """Rotation Factory.""" name = Sequence(lambda n: 'policy{0}'.format(n)) days = 30 class Meta: """Factory configuration.""" model = RotationPolicy class DestinationFactory(BaseFactory): """Destination factory.""" plugin_name = 'test-destination' label = Sequence(lambda n: 'destination{0}'.format(n)) class Meta: """Factory Configuration.""" model = Destination class SourceFactory(BaseFactory): """Source factory.""" plugin_name = 'test-source' label = Sequence(lambda n: 'source{0}'.format(n)) class Meta: """Factory Configuration.""" model = Source class NotificationFactory(BaseFactory): """Notification factory.""" plugin_name = 'test-notification' label = Sequence(lambda n: 'notification{0}'.format(n)) class Meta: """Factory Configuration.""" model = Notification class RoleFactory(BaseFactory): """Role factory.""" name = Sequence(lambda n: 'role{0}'.format(n)) class Meta: """Factory Configuration.""" model = Role @post_generation def users(self, create, extracted, **kwargs): if not create: return if extracted: for user in extracted: self.users.append(user) class UserFactory(BaseFactory): """User Factory.""" username = Sequence(lambda n: 'user{0}'.format(n)) email = Sequence(lambda n: 'user{0}@example.com'.format(n)) active = True password = FuzzyText(length=24) class Meta: """Factory Configuration.""" model = User @post_generation def roles(self, create, extracted, **kwargs): if not create: return if extracted: for role in extracted: self.roles.append(role) @post_generation def certificates(self, create, extracted, **kwargs): if not create: return if extracted: for cert in extracted: self.certificates.append(cert) @post_generation def authorities(self, create, extracted, **kwargs): if not create: return if extracted: for authority in extracted: self.authorities.append(authority) class PolicyFactory(BaseFactory): """Policy Factory.""" name = Sequence(lambda n: 'endpoint{0}'.format(n)) class Meta: """Factory Configuration.""" model = Policy class EndpointFactory(BaseFactory): """Endpoint Factory.""" owner = 'joe@example.com' name = Sequence(lambda n: 'endpoint{0}'.format(n)) type = FuzzyChoice(['elb']) active = True port = FuzzyInteger(0, high=65535) dnsname = 'endpoint.example.com' policy = SubFactory(PolicyFactory) certificate = SubFactory(CertificateFactory) source = SubFactory(SourceFactory) class Meta: """Factory Configuration.""" model = Endpoint