#!/usr/bin/env python3

# Copyright 2014-2023 Security Onion Solutions, LLC
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import subprocess
import sys
import time
import yaml

lockFile = "/tmp/so-firewall.lock"
hostgroupsFilename = "/opt/so/saltstack/local/pillar/firewall/soc_firewall.sls"
defaultsFilename = "/opt/so/saltstack/default/salt/firewall/defaults.yaml"

def showUsage(options, args):
  print('Usage: {} [OPTIONS] <COMMAND> [ARGS...]'.format(sys.argv[0]))
  print('  Options:')
  print('   --apply         - After updating the firewall configuration files, apply the new firewall state')
  print('')
  print('  General commands:')
  print('    help           - Prints this usage information.')
  print('    apply          - Apply the firewall state.')
  print('')
  print('  Host commands:')
  print('    includehost    - Includes the given IP in the given group. Args: <GROUP_NAME> <IP>')
  print('    addhostgroup   - Adds a new, custom host group. Args: <GROUP_NAME>')
  print('')
  print('  Where:')
  print('   GROUP_NAME    - The name of an alias group (Ex: analyst)')
  print('   IP            - Either a single IP address (Ex: 8.8.8.8) or a CIDR block (Ex: 10.23.0.0/16).')
  sys.exit(1)

def checkApplyOption(options):
  if "--apply" in options:
    return apply(None, None)

def loadYaml(filename):
  file = open(filename, "r")
  content = file.read()
  return yaml.safe_load(content)

def writeYaml(filename, content):
  file = open(filename, "w")
  return yaml.dump(content, file)

def addIp(name, ip):
  content = loadYaml(hostgroupsFilename)
  defaults = loadYaml(defaultsFilename)
  allowedHostgroups = defaults['firewall']['hostgroups']
  unallowedHostgroups = ['anywhere', 'dockernet', 'localhost', 'self']
  for hg in unallowedHostgroups:
    allowedHostgroups.pop(hg)
  if not content:
    content = {'firewall': {'hostgroups': {name: []}}}
  if name in allowedHostgroups:
    if name not in content['firewall']['hostgroups']:
      hostgroup = content['firewall']['hostgroups'].update({name: [ip]})
    else:
      hostgroup = content['firewall']['hostgroups'][name]
  else:
    print('Host group not defined in salt/firewall/defaults.yaml or hostgroup name is unallowed.', file=sys.stderr)
    return 4
  ips = hostgroup
  if ips is None:
    ips = []
    hostgroup = ips
  if ip not in ips:
    ips.append(ip)
  else:
    print('Already exists', file=sys.stderr)
    return 3
  writeYaml(hostgroupsFilename, content)
  return 0

def includehost(options, args):
  if len(args) != 2:
    print('Missing host group name or ip argument', file=sys.stderr)
    showUsage(options, args)
  result = addIp(args[0], args[1])
  code = result
  if code == 0:
    code = checkApplyOption(options)
  return code

def apply(options, args):
  proc = subprocess.run(['salt-call', 'state.apply', 'firewall', 'queue=True'])
  return proc.returncode

def main():
  options = []
  args = sys.argv[1:]
  for option in args:
    if option.startswith("--"):
      options.append(option)
      args.remove(option)

  if len(args) == 0:
    showUsage(options, None)

  commands = {
    "help": showUsage,
    "includehost": includehost,
    "apply": apply
  }

  code=1

  try:
    lockAttempts = 0
    maxAttempts = 30
    while lockAttempts < maxAttempts:
      lockAttempts = lockAttempts + 1
      try:
        f = open(lockFile, "x")
        f.close()
        break
      except:
        time.sleep(2)

    if lockAttempts == maxAttempts:
      print("Lock file (" + lockFile + ") could not be created; proceeding without lock.")

    cmd = commands.get(args[0], showUsage)
    code = cmd(options, args[1:])
  finally:
    try:
      os.remove(lockFile)
    except:
      print("Lock file (" + lockFile + ") already removed")

  sys.exit(code)

if __name__ == "__main__":
  main()