[Phase 2] SQLite device inventory (SQLAlchemy models)

backend/app/db.py

@@ -1,9 +1,15 @@
-"""Database connection and initialisation."""
+"""Database connection, initialisation, and session helpers."""
+
+from __future__ import annotations
 
 import os
+from typing import TYPE_CHECKING
 
 from sqlalchemy import create_engine
-from sqlalchemy.orm import DeclarativeBase, sessionmaker
+from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
+
+if TYPE_CHECKING:
+    from app.models.device import Device
 
 DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./networkcrawler.db")
 
@@ -33,3 +39,92 @@ def get_db():
         yield db
     finally:
         db.close()
+
+
+def upsert_device(
+    session: Session,
+    *,
+    ip_address: str,
+    mac_address: str | None = None,
+    vendor: str | None = None,
+    hostname: str | None = None,
+    os_guess: str | None = None,
+) -> Device:
+    """Insert or update a Device row keyed on ip_address.
+
+    If a Device with the given ip_address already exists, only non-None
+    fields are written so that richer data from a previous scan is never
+    overwritten with None.  The caller is responsible for committing.
+
+    Returns the Device instance (either existing or newly created).
+    """
+    from sqlalchemy import select
+
+    from app.models.device import Device
+
+    stmt = select(Device).where(Device.ip_address == ip_address)
+    device: Device | None = session.execute(stmt).scalar_one_or_none()
+
+    if device is None:
+        device = Device(
+            ip_address=ip_address,
+            mac_address=mac_address,
+            vendor=vendor,
+            hostname=hostname,
+            os_guess=os_guess,
+        )
+        session.add(device)
+    else:
+        if mac_address is not None:
+            device.mac_address = mac_address
+        if vendor is not None:
+            device.vendor = vendor
+        if hostname is not None:
+            device.hostname = hostname
+        if os_guess is not None:
+            device.os_guess = os_guess
+
+    return device
+
+
+def upsert_port(
+    session: Session,
+    *,
+    device_id: int,
+    port_number: int,
+    protocol: str = "tcp",
+    service_name: str | None = None,
+    version_banner: str | None = None,
+) -> Device:
+    """Insert or update a Port row keyed on (device_id, port_number, protocol).
+
+    Non-None fields overwrite existing values.  Caller must commit.
+    Returns the Port instance.
+    """
+    from sqlalchemy import select
+
+    from app.models.device import Port
+
+    stmt = select(Port).where(
+        Port.device_id == device_id,
+        Port.port_number == port_number,
+        Port.protocol == protocol,
+    )
+    port = session.execute(stmt).scalar_one_or_none()
+
+    if port is None:
+        port = Port(
+            device_id=device_id,
+            port_number=port_number,
+            protocol=protocol,
+            service_name=service_name,
+            version_banner=version_banner,
+        )
+        session.add(port)
+    else:
+        if service_name is not None:
+            port.service_name = service_name
+        if version_banner is not None:
+            port.version_banner = version_banner
+
+    return port

backend/app/models/device.py

@@ -1,17 +1,24 @@
-"""Device and Port SQLAlchemy models (stub — full implementation in Phase 2)."""
+"""Device and Port SQLAlchemy ORM models."""
 
-from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, func
+from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, UniqueConstraint, func
 from sqlalchemy.orm import relationship
 
 from app.db import Base
 
 
 class Device(Base):
+    """A network device discovered by arp-scan and/or nmap."""
+
     __tablename__ = "devices"
+    __table_args__ = (
+        # One row per IP address — upserts update in place rather than inserting duplicates.
+        UniqueConstraint("ip_address", name="uq_devices_ip_address"),
+    )
 
     id = Column(Integer, primary_key=True, index=True)
     ip_address = Column(String, nullable=False, index=True)
     mac_address = Column(String, nullable=True)
+    vendor = Column(String, nullable=True)  # hardware vendor from arp-scan OUI lookup
     hostname = Column(String, nullable=True)
     os_guess = Column(String, nullable=True)
     first_seen = Column(DateTime, default=func.now())
@@ -21,7 +28,13 @@ class Device(Base):
 
 
 class Port(Base):
+    """An open TCP/UDP port observed on a Device during an nmap scan."""
+
     __tablename__ = "ports"
+    __table_args__ = (
+        # One row per (device, port, protocol) tuple.
+        UniqueConstraint("device_id", "port_number", "protocol", name="uq_ports_device_port_proto"),
+    )
 
     id = Column(Integer, primary_key=True, index=True)
     device_id = Column(Integer, ForeignKey("devices.id"), nullable=False)