Atomic postures and mfa-flows assignments

crates/defguard_common/src/db/models/mfa_flow.rs

@@ -49,19 +49,20 @@ pub struct MfaFlowSnapshot {
     pub steps: Vec<MfaFlowStep<Id>>,
 }
 
-/// Assignment of an MFA flow to a location, enriched for API consumption.
+/// MFA flow assignment with location metadata.
 #[derive(Clone, Debug, Serialize)]
 pub struct LocationMfaFlowItem {
     pub id: Id,
     pub title: String,
     pub step_count: i64,
+    pub group_ids: Vec<Id>,
     pub group_names: Vec<String>,
     pub position: i32,
     pub is_default: bool,
 }
 
 /// Input for a single flow assignment to a location.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize, ToSchema)]
 pub struct LocationMfaFlowAssignment {
     pub flow_id: Id,
     pub is_default: bool,
@@ -117,6 +118,19 @@ pub enum MfaFlowAssignmentError {
     Sqlx(#[from] sqlx::Error),
 }
 
+/// License-related errors that can occur during MFA flow assignment.
+#[derive(Debug, Error)]
+pub enum MfaFlowAssignmentLicenseError {
+    #[error("MFA flow group assignments require an Enterprise license")]
+    GroupAssignmentNotAllowed,
+    #[error("Multi-step MFA flows require a Business license")]
+    MultipleStepsNotAllowed,
+    #[error("Multiple MFA flows can only be assigned with Business license")]
+    MultipleMfaFlowsNotAllowed,
+    #[error(transparent)]
+    Database(#[from] sqlx::Error),
+}
+
 /// Errors that can occur when updating an MFA flow.
 #[derive(Debug, Error)]
 pub enum MfaFlowUpdateError {
@@ -516,6 +530,44 @@ impl MfaFlow<Id> {
         Ok(())
     }
 
+    /// Validates whether the current license permits the requested MFA flow assignments.
+    pub async fn validate_mfa_flow_assignments_license(
+        conn: &mut PgConnection,
+        assignments: &[LocationMfaFlowAssignment],
+        has_enterprise_access: bool,
+        is_business_license_active: bool,
+    ) -> Result<(), MfaFlowAssignmentLicenseError> {
+        // Enterprise can make all assignments.
+        if has_enterprise_access {
+            return Ok(());
+        }
+
+        // Business and Free can't assign groups.
+        if assignments.iter().any(|a| !a.group_ids.is_empty()) {
+            return Err(MfaFlowAssignmentLicenseError::GroupAssignmentNotAllowed);
+        }
+
+        // Business can assign multiple and multi-step flows.
+        if is_business_license_active {
+            return Ok(());
+        }
+
+        // Free can't assign multi-step flows.
+        if let Some(assignment) = assignments.first() {
+            let steps = MfaFlowStep::find_by_flow(&mut *conn, assignment.flow_id).await?;
+            if steps.len() > 1 {
+                return Err(MfaFlowAssignmentLicenseError::MultipleStepsNotAllowed);
+            }
+        }
+
+        // Free can't assign multiple flows.
+        if assignments.len() > 1 {
+            return Err(MfaFlowAssignmentLicenseError::MultipleMfaFlowsNotAllowed);
+        }
+
+        Ok(())
+    }
+
     /// Returns the enriched assignment list for a location, ordered by position.
     pub async fn for_location<'e, E: PgExecutor<'e>>(
         executor: E,
@@ -525,7 +577,10 @@ impl MfaFlow<Id> {
             LocationMfaFlowItem,
             "SELECT mf.id, mf.title, \
              COALESCE(s.step_count, 0) AS \"step_count!: i64\", \
-             COALESCE(array_agg(g.name ORDER BY g.name) \
+             COALESCE(array_agg(lmfg.group_id ORDER BY lmfg.group_id) \
+                      FILTER (WHERE lmfg.group_id IS NOT NULL), '{}') \
+                      AS \"group_ids!: Vec<Id>\", \
+             COALESCE(array_agg(g.name ORDER BY lmfg.group_id) \
                       FILTER (WHERE g.name IS NOT NULL), '{}') \
                       AS \"group_names!: Vec<String>\", \
              lmf.position, lmf.is_default \

crates/defguard_common/src/db/models/wizard.rs

@@ -122,15 +122,13 @@ impl Wizard {
         .fetch_one(executor)
         .await?;
 
-        let active_wizard;
-
-        if has_auto_adopt_flags {
-            active_wizard = ActiveWizard::AutoAdoption;
+        let active_wizard = if has_auto_adopt_flags {
+            ActiveWizard::AutoAdoption
         } else if is_fresh_instance {
-            active_wizard = ActiveWizard::Initial;
+            ActiveWizard::Initial
         } else {
-            active_wizard = ActiveWizard::Migration;
-        }
+            ActiveWizard::Migration
+        };
 
         wizard.active_wizard = active_wizard;
 

crates/defguard_core/src/enterprise/handlers/device_posture.rs

@@ -1114,84 +1114,12 @@ pub async fn duplicate_device_posture(
     Ok(ApiResponse::json(response, StatusCode::CREATED))
 }
 
-/// Request body for assigning posture checks to a VPN location.
-#[derive(Clone, Debug, Deserialize, Serialize, ToSchema)]
-pub struct AssignPosturesData {
-    pub postures: Vec<Id>,
-}
-
 /// Request body for assigning VPN locations to a posture check.
 #[derive(Clone, Debug, Deserialize, Serialize, ToSchema)]
 pub struct AssignLocationsData {
     pub locations: Vec<Id>,
 }
 
-/// Assign device posture check policies to a location
-///
-/// Replaces the current assignment.
-#[utoipa::path(
-    put,
-    path = "/api/v1/network/{id}/postures",
-    tag = "device posture",
-    params(
-        ("id" = i64, Path, description = "ID of the location.")
-    ),
-    request_body = AssignPosturesData,
-    responses(
-        (status = 200, description = "Device posture check policies assigned to the location.", body = [Id]),
-        (status = 400, description = "Posture checks cannot be assigned to a service location.", body = ApiErrorResponse, example = json!({"msg": "Posture checks cannot be assigned to service locations"})),
-        (status = 401, description = "Session is missing or invalid.", body = ApiErrorResponse, example = json!({"msg": "Session is required"})),
-        (status = 403, description = "Requires admin privileges and an active enterprise license.", body = ApiErrorResponse, example = json!({"msg": "requires privileged access"})),
-        (status = 404, description = "Location not found.", body = ApiErrorResponse, example = json!({"msg": "Location 1 not found"})),
-        (status = 500, description = "Unable to assign device posture check policies to the location.", body = ApiErrorResponse, example = json!({"msg": "Internal server error"}))
-    ),
-    security(
-        ("cookie" = []),
-        ("api_token" = [])
-    )
-)]
-pub async fn set_postures_for_location(
-    _license: LicenseGated<DevicePostureFeature>,
-    _admin: AdminRole,
-    session: SessionInfo,
-    context: ApiRequestContext,
-    Path(location_id): Path<Id>,
-    State(appstate): State<AppState>,
-    Json(data): Json<AssignPosturesData>,
-) -> ApiResult {
-    debug!(
-        "User {} assigning device posture checks {:?} to location {location_id}",
-        session.user.username, data.postures
-    );
-
-    let location = WireguardNetwork::find_by_id(&appstate.pool, location_id)
-        .await?
-        .ok_or_else(|| WebError::ObjectNotFound(format!("Location {location_id} not found")))?;
-
-    let mut tx = appstate.pool.begin().await?;
-    let old_postures = DevicePostureLocation::find_by_location(&mut *tx, location_id).await?;
-    let result =
-        DevicePostureLocation::set_for_location(&mut tx, location_id, &data.postures).await?;
-    let gateway_commands = if same_id_set(&old_postures, &result) {
-        Vec::new()
-    } else {
-        build_location_peer_refresh_commands(&mut tx, [location_id]).await?
-    };
-    tx.commit().await?;
-
-    appstate.send_multiple_gateway_commands(gateway_commands);
-
-    appstate.emit_event(ApiEvent {
-        context,
-        event: Box::new(ApiEventType::LocationPosturesAssigned {
-            location,
-            posture_ids: result.clone(),
-        }),
-    })?;
-
-    Ok(ApiResponse::json(result, StatusCode::OK))
-}
-
 /// Assign locations to a device posture check policy
 ///
 /// Replaces the current assignment.