Merge pull request #715 from Ostorlab/feature/add-domain-name-model

Add domain asset model

src/ostorlab/runtimes/local/models/alembic/versions/b31a27c8584f_add_domain_asset_model.py

@@ -0,0 +1,49 @@
+"""Add domain asset model
+
+Revision ID: b31a27c8584f
+Revises: 3aa6ae380275
+Create Date: 2024-06-13 17:05:48.425291
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "b31a27c8584f"
+down_revision = "3aa6ae380275"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "domain_asset",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["id"], ["asset.id"], name=op.f("fk_domain_asset_id_asset")
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_domain_asset")),
+    )
+    op.create_table(
+        "domain_name",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("name", sa.String(length=255), nullable=True),
+        sa.Column("domain_asset_id", sa.Integer(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["domain_asset_id"],
+            ["domain_asset.id"],
+            name=op.f("fk_domain_name_domain_asset_id_domain_asset"),
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_domain_name")),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("domain_name")
+    op.drop_table("domain_asset")
+    # ### end Alembic commands ###

src/ostorlab/runtimes/local/models/models.py

@@ -1118,3 +1118,78 @@ def delete(network_id: int) -> None:
             session.query(Network).filter_by(id=network_id).delete()
             session.query(IPRange).filter_by(network_asset_id=network_id).delete()
             session.commit()
+
+
+class DomainName(Base):
+    __tablename__ = "domain_name"
+    id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True)
+    name = sqlalchemy.Column(sqlalchemy.String(255))
+    domain_asset_id = sqlalchemy.Column(
+        sqlalchemy.Integer, sqlalchemy.ForeignKey("domain_asset.id")
+    )
+
+    @staticmethod
+    def create(name: str, domain_asset_id: Optional[int] = None) -> "DomainName":
+        """Persist the domain name information in the database.
+
+        Args:
+            name: The domain name.
+            domain_asset_id: The domain asset id.
+
+        Returns:
+            DomainName object.
+        """
+        with Database() as session:
+            domain_name = DomainName(name=name, domain_asset_id=domain_asset_id)
+            session.add(domain_name)
+            session.commit()
+            return domain_name
+
+
+class DomainAsset(Asset):
+    __tablename__ = "domain_asset"
+    id = sqlalchemy.Column(
+        sqlalchemy.Integer, sqlalchemy.ForeignKey("asset.id"), primary_key=True
+    )
+
+    __mapper_args__ = {
+        "polymorphic_identity": "domain_asset",
+    }
+
+    @staticmethod
+    def create(
+        domains: List[Dict[str, str]], scan_id: Optional[int] = None
+    ) -> "DomainAsset":
+        """Persist the domain asset information in the database.
+
+        Args:
+            domains: list of domain names.
+            scan_id: The scan id.
+
+        Returns:
+            DomainAsset object.
+        """
+        with Database() as session:
+            domain_asset_instance = DomainAsset(scan_id=scan_id)
+            session.add(domain_asset_instance)
+            session.commit()
+
+            for domain in domains:
+                DomainName.create(
+                    name=domain.get("name"), domain_asset_id=domain_asset_instance.id
+                )
+            return domain_asset_instance
+
+    @staticmethod
+    def delete(domain_asset_id: int) -> None:
+        """Delete the domain asset information from the database.
+
+        Args:
+            domain_asset_id: The domain asset id.
+        """
+        with Database() as session:
+            session.query(DomainAsset).filter_by(id=domain_asset_id).delete()
+            session.query(DomainName).filter_by(
+                domain_asset_id=domain_asset_id
+            ).delete()
+            session.commit()

src/ostorlab/serve_app/oxo.py

@@ -31,6 +31,7 @@
 from ostorlab.assets import ipv4 as ipv4_address_asset
 from ostorlab.assets import ipv6 as ipv6_address_asset
 from ostorlab.assets import link as link_asset
+from ostorlab.assets import domain_name as domain_name_asset
 from ostorlab.assets import asset as ostorlab_asset
 
 DEFAULT_NUMBER_ELEMENTS = 15
@@ -372,6 +373,9 @@ def mutate(
             if asset.ip is not None:
                 new_asset = models.Network.create(networks=asset.ip)
                 created_assets.append(new_asset)
+            if asset.domain is not None:
+                new_asset = models.DomainAsset.create(domains=asset.domain)
+                created_assets.append(new_asset)
         if len(errors) > 0:
             error_messages = "\n".join(errors)
             raise graphql.GraphQLError(f"Invalid assets: {error_messages}")
@@ -396,6 +400,8 @@ def _validate(asset: types.OxoAssetInputType) -> Optional[str]:
             assets.append(asset.link)
         if asset.ip is not None:
             assets.append(asset.ip)
+        if asset.domain is not None:
+            assets.append(asset.domain)
 
         if len(assets) == 0:
             return f"Asset {asset} input is missing target."
@@ -654,6 +660,16 @@ def _prepare_assets(asset_ids: List[int]) -> List[ostorlab_asset.Asset]:
                         scan_assets.append(
                             link_asset.Link(url=link.url, method=link.method)
                         )
+                elif asset.type == "domain_asset":
+                    domains = (
+                        session.query(models.DomainName)
+                        .filter_by(domain_asset_id=asset.id)
+                        .all()
+                    )
+                    for domain in domains:
+                        scan_assets.append(
+                            domain_name_asset.DomainName(name=domain.name)
+                        )
                 else:
                     raise graphql.GraphQLError("Unsupported asset type.")
 

src/ostorlab/serve_app/types.py

@@ -328,6 +328,43 @@ def resolve_networks(self, info) -> List[OxoIPRangeAssetType]:
             return [OxoIPRangeAssetType(host=ip.host, mask=ip.mask) for ip in ips]
 
 
+class OxoDomainNameAssetType(graphene_sqlalchemy.SQLAlchemyObjectType):
+    class Meta:
+        model = models.DomainName
+        only_fields = "name"
+
+
+class OxoDomainNameAssetsType(graphene_sqlalchemy.SQLAlchemyObjectType):
+    domain_names = graphene.List(OxoDomainNameAssetType, required=False)
+
+    class Meta:
+        model = models.DomainAsset
+        only_fields = ("id",)
+
+    def resolve_domain_names(
+        self, info: graphql_base.ResolveInfo
+    ) -> List[OxoDomainNameAssetType]:
+        """Resolve domain names query.
+
+        Args:
+            self: The domain asset object.
+            info: GraphQL resolve info.
+
+        Returns:
+            List of domain names.
+        """
+        with models.Database() as session:
+            domain_names = (
+                session.query(models.DomainName)
+                .filter_by(domain_asset_id=self.id)
+                .all()
+            )
+            return [
+                OxoDomainNameAssetType(name=domain_name.name)
+                for domain_name in domain_names
+            ]
+
+
 class OxoAssetType(graphene.Union):
     class Meta:
         model = models.Asset
@@ -338,6 +375,7 @@ class Meta:
             OxoIOSStoreAssetType,
             OxoUrlsAssetType,
             OxoNetworkAssetType,
+            OxoDomainNameAssetsType,
         )
 
 
@@ -785,6 +823,10 @@ class OxoLinkInputType(graphene.InputObjectType):
     method = graphene.String(required=False, default_value="GET")
 
 
+class OxoDomainNameInputType(graphene.InputObjectType):
+    name = graphene.String(required=True)
+
+
 class OxoAssetInputType(graphene.InputObjectType):
     android_apk_file = graphene.List(OxoAndroidFileAssetInputType)
     android_aab_file = graphene.List(OxoAndroidFileAssetInputType)
@@ -793,6 +835,7 @@ class OxoAssetInputType(graphene.InputObjectType):
     ios_store = graphene.List(OxoIOSStoreAssetInputType)
     link = graphene.List(OxoLinkInputType)
     ip = graphene.List(OxoIPRangeInputType)
+    domain = graphene.List(OxoDomainNameInputType)
 
 
 class AgentArgumentInputType(graphene.InputObjectType):

tests/conftest.py

@@ -1076,6 +1076,18 @@ def url_asset(mocker: plugin.MockerFixture, db_engine_path: str) -> models.Urls:
     return asset
 
 
+@pytest.fixture
+def domain_asset(
+    mocker: plugin.MockerFixture, db_engine_path: str
+) -> models.DomainName:
+    """Create a DomainName asset."""
+    mocker.patch.object(models, "ENGINE_URL", db_engine_path)
+    asset = models.DomainAsset.create(
+        domains=[{"name": "google.com"}, {"name": "tesla.com"}]
+    )
+    return asset
+
+
 @pytest.fixture
 def android_file_asset(
     mocker: plugin.MockerFixture, db_engine_path: str
@@ -1198,3 +1210,23 @@ def run_scan_mock(mocker: plugin.MockerFixture) -> None:
         return_value=True,
     )
     mocker.patch("ostorlab.runtimes.local.runtime.LocalRuntime._inject_assets")
+
+
+@pytest.fixture
+def run_scan_mock2(mocker: plugin.MockerFixture) -> None:
+    """Mock functions required to run a scan."""
+    mocker.patch(
+        "ostorlab.cli.docker_requirements_checker.is_docker_installed",
+        return_value=True,
+    )
+    mocker.patch(
+        "ostorlab.cli.docker_requirements_checker.is_docker_working", return_value=True
+    )
+    mocker.patch(
+        "ostorlab.cli.docker_requirements_checker.is_swarm_initialized",
+        return_value=True,
+    )
+    mocker.patch("docker.from_env")
+    mocker.patch(
+        "ostorlab.runtimes.local.runtime.LocalRuntime.can_run", return_value=True
+    )

tests/runtimes/local/models/models_test.py

@@ -332,6 +332,28 @@ def testAssetModels_whenCreateNetwork_assetCreated(
         assert ips[1].mask == "32"
 
 
+def testAssetModels_whenCreateDomainAsset_assetCreated(
+    mocker: plugin.MockerFixture, db_engine_path: str
+) -> None:
+    """Ensure we correctly persist the domain information."""
+    mocker.patch.object(models, "ENGINE_URL", db_engine_path)
+    models.DomainAsset.create(
+        domains=[{"name": "domain1.test.ma"}, {"name": "domain2.test.ma"}]
+    )
+
+    with models.Database() as session:
+        assert session.query(models.DomainAsset).count() == 1
+        domain_id = session.query(models.DomainAsset).all()[0].id
+        domain_names = (
+            session.query(models.DomainName)
+            .filter(models.DomainName.domain_asset_id == domain_id)
+            .all()
+        )
+        assert len(domain_names) == 2
+        assert domain_names[0].name == "domain1.test.ma"
+        assert domain_names[1].name == "domain2.test.ma"
+
+
 def testAssetModels_whenCreateUrl_assetCreated(
     mocker: plugin.MockerFixture, db_engine_path: str
 ) -> None: