#!/usr/bin/env python3

# Copyright 2014-2022 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 sys, argparse, re, docker
from packaging.version import Version, InvalidVersion
from itertools import groupby, chain


def get_image_name(string) -> str:
  return ':'.join(string.split(':')[:-1])


def get_so_image_basename(string) -> str:
  return get_image_name(string).split('/so-')[-1]


def get_image_version(string) -> str:
  ver = string.split(':')[-1]
  if ver == 'latest':
    # Version doesn't like "latest", so use a high semver
    return '99999.9.9'
  else:
    try:
      Version(ver)
    except InvalidVersion:
      # Also return a very high semver for any version 
      # with a dash in it since it will likely be a dev version of some kind
      if '-' in ver:
        return '999999.9.9'
    return ver


def main(quiet):
  client = docker.from_env()

  # Prune old/stopped containers
  if not quiet: print('Pruning old containers')
  client.containers.prune()

  image_list = client.images.list(filters={ 'dangling': False })

  # Map list of image objects to flattened list of tags (format: "name:version")
  tag_list = list(chain.from_iterable(list(map(lambda x: x.attrs.get('RepoTags'), image_list))))

  # Filter to only SO images (base name begins with "so-")
  tag_list = list(filter(lambda x: re.match(r'^.*\/so-[^\/]*$', get_image_name(x)), tag_list))

  # Group tags into lists by base name (sort by same projection first)
  tag_list.sort(key=lambda x: get_so_image_basename(x))
  grouped_tag_lists = [ list(it) for _, it in groupby(tag_list, lambda x: get_so_image_basename(x)) ]

  no_prunable = True
  for t_list in grouped_tag_lists:
    try:
      # Group tags by version, in case multiple images exist with the same version string
      t_list.sort(key=lambda x: Version(get_image_version(x)), reverse=True)
      grouped_t_list = [ list(it) for _,it in groupby(t_list, lambda x: get_image_version(x)) ]

      # Keep the 2 most current version groups
      if len(grouped_t_list) <= 2:
        continue
      else:
        no_prunable = False
        for group in grouped_t_list[2:]:
          for tag in group:
            if not quiet: print(f'Removing image {tag}')
            try:
              client.images.remove(tag, force=True)
            except docker.errors.ClientError as e:
              print(f'Could not remove image {tag}, continuing...')
    except (docker.errors.APIError, InvalidVersion) as e:
      print(f'so-{get_so_image_basename(t_list[0])}: {e}', file=sys.stderr)
      exit(1)
    except Exception as e:
      print('Unhandled exception occurred:')
      print(f'so-{get_so_image_basename(t_list[0])}: {e}', file=sys.stderr)
      exit(1)

  if no_prunable and not quiet:
    print('No Security Onion images to prune')


if __name__ == "__main__":
  main_parser = argparse.ArgumentParser(add_help=False)
  main_parser.add_argument('-q', '--quiet', action='store_const', const=True, required=False)
  args = main_parser.parse_args(sys.argv[1:])
 
  main(args.quiet)