From a639671b1f833ca1fd2dce03a2126cf0c14a37b9 Mon Sep 17 00:00:00 2001 From: Parth P Panchal Date: Sun, 6 Oct 2019 07:56:36 +0530 Subject: [PATCH] Passes filter tag as a parameter --- havoc/app.py | 1 + havoc/core.py | 7 ++++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/havoc/app.py b/havoc/app.py index eaf41bf..3e01024 100644 --- a/havoc/app.py +++ b/havoc/app.py @@ -77,6 +77,7 @@ def run_daemon(havoc, log, run_every='5min'): @click.option('--logfile', default='/var/log/havoc.log', help="Logging file when the application is in daemon mode") @click.option('--dry-run', is_flag=True, help="If use, HAvOC will display the result and not change haproxy.cfg") @click.option('--debug', is_flag=True, help="Debug mode") +@click.option('--filter-tag', default="pool", help="Filter tag") def cli(**options): """HAvOC (HAproxy clOud Configuration) diff --git a/havoc/core.py b/havoc/core.py index 580fc29..e273449 100644 --- a/havoc/core.py +++ b/havoc/core.py @@ -28,7 +28,7 @@ def __init__(self, ec2, nova, options, log): self.options = options - def get_ec2_instances_in_pool(self, pool, suffix): + def get_ec2_instances_in_pool(self, pool, suffix, filter_tag): self.log.debug("AWS EC2 : Trying to find vm in %s", pool) instances = [] if self.ec2 is None: @@ -36,7 +36,7 @@ def get_ec2_instances_in_pool(self, pool, suffix): return instances - filters = {'tag:pool': pool,'instance-state-name': 'running'} + filters = {'tag:{}'.format(filter_tag): pool, 'instance-state-name': 'running'} if self.options['aws_vpc'] is not None: filters['tag:vpc'] = self.options['aws_vpc'] @@ -156,13 +156,14 @@ def reload_haproxy(self): return False def run(self): + filter_tag = self.options['filter_tag'] instances = {} # Retrieve instances for pool in self.options['pools'].split(','): self.log.debug("POOL : %s", pool) #TODO: Create CLI parameters for suffixes - ec2_instances = self.get_ec2_instances_in_pool(pool, "_aws") + ec2_instances = self.get_ec2_instances_in_pool(pool, "_aws", filter_tag) os_instances = self.get_os_instances_in_pool(pool, "_os") instances[pool] = ec2_instances + os_instances