models.py

    ),
)


# Table required for Many-to-Many relationships between Ansible groups and hosts
ansible_groups_hosts_table = db.Table(
    "ansible_groups_hosts",
    db.Column(
        "ansible_group_id",
        db.Integer,
        db.ForeignKey("ansible_group.id"),
        primary_key=True,
    ),
    db.Column("host_id", db.Integer, db.ForeignKey("host.id"), primary_key=True),
)


class AnsibleGroupType(Enum):
    STATIC = "STATIC"
    NETWORK_SCOPE = "NETWORK_SCOPE"
    NETWORK = "NETWORK"
    DEVICE_TYPE = "DEVICE_TYPE"
    IOC = "IOC"

    def __str__(self):
        return self.name

    @classmethod
    def choices(cls):
        return [(item, item.name) for item in AnsibleGroupType]

    @classmethod
    def coerce(cls, value):
        return value if type(value) == AnsibleGroupType else AnsibleGroupType[value]


class AnsibleGroup(CreatedMixin, db.Model):
    __versioned__ = {}
    __tablename__ = "ansible_group"
    # Define id here so that it can be used in the primary and secondary join
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(CIText, nullable=False, unique=True)
    vars = db.Column(postgresql.JSONB)
    type = db.Column(
        db.Enum(AnsibleGroupType, name="ansible_group_type"),
        default=AnsibleGroupType.STATIC,
        nullable=False,
    )

    children = db.relationship(
        "AnsibleGroup",
        secondary=ansible_groups_parent_child_table,
        primaryjoin=id == ansible_groups_parent_child_table.c.parent_group_id,
        secondaryjoin=id == ansible_groups_parent_child_table.c.child_group_id,
        backref=db.backref("parents"),
    )

    def __str__(self):
        return str(self.name)

    @property
    def is_dynamic(self):
        return self.type != AnsibleGroupType.STATIC

    @property
    def hosts(self):
        if self.type == AnsibleGroupType.STATIC:
            return self._hosts
        if self.type == AnsibleGroupType.NETWORK_SCOPE:
            return (
                Host.query.join(Host.interfaces)
                .join(Interface.network)
                .join(Network.scope)
                .filter(NetworkScope.name == self.name, Interface.name == Host.name)
                .order_by(Host.name)
                .all()
            )
        if self.type == AnsibleGroupType.NETWORK:
            return (
                Host.query.join(Host.interfaces)
                .join(Interface.network)
                .filter(Network.vlan_name == self.name, Interface.name == Host.name)
                .order_by(Host.name)
                .all()
            )
        if self.type == AnsibleGroupType.DEVICE_TYPE:
            return (
                Host.query.join(Host.device_type)
                .filter(DeviceType.name == self.name)
                .order_by(Host.name)
                .all()
            )
        if self.type == AnsibleGroupType.IOC:
            return Host.query.filter(Host.is_ioc.is_(True)).order_by(Host.name).all()

    @hosts.setter
    def hosts(self, value):
        # For dynamic group type, _hosts can only be set to []
        if self.is_dynamic and value:
            raise AttributeError("can't set dynamic hosts")
        self._hosts = value

    def to_dict(self, recursive=False):
        d = super().to_dict()
        d.update(
            {
                "name": self.name,
                "vars": self.vars,
                "type": self.type.name,
                "hosts": [host.fqdn for host in self.hosts],
                "children": [str(child) for child in self.children],
            }
        )
        return d


class Host(CreatedMixin, SearchableMixin, db.Model):
    __versioned__ = {}
    __mapping__ = {
        "created_at": {"type": "date", "format": "yyyy-MM-dd HH:mm"},
        "updated_at": {"type": "date", "format": "yyyy-MM-dd HH:mm"},
        "user": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
        "name": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
        "fqdn": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
        "is_ioc": {"type": "boolean"},
        "device_type": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
        "model": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
        "description": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
        "items": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
        "interfaces": {
            "properties": {
                "id": {"enabled": False},
                "created_at": {"type": "date", "format": "yyyy-MM-dd HH:mm"},
                "updated_at": {"type": "date", "format": "yyyy-MM-dd HH:mm"},
                "user": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
                "is_main": {"type": "boolean"},
                "network": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
                "ip": {"type": "ip"},
                "netmask": {"enabled": False},
                "name": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
                "mac": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
                "host": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
                "cnames": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
                "domain": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
                "device_type": {
                    "type": "text",
                    "fields": {"keyword": {"type": "keyword"}},
                },
                "model": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
            }
        },
        "ansible_vars": {"enabled": False},
        "ansible_groups": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
    }

    # id shall be defined here to be used by SQLAlchemy-Continuum
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.Text, nullable=False, unique=True)
    description = db.Column(db.Text)
    device_type_id = db.Column(
        db.Integer, db.ForeignKey("device_type.id"), nullable=False
    )
    is_ioc = db.Column(db.Boolean, nullable=False, default=False)
    ansible_vars = db.Column(postgresql.JSONB)

    # 1. Set cascade to all (to add delete) and delete-orphan to delete all interfaces
    # when deleting a host
    # 2. Return interfaces sorted by name so that the main one (the one starting with
    # the same name as the host) is always the first one.
    # As an interface name always has to start with the name of the host, the one
    # matching the host name will always come first.
    interfaces = db.relationship(
        "Interface",
        backref=db.backref("host", lazy="joined"),
        cascade="all, delete-orphan",
        lazy="joined",
        order_by="Interface.name",
    )
    items = db.relationship(
        "Item", backref=db.backref("host", lazy="joined"), lazy="joined"
    )
    ansible_groups = db.relationship(
        "AnsibleGroup",
        secondary=ansible_groups_hosts_table,
        lazy="joined",
        backref=db.backref("_hosts"),
    )

    def __init__(self, **kwargs):
        # Automatically convert device_type as an instance of its class if passed as a string
        if "device_type" in kwargs:
            kwargs["device_type"] = utils.convert_to_model(
                kwargs["device_type"], DeviceType
            )
        # Automatically convert items to a list of instances if passed as a list of ics_id
        if "items" in kwargs:
            kwargs["items"] = [
                utils.convert_to_model(item, Item, filter="ics_id")
                for item in kwargs["items"]
            ]
        # Automatically convert ansible groups to a list of instances if passed as a list of strings
        if "ansible_groups" in kwargs:
            kwargs["ansible_groups"] = [
                utils.convert_to_model(group, AnsibleGroup)
                for group in kwargs["ansible_groups"]
            ]
        super().__init__(**kwargs)

    @property
    def model(self):
        """Return the model of the first linked item"""
        try:
            return utils.format_field(self.items[0].model)
        except IndexError:
            return None

    @property
    def main_interface(self):
        """Return the host main interface

        The main interface is the one that has the same name as the host
        or the first one found
        """
        # As interfaces are sorted, the first one is always the main one
        try:
            return self.interfaces[0]
        except IndexError:
            return None

    @property
    def main_network(self):
        """Return the host main interface network"""
        try:
            return self.main_interface.network
        except AttributeError:
            return None

    @property
    def fqdn(self):
        """Return the host fully qualified domain name

        The domain is based on the main interface
        """
        if self.main_interface:
            return f"{self.name}.{self.main_interface.network.domain}"
        else:
            return self.name

    def __str__(self):
        return str(self.name)

    @validates("name")
    def validate_name(self, key, string):
        """Ensure the name matches the required format"""
        if string is None:
            return None
        # Force the string to lowercase
        lower_string = string.lower()
        if HOST_NAME_RE.fullmatch(lower_string) is None:
            raise ValidationError("Host name shall match [a-z0-9\-]{2,20}")
        existing_cname = Cname.query.filter_by(name=lower_string).first()
        if existing_cname:
            raise ValidationError(f"Host name matches an existing cname")
        existing_interface = Interface.query.filter(
            Interface.name == lower_string, Interface.host_id != self.id
        ).first()
        if existing_interface:
            raise ValidationError(f"Host name matches an existing interface")
        return lower_string

    def stack_members(self):
        """Return all items part of the stack sorted by stack member number"""
        members = [item for item in self.items if item.stack_member is not None]
        return sorted(members, key=lambda x: x.stack_member)

    def stack_members_numbers(self):
        """Return the list of stack member numbers"""
        return [item.stack_member for item in self.stack_members()]

    def free_stack_members(self):
        """Return the list of free stack member numbers"""
        return [nb for nb in range(0, 10) if nb not in self.stack_members_numbers()]

    def to_dict(self, recursive=False):
        d = super().to_dict()
        d.update(
            {
                "name": self.name,
                "fqdn": self.fqdn,
                "is_ioc": self.is_ioc,
                "device_type": str(self.device_type),
                "model": self.model,
                "description": self.description,
                "items": [str(item) for item in self.items],
                "interfaces": [str(interface) for interface in self.interfaces],
                "ansible_vars": self.ansible_vars,
                "ansible_groups": [str(group) for group in self.ansible_groups],
            }
        )
        if recursive:
            # Replace the list of interface names by the full representation
            # so that we can index everything in elasticsearch
            d["interfaces"] = [interface.to_dict() for interface in self.interfaces]
        return d


class Interface(CreatedMixin, db.Model):
    network_id = db.Column(db.Integer, db.ForeignKey("network.id"), nullable=False)
    ip = db.Column(postgresql.INET, nullable=False, unique=True)
    name = db.Column(db.Text, nullable=False, unique=True)
    mac = db.Column(postgresql.MACADDR, nullable=True, unique=True)
    host_id = db.Column(db.Integer, db.ForeignKey("host.id"), nullable=False)

    # Add delete and delete-orphan options to automatically delete cnames when:
    # - deleting an interface
    # - de-associating a cname (removing it from the interface.cnames list)
    cnames = db.relationship(
        "Cname",
        backref=db.backref("interface", lazy="joined"),
        cascade="all, delete, delete-orphan",
        lazy="joined",
    )

    def __init__(self, **kwargs):
        # Always set self.host and not self.host_id to call validate_name
        host_id = kwargs.pop("host_id", None)
        if host_id is not None:
            host = Host.query.get(host_id)
        elif "host" in kwargs:
            # Automatically convert host to an instance of Host if it was passed
            # as a string
            host = utils.convert_to_model(kwargs.pop("host"), Host, "name")
        else:
            host = None
        # Always set self.network and not self.network_id to call validate_interfaces
        network_id = kwargs.pop("network_id", None)
        if network_id is not None:
            kwargs["network"] = Network.query.get(network_id)
        elif "network" in kwargs:
            # Automatically convert network to an instance of Network if it was passed
            # as a string
            kwargs["network"] = utils.convert_to_model(
                kwargs["network"], Network, "vlan_name"
            )
        # WARNING! Setting self.network will call validate_interfaces in the Network class
        # For the validation to work, self.ip must be set before!
        # Ensure that ip is passed before network
        try:
            ip = kwargs.pop("ip")
        except KeyError:
            super().__init__(host=host, **kwargs)
        else:
            super().__init__(host=host, ip=ip, **kwargs)

    @validates("name")
    def validate_name(self, key, string):
        """Ensure the name matches the required format"""
        if string is None:
            return None
        # Force the string to lowercase
        lower_string = string.lower()
        if HOST_NAME_RE.fullmatch(lower_string) is None:
            raise ValidationError("Interface name shall match [a-z0-9\-]{2,20}")
        if self.host and not lower_string.startswith(self.host.name):
            raise ValidationError(
                f"Interface name shall start with the host name '{self.host}'"
            )
        existing_cname = Cname.query.filter_by(name=lower_string).first()
        if existing_cname:
            raise ValidationError(f"Interface name matches an existing cname")
        existing_host = Host.query.filter(
            Host.name == lower_string, Host.id != self.host.id
        ).first()
        if existing_host:
            raise ValidationError(f"Interface name matches an existing host")
        return lower_string

    @validates("mac")
    def validate_mac(self, key, string):
        """Ensure the mac is a valid MAC address"""
        if not string:
            return None
        if MAC_ADDRESS_RE.fullmatch(string) is None:
            raise ValidationError(f"'{string}' does not appear to be a MAC address")
        return string

    @validates("cnames")
    def validate_cnames(self, key, cname):
        """Ensure the cname is unique by domain"""
        existing_cnames = Cname.query.filter_by(name=cname.name).all()
        for existing_cname in existing_cnames:
            if existing_cname.domain == str(self.network.domain):
                raise ValidationError(
                    f"Duplicate cname on the {self.network.domain} domain"
                )
        return cname

    @property
    def address(self):
        return ipaddress.ip_address(self.ip)

    @property
    def is_ioc(self):
        return self.is_main and self.host.is_ioc

    @property
    def is_main(self):
        return self.name == self.host.main_interface.name

    def __str__(self):
        return str(self.name)

    def __repr__(self):
        return f"Interface(id={self.id}, network_id={self.network_id}, ip={self.ip}, name={self.name}, mac={self.mac})"

    def to_dict(self, recursive=False):
        d = super().to_dict()
        d.update(
            {
                "is_main": self.is_main,
                "network": str(self.network),
                "ip": self.ip,
                "netmask": str(self.network.netmask),
                "name": self.name,
                "mac": utils.format_field(self.mac),
                "host": utils.format_field(self.host),
                "cnames": [str(cname) for cname in self.cnames],
                "domain": str(self.network.domain),
            }
        )
        if self.host:
            d["device_type"] = str(self.host.device_type)
            d["model"] = utils.format_field(self.host.model)
        else:
            d["device_type"] = None
            d["model"] = None
        return d


class Mac(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    address = db.Column(postgresql.MACADDR, nullable=False, unique=True)
    item_id = db.Column(db.Integer, db.ForeignKey("item.id"))

    def __str__(self):
        return str(self.address)

    @validates("address")
    def validate_address(self, key, string):
        """Ensure the address is a valid MAC address"""
        if string is None:
            return None
        if MAC_ADDRESS_RE.fullmatch(string) is None:
            raise ValidationError(f"'{string}' does not appear to be a MAC address")
        return string

    def to_dict(self, recursive=False):
        return {
            "id": self.id,
            "address": self.address,
            "item": utils.format_field(self.item),
        }


class Cname(CreatedMixin, db.Model):
    name = db.Column(db.Text, nullable=False)
    interface_id = db.Column(db.Integer, db.ForeignKey("interface.id"), nullable=False)

    def __init__(self, **kwargs):
        # Always set self.interface and not self.interface_id to call validate_cnames
        interface_id = kwargs.pop("interface_id", None)
        if interface_id is not None:
            kwargs["interface"] = Interface.query.get(interface_id)
        super().__init__(**kwargs)

    def __str__(self):
        return str(self.name)

    @property
    def domain(self):
        """Return the cname domain name"""
        return str(self.interface.network.domain)

    @property
    def fqdn(self):
        """Return the cname fully qualified domain name"""
        return f"{self.name}.{self.domain}"

    @validates("name")
    def validate_name(self, key, string):
        """Ensure the name matches the required format"""
        if string is None:
            return None
        # Force the string to lowercase
        lower_string = string.lower()
        if HOST_NAME_RE.fullmatch(lower_string) is None:
            raise ValidationError("cname shall match [a-z0-9\-]{2,20}")
        existing_interface = Interface.query.filter_by(name=lower_string).first()
        if existing_interface:
            raise ValidationError(f"cname matches an existing interface")
        existing_host = Host.query.filter_by(name=lower_string).first()
        if existing_host:
            raise ValidationError(f"cname matches an existing host")
        return lower_string

    def to_dict(self, recursive=False):
        d = super().to_dict()
        d.update({"name": self.name, "interface": str(self.interface)})
        return d


class Domain(CreatedMixin, db.Model):
    name = db.Column(db.Text, nullable=False, unique=True)

    scopes = db.relationship(
        "NetworkScope", backref=db.backref("domain", lazy="joined"), lazy=True
    )
    networks = db.relationship(
        "Network", backref=db.backref("domain", lazy="joined"), lazy=True
    )

    def __str__(self):
        return str(self.name)

    def to_dict(self, recursive=False):
        d = super().to_dict()
        d.update(
            {
                "name": self.name,
                "scopes": [str(scope) for scope in self.scopes],
                "networks": [str(network) for network in self.networks],
            }
        )
        return d


class NetworkScope(CreatedMixin, db.Model):
    __tablename__ = "network_scope"
    name = db.Column(CIText, nullable=False, unique=True)
    first_vlan = db.Column(db.Integer, nullable=False, unique=True)
    last_vlan = db.Column(db.Integer, nullable=False, unique=True)
    supernet = db.Column(postgresql.CIDR, nullable=False, unique=True)
    domain_id = db.Column(db.Integer, db.ForeignKey("domain.id"), nullable=False)
    description = db.Column(db.Text)

    networks = db.relationship(
        "Network", backref=db.backref("scope", lazy="joined"), lazy=True
    )

    __table_args__ = (
        sa.CheckConstraint(
            "first_vlan < last_vlan", name="first_vlan_less_than_last_vlan"
        ),
    )

    def __str__(self):
        return str(self.name)

    @property
    def supernet_ip(self):
        return ipaddress.ip_network(self.supernet)

    def prefix_range(self):
        """Return the list of subnet prefix that can be used for this network scope"""
        return list(range(self.supernet_ip.prefixlen + 1, 31))

    def vlan_range(self):
        """Return the list of vlan ids that can be assigned for this network scope

        The range is defined by the first and last vlan
        """
        return range(self.first_vlan, self.last_vlan + 1)

    def used_vlans(self):
        """Return the list of vlan ids in use

        The list is sorted
        """
        return sorted(network.vlan_id for network in self.networks)

    def available_vlans(self):
        """Return the list of vlan ids available"""
        return [vlan for vlan in self.vlan_range() if vlan not in self.used_vlans()]

    def used_subnets(self):
        """Return the list of subnets in use

        The list is sorted
        """
        return sorted(network.network_ip for network in self.networks)

    def available_subnets(self, prefix):
        """Return the list of available subnets with the given prefix"""
        return [
            str(subnet)
            for subnet in self.supernet_ip.subnets(new_prefix=prefix)
            if subnet not in self.used_subnets()
        ]

    def to_dict(self, recursive=False):
        d = super().to_dict()
        d.update(
            {
                "name": self.name,
                "first_vlan": self.first_vlan,
                "last_vlan": self.last_vlan,
                "supernet": self.supernet,
                "description": self.description,
                "domain": str(self.domain),
                "networks": [str(network) for network in self.networks],
            }
        )
        return d


# Define RQ JobStatus as a Python enum
# We can't use the one defined in rq/job.py as it's
# not a real enum (it's a custom one) and is not
# compatible with sqlalchemy
class JobStatus(Enum):
    QUEUED = "queued"
    FINISHED = "finished"
    FAILED = "failed"
    STARTED = "started"
    DEFERRED = "deferred"


class Task(db.Model):
    # Use job id generated by RQ
    id = db.Column(postgresql.UUID, primary_key=True)
    created_at = db.Column(db.DateTime, default=utcnow())
    ended_at = db.Column(db.DateTime)
    name = db.Column(db.Text, nullable=False, index=True)
    command = db.Column(db.Text)
    status = db.Column(db.Enum(JobStatus, name="job_status"))
    awx_resource = db.Column(db.Text)
    awx_job_id = db.Column(db.Integer)
    exception = db.Column(db.Text)
    user_id = db.Column(
        db.Integer,
        db.ForeignKey("user_account.id"),
        nullable=False,
        default=utils.fetch_current_user_id,
    )
    depends_on_id = db.Column(postgresql.UUID, db.ForeignKey("task.id"))

    reverse_dependencies = db.relationship(
        "Task", backref=db.backref("depends_on", remote_side=[id])
    )

    @property
    def awx_job_url(self):
        if self.awx_job_id is None:
            return None
        if self.awx_resource == "job":
            route = "jobs/playbook"
        elif self.awx_resource == "workflow_job":
            route = "workflows"
        elif self.awx_resource == "inventory_source":
            route = "jobs/inventory"
        else:
            current_app.logger.warning(f"Unknown AWX resource: {self.awx_resource}")
            return None
        return urllib.parse.urljoin(
            current_app.config["AWX_URL"], f"/#/{route}/{self.awx_job_id}"
        )

    def __str__(self):
        return str(self.id)

    def to_dict(self, recursive=False):
        return {
            "id": self.id,
            "name": self.name,
            "created_at": utils.format_field(self.created_at),
            "ended_at": utils.format_field(self.ended_at),
            "status": self.status.name,
            "awx_resource": self.awx_resource,
            "awx_job_id": self.awx_job_id,
            "awx_job_url": self.awx_job_url,
            "depends_on": self.depends_on_id,
            "command": self.command,
            "exception": self.exception,
            "user": str(self.user),
        }


def trigger_core_services_update(session):
    """Trigger core services update on any Interface modification.

    Called by before flush hook
    """
    # In session.dirty, we need to check session.is_modified(instance) because when updating a Host,
    # the interface is added to the session even if not modified.
    # In session.deleted, session.is_modified(instance) is usually False (we shouldn't check it).
    # In session.new, it will always be True and we don't need to check it.
    for kind in ("new", "dirty", "deleted"):
        for instance in getattr(session, kind):
            if isinstance(instance, Interface) and (
                (kind == "dirty" and session.is_modified(instance))
                or (kind in ("new", "deleted"))
            ):
                utils.trigger_core_services_update()
                return True
    return False


def trigger_inventory_update(session):
    """Trigger an inventory update in AWX

    Update on any AnsibleGroup/Cname/Domain/Host/Network/NetworkScope
    modification, but not on Interface as it's triggered with core
    services update.

    Called by before flush hook
    """
    # In session.dirty, we need to check session.is_modified(instance) because the instance
    # could have been added to the session without being modified
    # In session.deleted, session.is_modified(instance) is usually False (we shouldn't check it).
    # In session.new, it will always be True and we don't need to check it.
    for kind in ("new", "dirty", "deleted"):
        for instance in getattr(session, kind):
            if isinstance(
                instance, (AnsibleGroup, Cname, Domain, Host, Network, NetworkScope)
            ) and (
                (kind == "dirty" and session.is_modified(instance))
                or (kind in ("new", "deleted"))
            ):
                utils.trigger_inventory_update()
                return True
    return False


@sa.event.listens_for(db.session, "before_flush")
def before_flush(session, flush_context, instances):
    """Before flush hook

    Used to trigger core services and inventory update.

    See http://docs.sqlalchemy.org/en/latest/orm/session_events.html#before-flush
    """
    if trigger_core_services_update(session):
        # This will also trigger an inventory update
        return
    trigger_inventory_update(session)


@sa.event.listens_for(Host.is_ioc, "set")
def receive_host_is_ioc_set(target, value, oldvalue, initiator):
    """Listen for the 'set' event on Host.is_ioc

    Trigger repository creation in GitLab
    """
    if value is True and oldvalue is not True:
        utils.trigger_ioc_repository_creation(target)


# call configure_mappers after defining all the models
# required by sqlalchemy_continuum
sa.orm.configure_mappers()
ItemVersion = version_class(Item)
# Set SQLAlchemy event listeners
db.event.listen(db.session, "before_flush", SearchableMixin.before_flush)
db.event.listen(
    db.session, "after_flush_postexec", SearchableMixin.after_flush_postexec
)
db.event.listen(db.session, "after_commit", SearchableMixin.after_commit)