From 210da4969d0043349c7d3d897ef24ab004bbb2e1 Mon Sep 17 00:00:00 2001
From: Jayson Jacobs <j4ys0n@gmail.com>
Date: Fri, 12 Jul 2024 15:44:22 -0600
Subject: [PATCH] add cidr groups

---
 example.env                           |  2 +-
 package.json                          |  5 +++--
 src/controllers/nginx.ts              | 20 ++++++++++++--------
 src/static/nginx-server-template.conf |  5 +----
 src/utils/nginx.ts                    |  7 +++++--
 5 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/example.env b/example.env
index e733537..e7b8386 100644
--- a/example.env
+++ b/example.env
@@ -1,4 +1,4 @@
-PORT=3000
+PORT=8080
 TARGET_URLS=http://localhost:1234/v1
 JWT_SECRET=your-jwt-secret-key-here
 AUTH_USERNAME=admin
diff --git a/package.json b/package.json
index 013ab96..6c8c4fb 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "llm-proxy",
-  "version": "1.1.5",
+  "version": "1.2.0",
   "description": "Manages Nginx for reverse proxy to multiple LLMs, with TLS & Bearer Auth tokens",
   "main": "dist/index.js",
   "scripts": {
@@ -18,7 +18,8 @@
     "openai",
     "certificate",
     "bearer auth",
-    "tls"
+    "tls",
+    "ai"
   ],
   "author": "Jayson Jacobs",
   "license": "Apache-2.0",
diff --git a/src/controllers/nginx.ts b/src/controllers/nginx.ts
index b5e1238..28d2325 100644
--- a/src/controllers/nginx.ts
+++ b/src/controllers/nginx.ts
@@ -13,11 +13,11 @@ export class NginxController {
   }
 
   public registerRoutes(): void {
-    this.app.post('/nginx/reload', ...this.requestHandlers, this.reloadNginx.bind(this))
+    this.app.get('/nginx/reload', ...this.requestHandlers, this.reloadNginx.bind(this))
     this.app.post('/nginx/config/update', ...this.requestHandlers, this.updateConfig.bind(this))
     this.app.get('/nginx/config/get', ...this.requestHandlers, this.getConfig.bind(this))
     this.app.get('/nginx/config/get-default', ...this.requestHandlers, this.getDefaultConfig.bind(this))
-    this.app.get('/nginx/config/write-default', ...this.requestHandlers, this.writeDefaultConfig.bind(this))
+    this.app.post('/nginx/config/write-default', ...this.requestHandlers, this.writeDefaultConfig.bind(this))
     this.app.post('/nginx/certificates/obtain', ...this.requestHandlers, this.obtainCertificates.bind(this))
     this.app.get('/nginx/certificates/renew', ...this.requestHandlers, this.renewCertificates.bind(this))
     log('info', 'NginxController initialized')
@@ -60,13 +60,17 @@ export class NginxController {
   }
 
   private async writeDefaultConfig(req: Request, res: Response): Promise<void> {
-    if (req.body != null && req.body.domain != null) {
-      const domain = req.body.domain
-      const { success, message } = await this.nginxManager.writeDefaultTemplate(domain)
-      if (success) {
-        res.json({ success, message: 'Default config written successfully' })
+    if (req.body != null && req.body.domain != null && req.body.cidrGroups != null) {
+      const { domain, cidrGroups } = req.body
+      if (Array.isArray(cidrGroups) && typeof domain === 'string') {
+        const { success, message } = await this.nginxManager.writeDefaultTemplate(domain, cidrGroups)
+        if (success) {
+          res.json({ success, message: 'Default config written successfully' })
+        } else {
+          res.status(500).json({ success, message })
+        }
       } else {
-        res.status(500).json({ success, message })
+        res.status(400).json({ success: false, message: 'Invalid request body' })  
       }
     } else {
       res.status(400).json({ success: false, message: 'Invalid request body' })
diff --git a/src/static/nginx-server-template.conf b/src/static/nginx-server-template.conf
index 79f1773..3f095ac 100644
--- a/src/static/nginx-server-template.conf
+++ b/src/static/nginx-server-template.conf
@@ -44,10 +44,7 @@ server {
     proxy_buffering off;
     client_max_body_size 0;
     proxy_read_timeout 36000s;
-    allow 10.1.0.0/16;
-    allow 10.6.0.0/24;
-    allow 10.9.9.0/24;
-    allow 10.99.10.0/24;
+{{allowedIPs}}
     deny all;
   }
 
diff --git a/src/utils/nginx.ts b/src/utils/nginx.ts
index a3c8d79..221d7c1 100644
--- a/src/utils/nginx.ts
+++ b/src/utils/nginx.ts
@@ -62,9 +62,12 @@ export class NginxManager {
     return this.putFile(this.configPath, newConfig)
   }
 
-  async writeDefaultTemplate(domain: string): Promise<NginxResponse> {
+  async writeDefaultTemplate(domain: string, cidrGroups: string[]): Promise<NginxResponse> {
     const templateContent = await readFile(CONFIG_TEMPLATE_PATH, 'utf-8')
-    const content = templateContent.replace(/{{domainName}}/g, domain)
+    const allowedIPs = cidrGroups.map((g) => `    allow ${g};\n`).reduce((acc, curr) => acc + curr, '')
+    const content = templateContent
+    .replace(/{{domainName}}/g, domain)
+    .replace(/{{allowedIPs}}/g, allowedIPs)
     return this.putFile(this.configPath, content)
   }