2fa policy check working

server/routers/org/updateOrg.ts

@@ -1,7 +1,7 @@
 import { Request, Response, NextFunction } from "express";
 import { z } from "zod";
 import { db } from "@server/db";
-import { orgs } from "@server/db";
+import { orgs, users } from "@server/db";
 import { eq } from "drizzle-orm";
 import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
@@ -9,6 +9,11 @@ import createHttpError from "http-errors";
 import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { OpenAPITags, registry } from "@server/openApi";
+import { build } from "@server/build";
+import { UserType } from "@server/types/UserTypes";
+import license from "#dynamic/license/license";
+import { getOrgTierData } from "#dynamic/lib/billing";
+import { TierId } from "@server/lib/billing/tiers";
 
 const updateOrgParamsSchema = z
     .object({
@@ -18,7 +23,8 @@ const updateOrgParamsSchema = z
 
 const updateOrgBodySchema = z
     .object({
-        name: z.string().min(1).max(255).optional()
+        name: z.string().min(1).max(255).optional(),
+        requireTwoFactor: z.boolean().optional()
     })
     .strict()
     .refine((data) => Object.keys(data).length > 0, {
@@ -71,10 +77,30 @@ export async function updateOrg(
 
         const { orgId } = parsedParams.data;
 
+        const isLicensed = await isLicensedOrSubscribed(orgId);
+        if (!isLicensed) {
+            parsedBody.data.requireTwoFactor = undefined;
+        }
+
+        if (
+            req.user &&
+            req.user.type === UserType.Internal &&
+            parsedBody.data.requireTwoFactor === true &&
+            !req.user.twoFactorEnabled
+        ) {
+            return next(
+                createHttpError(
+                    HttpCode.FORBIDDEN,
+                    "You must enable two-factor authentication for your account before enforcing it for all users"
+                )
+            );
+        }
+
         const updatedOrg = await db
             .update(orgs)
             .set({
-                name: parsedBody.data.name
+                name: parsedBody.data.name,
+                requireTwoFactor: parsedBody.data.requireTwoFactor
             })
             .where(eq(orgs.orgId, orgId))
             .returning();
@@ -102,3 +128,22 @@ export async function updateOrg(
         );
     }
 }
+
+async function isLicensedOrSubscribed(orgId: string): Promise<boolean> {
+    if (build === "enterprise") {
+        const isUnlocked = await license.isUnlocked();
+        if (!isUnlocked) {
+            return false;
+        }
+    }
+
+    if (build === "saas") {
+        const { tier } = await getOrgTierData(orgId);
+        const subscribed = tier === TierId.STANDARD;
+        if (!subscribed) {
+            return false;
+        }
+    }
+
+    return true;
+}