From f356be110e7e21f3f707ce6eaba16da8718f203c Mon Sep 17 00:00:00 2001
From: Eyad Arafat <eyad.mba@gmail.com>
Date: Thu, 4 Apr 2024 01:42:38 -0500
Subject: [PATCH] refactor: modular format registration

---
 src/IntegrationSchema.ts | 14 ++------------
 src/index.ts             |  4 +++-
 src/registerFormats.ts   | 24 ++++++++++++++++++++++++
 3 files changed, 29 insertions(+), 13 deletions(-)
 create mode 100644 src/registerFormats.ts

diff --git a/src/IntegrationSchema.ts b/src/IntegrationSchema.ts
index 3e95166..c031bcb 100644
--- a/src/IntegrationSchema.ts
+++ b/src/IntegrationSchema.ts
@@ -1,10 +1,5 @@
 import Ajv from 'ajv';
-import addFormats from 'ajv-formats';
-
-const ipv4 = addFormats.get('ipv4') as RegExp;
-const ipv6 = addFormats.get('ipv6') as RegExp;
-const ipv4CidrRegex = /^(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\/(3[0-2]|[12]?[0-9])$/;
-const ipv6CidrRegex = /^([0-9a-fA-F]{1,4}:){7,7}[0-9a-fA-F]{1,4}\/(12[0-8]|1[01][0-9]|[1-9]?[0-9])$|^([0-9a-fA-F]{1,4}:){1,7}:\/(12[0-8]|1[01][0-9]|[1-9]?[0-9])$|^::\/(12[0-8]|1[01][0-9]|[1-9]?[0-9])$|^([0-9a-fA-F]{1,4}:){1,7}[0-9a-fA-F]{1,4}$/;
+import { registerFormats } from './registerFormats';
 
 // JSON Schema allows an object to contain properties that are not specified by
 // the schema. This can be disabled with `additionalProperties: false`. Ajv then
@@ -24,14 +19,9 @@ const ipv6CidrRegex = /^([0-9a-fA-F]{1,4}:){7,7}[0-9a-fA-F]{1,4}\/(12[0-8]|1[01]
 export const IntegrationSchema = new Ajv({
   // Ignore "excludes", "multiple"; used in JupiterOne UI?
   strictSchema: false,
-  formats: {
-    ip: (x) => ipv4.test(x) || ipv6.test(x),
-    ipCidr: (x) => ipv4CidrRegex.test(x) || ipv6CidrRegex.test(x),
-  },
 });
 
-// Install ajv-formats
-addFormats(IntegrationSchema);
+registerFormats(IntegrationSchema);
 
 // Schema Imports : generated by tools/generate-schema-imports.sh
 import WorkloadJson from './schemas/Workload.json';
diff --git a/src/index.ts b/src/index.ts
index 0176f5a..50a1aa0 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,3 +1,6 @@
+export { registerSchemas } from './registerSchemas';
+export { registerFormats } from './registerFormats';
+
 export { EVERYONE, INTERNET } from './globalEntities';
 
 export { RelationshipClass } from './RelationshipClass';
@@ -21,5 +24,4 @@ export type EntityClass = keyof typeof entitySchemas;
 export const entityClasses = Object.keys(entitySchemas) as EntityClass[];
 
 export { validateEntityWithSchema } from './validateEntityWithSchema';
-export { registerSchemas } from './registerSchemas';
 export { getSchema } from './getSchema';
diff --git a/src/registerFormats.ts b/src/registerFormats.ts
new file mode 100644
index 0000000..d39a818
--- /dev/null
+++ b/src/registerFormats.ts
@@ -0,0 +1,24 @@
+import addFormats from 'ajv-formats';
+import Ajv from 'ajv';
+
+const ipv4 = addFormats.get('ipv4') as RegExp;
+const ipv6 = addFormats.get('ipv6') as RegExp;
+const ipv4CidrRegex = /^(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\/(3[0-2]|[12]?[0-9])$/;
+const ipv6CidrRegex = /^([0-9a-fA-F]{1,4}:){7,7}[0-9a-fA-F]{1,4}\/(12[0-8]|1[01][0-9]|[1-9]?[0-9])$|^([0-9a-fA-F]{1,4}:){1,7}:\/(12[0-8]|1[01][0-9]|[1-9]?[0-9])$|^::\/(12[0-8]|1[01][0-9]|[1-9]?[0-9])$|^([0-9a-fA-F]{1,4}:){1,7}[0-9a-fA-F]{1,4}$/;
+
+const isValidIpFormat = {
+  name: 'ip',
+  fn: (x) => ipv4.test(x) || ipv6.test(x),
+};
+
+const isValidIpCidrFormat = {
+  name: 'ipCidr',
+  fn: (x) => ipv4CidrRegex.test(x) || ipv6CidrRegex.test(x),
+};
+
+export function registerFormats(ajvInstance: Ajv): void {
+  ajvInstance.addFormat(isValidIpFormat.name, isValidIpFormat.fn);
+  ajvInstance.addFormat(isValidIpCidrFormat.name, isValidIpCidrFormat.fn);
+
+  addFormats(ajvInstance);
+}