From f374dcbb587fdf34a1b2819a225c545235a7614c Mon Sep 17 00:00:00 2001
From: William Wernert <william.wernert@securityonionsolutions.com>
Date: Thu, 21 Oct 2021 13:54:06 -0400
Subject: [PATCH] Check for IP environment variable in so-allow and so-deny

---
 salt/common/tools/sbin/so-allow | 32 ++++++++++++++++++--------------
 salt/common/tools/sbin/so-deny  | 32 ++++++++++++++++++--------------
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/salt/common/tools/sbin/so-allow b/salt/common/tools/sbin/so-allow
index 1d240d840..769dcc1e9 100755
--- a/salt/common/tools/sbin/so-allow
+++ b/salt/common/tools/sbin/so-allow
@@ -183,21 +183,25 @@ def main():
       return_code = e.errno
     sys.exit(return_code)
   elif args.roles is not None and args.ip is None:
-    main_parser.print_help()
-  else:
-    if validate_ip_cidr(args.ip):
-      try:
-        for role in args.roles:
-          return_code = apply(role, args.ip)
-          if return_code > 0:
-            break
-      except Exception as e:
-        print(f'Unexpected exception occurred: {e}', file=sys.stderr)
-        return_code = e.errno
+    if os.environ.get('IP') is None:
+      main_parser.print_help()
+      sys.exit(1)
     else:
-      print(f'Invalid IP address or CIDR block \'{args.ip}\', please try again.', file=sys.stderr)
-      return_code = 1
-
+      args.ip = os.environ['IP']
+  
+  if validate_ip_cidr(args.ip):
+    try:
+      for role in args.roles:
+        return_code = apply(role, args.ip)
+        if return_code > 0:
+          break
+    except Exception as e:
+      print(f'Unexpected exception occurred: {e}', file=sys.stderr)
+      return_code = e.errno
+  else:
+    print(f'Invalid IP address or CIDR block \'{args.ip}\', please try again.', file=sys.stderr)
+    return_code = 1
+    
   sys.exit(return_code)
 
 
diff --git a/salt/common/tools/sbin/so-deny b/salt/common/tools/sbin/so-deny
index c36a9b9d6..c13ea3f32 100755
--- a/salt/common/tools/sbin/so-deny
+++ b/salt/common/tools/sbin/so-deny
@@ -184,21 +184,25 @@ def main():
       return_code = e.errno
     sys.exit(return_code)
   elif args.roles is not None and args.ip is None:
-    main_parser.print_help()
-  else:
-    if validate_ip_cidr(args.ip):
-      try:
-        for role in args.roles:
-          return_code = apply(role, args.ip)
-          if return_code > 0:
-            break
-      except Exception as e:
-        print(f'Unexpected exception occurred: {e}', file=sys.stderr)
-        return_code = e.errno
+    if os.environ.get('IP') is None:
+      main_parser.print_help()
+      sys.exit(1)
     else:
-      print(f'Invalid IP address or CIDR block \'{args.ip}\', please try again.', file=sys.stderr)
-      return_code = 1
-
+      args.ip = os.environ['IP']
+  
+  if validate_ip_cidr(args.ip):
+    try:
+      for role in args.roles:
+        return_code = apply(role, args.ip)
+        if return_code > 0:
+          break
+    except Exception as e:
+      print(f'Unexpected exception occurred: {e}', file=sys.stderr)
+      return_code = e.errno
+  else:
+    print(f'Invalid IP address or CIDR block \'{args.ip}\', please try again.', file=sys.stderr)
+    return_code = 1
+    
   sys.exit(return_code)