Actually add new so-allow script

salt/common/tools/sbin/so-allow

@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env python3
 
 # Copyright 2014,2015,2016,2017,2018,2019,2020,2021 Security Onion Solutions, LLC
 #
@@ -15,152 +15,195 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-. /usr/sbin/so-common
+import ipaddress
+import textwrap
+import os
+import subprocess
+import sys
+import argparse
+import re
+from lxml import etree as ET
+from xml.dom import minidom
+from datetime import datetime as dt
+from datetime import timezone as tz
 
-local_salt_dir=/opt/so/saltstack/local
-
-SKIP=0
-
-function usage {
-
-cat << EOF
-
-Usage: $0 [-abefhoprsw] [ -i IP ]
-
-This program allows you to add a firewall rule to allow connections from a new IP address or CIDR range.
-
-If you run this program with no arguments, it will present a menu for you to choose your options.
-
-If you want to automate and skip the menu, you can pass the desired options as command line arguments.
-
-EXAMPLES
-
-To add 10.1.2.3 to the analyst role:
-so-allow -a -i 10.1.2.3
-
-To add 10.1.2.0/24 to the osquery role:
-so-allow -o -i 10.1.2.0/24
-
-EOF
 
+LOCAL_SALT_DIR='/opt/so/saltstack/local'
+WAZUH_CONF='/nsm/wazuh/etc/ossec.conf'
+VALID_ROLES = {
+  'a': { 'role': 'analyst','desc': 'Analyst - 80/tcp, 443/tcp' },
+  'b': { 'role': 'beats_endpoint', 'desc': 'Logstash Beat - 5044/tcp' },
+  'e': { 'role': 'elasticsearch_rest', 'desc': 'Elasticsearch REST API - 9200/tcp' },
+  'f': { 'role': 'strelka_frontend', 'desc': 'Strelka frontend - 57314/tcp' },
+  'o': { 'role': 'osquery_endpoint', 'desc': 'Osquery endpoint - 8090/tcp' },
+  's': { 'role': 'syslog', 'desc': 'Syslog device - 514/tcp/udp' },
+  'w': { 'role': 'wazuh_agent', 'desc': 'Wazuh agent - 1514/tcp/udp' },
+  'p': { 'role': 'wazuh_api', 'desc': 'Wazuh API - 55000/tcp' },
+  'r': { 'role': 'wazuh_authd', 'desc': 'Wazuh registration service - 1515/tcp' }
 }
 
-while getopts "ahfesprbowi:" OPTION
-do
-    case $OPTION in
-        h)
-            usage
-            exit 0
-            ;;
-        a)
-            FULLROLE="analyst"
-            SKIP=1
-            ;;
-        b)
-            FULLROLE="beats_endpoint"
-            SKIP=1
-            ;;
-	e)
-            FULLROLE="elasticsearch_rest"
-            SKIP=1
-            ;;
-        f)
-	    FULLROLE="strelka_frontend"
-            SKIP=1
-            ;;
-        i)  IP=$OPTARG
-            ;;
-        o)
-            FULLROLE="osquery_endpoint"
-            SKIP=1
-            ;;
-        w)
-            FULLROLE="wazuh_agent"
-            SKIP=1
-            ;;
-        s)
-            FULLROLE="syslog"
-            SKIP=1
-            ;;
-        p)
-            FULLROLE="wazuh_api"
-            SKIP=1
-            ;;
-        r)
-            FULLROLE="wazuh_authd"
-            SKIP=1
-            ;;
-        *)
-            usage
-            exit 0
-            ;;
-    esac
-done
 
-if [ "$SKIP" -eq 0 ]; then
+def validate_ip_cidr(ip_cidr: str) -> bool:
+  try:
+    ipaddress.ip_address(ip_cidr)
+  except ValueError:
+    try:
+      ipaddress.ip_network(ip_cidr)
+    except ValueError:
+      return False
+  return True
 
-    echo "This program allows you to add a firewall rule to allow connections from a new IP address."
-    echo ""
-    echo "Choose the role for the IP or Range you would like to add"
-    echo ""
-    echo "[a] - Analyst - ports 80/tcp and 443/tcp"
-    echo "[b] - Logstash Beat - port 5044/tcp"
-    echo "[e] - Elasticsearch REST API - port 9200/tcp"
-    echo "[f] - Strelka frontend - port 57314/tcp"
-    echo "[o] - Osquery endpoint - port 8090/tcp"
-    echo "[s] - Syslog device - 514/tcp/udp"
-    echo "[w] - Wazuh agent - port 1514/tcp/udp"
-    echo "[p] - Wazuh API - port 55000/tcp"
-    echo "[r] - Wazuh registration service - 1515/tcp"
-    echo ""
-    echo "Please enter your selection:"
-    read -r ROLE
-    echo "Enter a single ip address or range to allow (example: 10.10.10.10 or 10.10.0.0/16):"
-    read -r IP
 
-    if [ "$ROLE" == "a" ]; then
+def role_prompt() -> str:
-        FULLROLE=analyst
+  print()
-    elif [ "$ROLE" == "b" ]; then
+  print('Choose the role for the IP or Range you would like to allow')
-        FULLROLE=beats_endpoint
+  print()
-    elif [ "$ROLE" == "e" ]; then
+  for role in VALID_ROLES:
-        FULLROLE=elasticsearch_rest
+    print(f'[{role}] - {VALID_ROLES[role]["desc"]}')
-    elif [ "$ROLE" == "f" ]; then
+  print()
-        FULLROLE=strelka_frontend
+  role = input('Please enter your selection: ')
-    elif [ "$ROLE" == "o" ]; then
+  if role in VALID_ROLES.keys():
-        FULLROLE=osquery_endpoint
+    return VALID_ROLES[role]['role']
-    elif [ "$ROLE" == "w" ]; then
+  else:
-        FULLROLE=wazuh_agent
+    print(f'Invalid role \'{role}\', please try again.', file=sys.stderr)
-    elif [ "$ROLE" == "s" ]; then
+    sys.exit(1)
-        FULLROLE=syslog
-    elif [ "$ROLE" == "p" ]; then
-        FULLROLE=wazuh_api
-    elif [ "$ROLE" == "r" ]; then
-        FULLROLE=wazuh_authd
-    else
-        echo "I don't recognize that role"
-        exit 1
-    fi
   
-fi
 
-echo "Adding $IP to the $FULLROLE role. This can take a few seconds"
+def ip_prompt() -> str:
-/usr/sbin/so-firewall includehost $FULLROLE $IP
+  ip = input('Enter a single ip address or range to allow (ex: 10.10.10.10 or 10.10.0.0/16): ')
-salt-call state.apply firewall queue=True
+  if validate_ip_cidr(ip):
+    return ip
+  else:
+    print(f'Invalid IP address or CIDR block \'{ip}\', please try again.', file=sys.stderr)
+    sys.exit(1)
+
+
+def wazuh_enabled() -> bool:
+  for file in os.listdir(f'{LOCAL_SALT_DIR}/pillar'):
+    with open(file, 'r') as pillar:
+      if 'wazuh: 1' in pillar.read():
+        return True
+  return False
+
+
+def root_to_str(root: ET.ElementTree) -> str:
+    xml_str = ET.tostring(root, encoding='unicode', method='xml').replace('\n', '')
+    xml_str = re.sub(r'(?:(?<=>) *)', '', xml_str)
+    xml_str = re.sub(r'       -', '', xml_str)
+    xml_str = re.sub(r'      -->', ' -->', xml_str)
+    dom = minidom.parseString(xml_str)
+    return dom.toprettyxml(indent="  ")
+
+
+def add_wl(ip):
+    parser = ET.XMLParser(remove_blank_text=True)
+    with open(WAZUH_CONF, 'rb') as wazuh_conf:
+        tree = ET.parse(wazuh_conf, parser)
+    root = tree.getroot()
+
+    source_comment = ET.Comment(f'Address {ip} added by /usr/sbin/so-allow on {dt.utcnow().replace(tzinfo=tz.utc).strftime("%a %b %e %H:%M:%S %Z %Y")}')
+    new_global = ET.Element("global")
+    new_wl = ET.SubElement(new_global, 'white_list')
+    new_wl.text = ip
+
+    root.append(source_comment)
+    root.append(new_global)
+
+    with open(WAZUH_CONF, 'w') as add_out:
+        add_out.write(root_to_str(root))
+
+
+def apply(role: str, ip: str) -> int:
+  firewall_cmd = ['so-firewall', 'includehost', role, ip]
+  salt_cmd = ['salt-call', 'state.apply', '-l', 'quiet', 'firewall', 'queue=True']
+  restart_wazuh_cmd = ['so-wazuh-restart']
+  print(f'Adding {ip} to the {role} role. This can take a few seconds...')
+  cmd = subprocess.run(firewall_cmd)
+  if cmd.returncode == 0:
+    cmd = subprocess.run(salt_cmd, stdout=subprocess.DEVNULL)
+  else:
+    return cmd.returncode
+  if cmd.returncode == 0:
+    if wazuh_enabled and role=='analyst':
+      try:
+        add_wl(ip)
+        print(f'Added whitelist entry for {ip} from {WAZUH_CONF}', file=sys.stderr)
+      except Exception as e:
+        print(f'Failed to add whitelist entry for {ip} from {WAZUH_CONF}', file=sys.stderr)
+        print(e)
+        return 1
+      print('Restarting OSSEC Server...')
+      cmd = subprocess.run(restart_wazuh_cmd)
+    else:
+      return cmd.returncode
+  else:
+    print(f'Commmand \'{" ".join(salt_cmd)}\' failed.', file=sys.stderr)
+    return cmd.returncode
+  if cmd.returncode != 0:
+    print('Failed to restart OSSEC server.')
+  return cmd.returncode
+
+
+def main():
+  if os.geteuid() != 0:
+    print('You must run this script as root', file=sys.stderr)
+    sys.exit(1)
+
+  main_parser = argparse.ArgumentParser(
+    formatter_class=argparse.RawDescriptionHelpFormatter,
+    epilog=textwrap.dedent(f'''\
+        additional information:
+                      To use this script in interactive mode call it with no arguments
+        '''
+  ))
+
+  group = main_parser.add_argument_group(title='roles')
+  group.add_argument('-a', dest='roles', action='append_const', const=VALID_ROLES['a']['role'], help="Analyst - 80/tcp, 443/tcp")
+  group.add_argument('-b', dest='roles', action='append_const', const=VALID_ROLES['b']['role'], help="Logstash Beat - 5044/tcp")
+  group.add_argument('-e', dest='roles', action='append_const', const=VALID_ROLES['e']['role'], help="Elasticsearch REST API - 9200/tcp")
+  group.add_argument('-f', dest='roles', action='append_const', const=VALID_ROLES['f']['role'], help="Strelka frontend - 57314/tcp")
+  group.add_argument('-o', dest='roles', action='append_const', const=VALID_ROLES['o']['role'], help="Osquery endpoint - 8090/tcp")
+  group.add_argument('-s', dest='roles', action='append_const', const=VALID_ROLES['s']['role'], help="Syslog device - 514/tcp/udp")
+  group.add_argument('-w', dest='roles', action='append_const', const=VALID_ROLES['w']['role'], help="Wazuh agent - 1514/tcp/udp")
+  group.add_argument('-p', dest='roles', action='append_const', const=VALID_ROLES['p']['role'], help="Wazuh API - 55000/tcp")
+  group.add_argument('-r', dest='roles', action='append_const', const=VALID_ROLES['r']['role'], help="Wazuh registration service - 1515/tcp")
+  
+  ip_g = main_parser.add_argument_group(title='allow')
+  ip_g.add_argument('-i', help="IP or CIDR block to disallow connections from, requires at least one role argument", metavar='', dest='ip')
+
+  args = main_parser.parse_args(sys.argv[1:])
+
+  if args.roles is None:
+    role = role_prompt()
+    ip = ip_prompt()
+    try:
+      return_code = apply(role, ip)
+    except Exception as e:
+      print(f'Unexpected exception occurred: {e}', file=sys.stderr)
+      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
+    else:
+      print(f'Invalid IP address or CIDR block \'{args.ip}\', please try again.', file=sys.stderr)
+      return_code = 1
+
+  sys.exit(return_code)
+
+
+if __name__ == '__main__':
+  try:
+    main()
+  except KeyboardInterrupt:
+    sys.exit(1)
 
-# Check if Wazuh enabled
-if grep -q -R "wazuh: 1" $local_salt_dir/pillar/*; then
-    # If analyst, add to Wazuh AR whitelist
-    if [ "$FULLROLE" == "analyst" ]; then
-        WAZUH_MGR_CFG="/nsm/wazuh/etc/ossec.conf"
-        if ! grep -q "<white_list>$IP</white_list>" $WAZUH_MGR_CFG ; then
-            DATE=$(date)
-            sed -i 's/<\/ossec_config>//' $WAZUH_MGR_CFG
-            sed -i '/^$/N;/^\n$/D' $WAZUH_MGR_CFG
-            echo -e "<!--Address $IP added by /usr/sbin/so-allow on \"$DATE\"-->\n  <global>\n    <white_list>$IP</white_list>\n  </global>\n</ossec_config>" >> $WAZUH_MGR_CFG
-            echo "Added whitelist entry for $IP in $WAZUH_MGR_CFG."
-        echo
-            echo "Restarting OSSEC Server..."
-            /usr/sbin/so-wazuh-restart
-        fi
-    fi
-fi