def test_zmon_set_name(monkeypatch): put = MagicMock() put.return_value = MagicMock() monkeypatch.setattr('requests.Session.put', put) zmon = Zmon(URL, token=TOKEN) zmon.add_phone('user1@something', 'user1') put.assert_called_with(zmon.endpoint(client.GROUPS, 'user1@something', client.PHONE, 'user1'))
def test_zmon_delete_check_definition(monkeypatch, result): delete = MagicMock() delete.return_value.ok = result monkeypatch.setattr('requests.Session.delete', delete) zmon = Zmon(URL, token=TOKEN) res = zmon.delete_check_definition(1) assert res.ok is result delete.assert_called_with(zmon.endpoint(client.CHECK_DEF, 1))
def test_zmon_remove_member(monkeypatch): delete = MagicMock() delete.return_value.text = '1' monkeypatch.setattr('requests.Session.delete', delete) zmon = Zmon(URL, token=TOKEN) deleted = zmon.remove_member('group', 'user1') assert deleted is True delete.assert_called_with(zmon.endpoint(client.GROUPS, 'group', client.MEMBER, 'user1'))
def test_zmon_add_member(monkeypatch): put = MagicMock() put.return_value.text = '1' monkeypatch.setattr('requests.Session.put', put) zmon = Zmon(URL, token=TOKEN) added = zmon.add_member('group', 'user1') assert added is True put.assert_called_with(zmon.endpoint(client.GROUPS, 'group', client.MEMBER, 'user1'))
def test_zmon_remove_phone(monkeypatch): delete = MagicMock() delete.return_value.text = '1' monkeypatch.setattr('requests.Session.delete', delete) zmon = Zmon(URL, token=TOKEN) deleted = zmon.remove_phone('user1@something', '12345') assert deleted is True delete.assert_called_with(zmon.endpoint(client.GROUPS, 'user1@something', client.PHONE, '12345'))
def test_zmon_delete_entity(monkeypatch, result): delete = MagicMock() delete.return_value.text = result monkeypatch.setattr('requests.Session.delete', delete) zmon = Zmon(URL, token=TOKEN) deleted = zmon.delete_entity(1) assert deleted is (result == '1') delete.assert_called_with(zmon.endpoint(client.ENTITIES, 1))
def test_zmon_get_alert_defintions(monkeypatch, resp, result): get = MagicMock() get.return_value.json.return_value = resp monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) res = zmon.get_alert_definitions() assert res == result get.assert_called_with(zmon.endpoint(client.ACTIVE_ALERT_DEF))
def test_zmon_add_phone(monkeypatch): put = MagicMock() put.return_value.text = '1' monkeypatch.setattr('requests.Session.put', put) zmon = Zmon(URL, token=TOKEN) added = zmon.add_phone('user1@something', '12345') assert added is True put.assert_called_with(zmon.endpoint(client.GROUPS, 'user1@something', client.PHONE, '12345'))
def test_zmon_get_alert_defintion(monkeypatch): get = MagicMock() result = {'id': 1, 'type': 'dummy'} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) check = zmon.get_alert_definition(1) assert check == result get.assert_called_with(zmon.endpoint(client.ALERT_DEF, 1))
def test_zmon_get_groups(monkeypatch): get = MagicMock() result = [1, 2, 3] get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) check = zmon.get_groups() assert check == result get.assert_called_with(zmon.endpoint(client.GROUPS))
def test_zmon_get_grafana_dashboard(monkeypatch): get = MagicMock() result = {'dashboard': {}} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) check = zmon.get_grafana_dashboard(1) assert check == result get.assert_called_with(zmon.endpoint(client.GRAFANA, 1))
def test_zmon_status(monkeypatch): get = MagicMock() result = {'status': 'success'} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) status = zmon.status() assert status == result get.assert_called_with(zmon.endpoint(client.STATUS))
def test_zmon_alert_data(monkeypatch): get = MagicMock() result = {'entity-1': []} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) check = zmon.get_alert_data(1) assert check == result get.assert_called_with(zmon.endpoint(client.ALERT_DATA, 1, 'all-entities'))
def test_zmon_update_dashboard(monkeypatch, d): post = MagicMock() result = 1 post.return_value.json.return_value = result monkeypatch.setattr('requests.Session.post', post) zmon = Zmon(URL, token=TOKEN) res = zmon.update_dashboard(d) assert res == result url = zmon.endpoint(client.DASHBOARD, 1) if d['id'] else zmon.endpoint(client.DASHBOARD) post.assert_called_with(url, json=d)
def test_zmon_get_dashboard(monkeypatch): get = MagicMock() result = {'id': 1, 'type': 'dummy'} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) res = zmon.get_dashboard(1) assert res == result get.assert_called_with(zmon.endpoint(client.DASHBOARD, 1))
def test_zmon_get_entity(monkeypatch): get = MagicMock() result = {'id': 1, 'type': 'dummy'} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) res = zmon.get_entity(1) assert res == result get.assert_called_with(zmon.endpoint(client.ENTITIES, 1, trailing_slash=False))
def test_zmon_get_entities(monkeypatch, q, result): get = MagicMock() get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) res = zmon.get_entities(query=q) assert res == result params = {'query': json.dumps(q)} if q else None get.assert_called_with(zmon.endpoint(client.ENTITIES), params=params)
def test_zmon_delete_alert_definition(monkeypatch): delete = MagicMock() result = {'status': 'success'} delete.return_value.json.return_value = result monkeypatch.setattr('requests.Session.delete', delete) zmon = Zmon(URL, token=TOKEN) res = zmon.delete_alert_definition(1) assert res == result delete.assert_called_with(zmon.endpoint(client.ALERT_DEF, 1))
def test_zmon_get_token(monkeypatch): post = MagicMock() result = '1111' post.return_value.text = result monkeypatch.setattr('requests.Session.post', post) zmon = Zmon(URL, token=TOKEN) check = zmon.get_onetime_token() assert check == result post.assert_called_with(zmon.endpoint(client.TOKENS), json={})
def test_zmon_search(monkeypatch): get = MagicMock() result = {'alerts': []} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) q = 'health check' search = zmon.search(q) assert search == result get.assert_called_with(zmon.endpoint(client.SEARCH), params={'query': q})
def test_zmon_headers(monkeypatch): zmon = Zmon(URL, token=TOKEN) assert zmon.session.headers['Authorization'] == 'Bearer {}'.format(TOKEN) zmon = Zmon(URL, username='******', password='******') assert zmon.session.auth == ('user1', 'password') zmon = Zmon(URL, token=TOKEN) assert zmon.session.headers['User-Agent'] == client.ZMON_USER_AGENT zmon = Zmon(URL, token=TOKEN, user_agent='zmon-client-wrapper/0.5') assert zmon.session.headers['User-Agent'] == 'zmon-client-wrapper/0.5' zmon = Zmon(URL, token=TOKEN, verify=False) assert zmon.session.verify is False
def test_zmon_list_tokens(monkeypatch): get = MagicMock() result = [1, 2, 3] get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) check = zmon.list_onetime_tokens() assert check == result get.assert_called_with(zmon.endpoint(client.TOKENS), timeout=DEFAULT_TIMEOUT)
def get_client(config): verify = config.get('verify', True) if 'user' in config and 'password' in config: return Zmon(config['url'], username=config['user'], password=config['password'], verify=verify) elif os.environ.get('ZMON_TOKEN'): return Zmon(config['url'], token=os.environ.get('ZMON_TOKEN'), verify=verify) elif 'token' in config: return Zmon(config['url'], token=config['token'], verify=verify) raise RuntimeError( 'Failed to intitialize ZMON client. Invalid configuration!')
def test_zmon_view_urls(monkeypatch): zmon = Zmon(URL, token=TOKEN) # Checks check = {'id': 1} assert '{}#/check-definitions/view/1/'.format( URL) == zmon.check_definition_url(check) # Alerts alert = {'id': 1} assert '{}#/alert-details/1/'.format(URL) == zmon.alert_details_url(alert) # Dashboard dashboard_id = 1 assert '{}#/dashboards/views/1/'.format(URL) == zmon.dashboard_url( dashboard_id) # Token token = '1234' assert '{}/tv/1234/'.format(URL) == zmon.token_login_url(token) # Grafana dashboard = {'id': 'grafana-dash'} assert '{}/visualization/dashboard/grafana-dash'.format( URL) == zmon.grafana_dashboard_url(dashboard)
def test_zmon_switch_active_user(monkeypatch, success): del_success, put_success = success delete = MagicMock() delete.return_value.ok = del_success if not del_success: delete.return_value.raise_for_status.side_effect = RuntimeError put = MagicMock() put.return_value.ok = put_success put.return_value.text = '1' if not put_success: put.return_value.raise_for_status.side_effect = RuntimeError monkeypatch.setattr('requests.Session.delete', delete) monkeypatch.setattr('requests.Session.put', put) zmon = Zmon(URL, token=TOKEN) if not del_success or not put_success: with pytest.raises(RuntimeError): zmon.switch_active_user('g', 'u') else: switched = zmon.switch_active_user('g', 'u') assert switched is True delete.assert_called_with(zmon.endpoint(client.GROUPS, 'g', 'active')) if del_success: put.assert_called_with(zmon.endpoint(client.GROUPS, 'g', 'active', 'u'), timeout=DEFAULT_TIMEOUT)
def get_clients(zmon_url, verify=True) -> Zmon: """Return Pykube and Zmon client instances as a tuple.""" # Get token if set as ENV variable. This is useful in development. zmon_token = os.getenv('ZMON_AGENT_TOKEN') if not zmon_token: zmon_token = tokens.get('uid') return Zmon(zmon_url, token=zmon_token, verify=verify)
def test_zmon_search_team(monkeypatch): get = MagicMock() result = {'alerts': []} get.return_value.json.return_value = result monkeypatch.setattr('requests.Session.get', get) zmon = Zmon(URL, token=TOKEN) q = 'health check' teams = ['team-1', 'team-2'] search = zmon.search(q, teams=teams) assert search == result get.assert_called_with(zmon.endpoint(client.SEARCH), params={ 'query': q, 'teams': 'team-1,team-2' }, timeout=DEFAULT_TIMEOUT)
def validate_sli_source(config, source, ignore_keys=False): if 'zmon_url' not in config: config = set_config_file() zmon = Zmon(config['zmon_url'], token=zign.api.get_token('uid', ['uid'])) check_id = int(source['check_id']) try: check = zmon.get_check_definition(check_id) except Exception: raise SLRClientError( 'Check definition {} does not seem to exist!'.format(check_id)) alerts = zmon.get_alert_definitions() filtered = [ alert for alert in alerts if alert['check_definition_id'] == check_id ] if not filtered: raise SLRClientError( 'Check definition has no alerts. Please create an Alert for this check on ZMON {}' .format(zmon.check_definition_url(check))) if ignore_keys: return keys = [k for k in source['keys'] if '.*' not in k] if not keys: # Do not validate keys if we have wildcards return sli_exists = False sample_data = set() for alert in filtered: if sli_exists: break alert_data = zmon.get_alert_data(alert['id']) values = { v['entity']: v['results'][0]['value'] for v in alert_data if len(v['results']) } for entity, data in values.items(): if type(data) is dict: flattened = flatten(data) data_keys = flattened.keys() sample_data.update(list(data_keys)) if not (set(keys) - set(data_keys)): sli_exists = True break if not sli_exists: raise SLRClientError( 'Some SLI keys do not exist. Please check the data returned from the check and the corresponding keys ' 'in the SLI source. Found the following keys returned from check {}: {}' .format(check_id, set(sample_data)))
def test_zmon_add_entity(monkeypatch, e, result): fail = True if type(result) is dict: fail = False put = MagicMock() resp = MagicMock() resp.ok = True put.return_value = resp monkeypatch.setattr('requests.Session.put', put) zmon = Zmon(URL, token=TOKEN) if fail: with pytest.raises(result): zmon.add_entity(e) else: r = zmon.add_entity(e) assert r.ok is True put.assert_called_with(zmon.endpoint(client.ENTITIES, trailing_slash=False), data=json.dumps(result), timeout=DEFAULT_TIMEOUT)
def validate_sli(config, data_source): if 'zmon_url' not in config: config = set_config_file() zmon = Zmon(config['zmon_url'], token=zign.api.get_token('uid', ['uid'])) check_id = int(data_source['definition']['check_id']) try: zmon.get_check_definition(check_id) except: raise SLRClientError( 'Check definition {} does not seem to exist!'.format(check_id)) alerts = zmon.get_alert_definitions() filtered = [ alert for alert in alerts if alert['check_definition_id'] == check_id ] if not filtered: raise SLRClientError( 'Check definition has no alerts. Please create an Alert for this check on ZMON.' ) keys = data_source['definition']['keys'] sli_exists = False sample_data = {} for alert in filtered: if sli_exists: break alert_data = zmon.get_alert_data(alert['id']) values = { v['entity']: v['results'][0]['value'] for v in alert_data if len(v['results']) } for entity, data in values.items(): if type(data) is dict: data_keys = data.keys() if data_keys: sample_data = data_keys if not (set(keys) - set(data_keys)): sli_exists = True break if not sli_exists: raise SLRClientError( 'Some SLI keys do not exist. Please check the data returned from the check and the corresponding keys ' 'in the data-source. Found the following keys returned from check {}: {}' .format(check_id, set(sample_data)))
def main(): argp = argparse.ArgumentParser(description='ZMON AWS Agent') argp.add_argument('-e', '--entity-service', dest='entityservice') argp.add_argument('-r', '--region', dest='region', default=None) argp.add_argument('-j', '--json', dest='json', action='store_true') argp.add_argument('-t', '--tracer', dest='tracer', default=os.environ.get('OPENTRACING_TRACER', 'noop')) argp.add_argument('--no-oauth2', dest='disable_oauth2', action='store_true', default=False) argp.add_argument('--postgresql-user', dest='postgresql_user', default=os.environ.get('AGENT_POSTGRESQL_USER')) argp.add_argument('--postgresql-pass', dest='postgresql_pass', default=os.environ.get('AGENT_POSTGRESQL_PASS')) args = argp.parse_args() if not args.disable_oauth2: tokens.configure() tokens.manage('uid', ['uid']) tokens.start() init_opentracing_tracer(args.tracer) root_span = opentracing.tracer.start_span(operation_name='aws_entity_discovery') with root_span: logging.basicConfig(level=logging.INFO) # 0. Fetch extra data for entities entity_extras = {} for ex in os.getenv('EXTRA_ENTITY_FIELDS', '').split(','): if '=' not in ex: continue k, v = ex.split('=', 1) if k and v: entity_extras[k] = v # 1. Determine region if not args.region: logger.info('Trying to figure out region..') try: response = requests.get('http://169.254.169.254/latest/meta-data/placement/availability-zone', timeout=2) except Exception: root_span.set_tag('error', True) root_span.log_kv({'exception': traceback.format_exc()}) logger.exception('Region was not specified as a parameter and' + 'can not be fetched from instance meta-data!') raise region = response.text[:-1] else: region = args.region root_span.set_tag('region', region) logger.info('Using region: {}'.format(region)) logger.info('Entity service URL: %s', args.entityservice) logger.info('Reading DNS data for hosted zones') aws.populate_dns_data() aws_account_id = aws.get_account_id(region) infrastructure_account = 'aws:{}'.format(aws_account_id) if aws_account_id else None if not infrastructure_account: logger.error('AWS agent: Cannot determine infrastructure account ID. Terminating!') return root_span.set_tag('account', infrastructure_account) # 2. ZMON entities if not args.disable_oauth2: token = os.getenv('ZMON_TOKEN', None) or tokens.get('uid') zmon_client = Zmon(args.entityservice, token=token, user_agent=get_user_agent()) query = {'infrastructure_account': infrastructure_account, 'region': region, 'created_by': 'agent'} entities = zmon_client.get_entities(query) # 3. Get running apps apps = aws.get_running_apps(region, entities) elbs = [] scaling_groups = [] elastigroups = [] certificates = [] rds = [] elasticaches = [] dynamodbs = [] sqs = [] postgresql_clusters = [] aws_limits = [] new_entities = [] to_be_removed = [] if len(apps) > 0: elbs = aws.get_running_elbs(region, infrastructure_account) scaling_groups = aws.get_auto_scaling_groups(region, infrastructure_account) elastigroups = elastigroup.get_elastigroup_entities(region, infrastructure_account) rds = aws.get_rds_instances(region, infrastructure_account, entities) elasticaches = aws.get_elasticache_nodes(region, infrastructure_account) dynamodbs = aws.get_dynamodb_tables(region, infrastructure_account) certificates = aws.get_certificates(region, infrastructure_account) aws_limits = aws.get_limits(region, infrastructure_account, apps, elbs, entities) sqs = aws.get_sqs_queues(region, infrastructure_account, entities) postgresql_clusters = postgresql.get_postgresql_clusters(region, infrastructure_account, scaling_groups, apps) account_alias = aws.get_account_alias(region) ia_entity = { 'type': 'local', 'infrastructure_account': infrastructure_account, 'account_alias': account_alias, 'region': region, 'id': 'aws-ac[{}:{}]'.format(infrastructure_account, region), 'created_by': 'agent', } account_alias_prefix = os.getenv('ACCOUNT_ALIAS_PREFIX', None) owner = account_alias if account_alias_prefix: owner = owner.replace(account_alias_prefix, '', 1) root_span.set_tag('team', owner) application_entities = aws.get_apps_from_entities(apps, infrastructure_account, region) if args.postgresql_user and args.postgresql_pass: postgresql_databases = postgresql.get_databases_from_clusters(postgresql_clusters, infrastructure_account, region, args.postgresql_user, args.postgresql_pass) else: # Pretend the list of DBs is empty, but also make sure we don't remove # any pre-existing database entities because we don't know about them. postgresql_databases = [] entities = [e for e in entities if e.get('type') != 'postgresql_database'] current_entities = ( elbs + scaling_groups + elastigroups + apps + application_entities + rds + postgresql_databases + postgresql_clusters + elasticaches + dynamodbs + certificates + sqs) current_entities.append(aws_limits) current_entities.append(ia_entity) for entity in current_entities: entity.update(entity_extras) # 4. Removing misssing entities existing_ids = get_existing_ids(entities) current_entities_ids = {e['id'] for e in current_entities} to_be_removed, delete_error_count = remove_missing_entities( existing_ids, current_entities_ids, zmon_client, json=args.json) root_span.log_kv({'total_entitites': str(len(current_entities))}) root_span.log_kv({'removed_entities': str(len(to_be_removed))}) logger.info('Found {} removed entities from {} entities ({} failed)'.format( len(to_be_removed), len(current_entities), delete_error_count)) # 5. Get new/updated entities new_entities, add_error_count = add_new_entities(current_entities, entities, zmon_client, json=args.json) root_span.log_kv({'new_entities': str(len(new_entities))}) logger.info('Found {} new entities from {} entities ({} failed)'.format( len(new_entities), len(current_entities), add_error_count)) # 6. Always add Local entity if not args.json: ia_entity['errors'] = {'delete_count': delete_error_count, 'add_count': add_error_count} update_local_entity(zmon_client, ia_entity) types = {e['type']: len([t for t in new_entities if t['type'] == e['type']]) for e in new_entities} for t, v in types.items(): logger.info('Found {} new entities of type: {}'.format(v, t)) # Check if it is a dry run! if args.json: d = { 'applications': application_entities, 'apps': apps, 'elastigroups': elastigroups, 'dynamodb': dynamodbs, 'elbs': elbs, 'elc': elasticaches, 'rds': rds, 'certificates': certificates, 'aws_limits': aws_limits, 'sqs_queues': sqs, 'new_entities': new_entities, 'to_be_removed': to_be_removed, 'posgresql_clusters': postgresql_clusters } print(json.dumps(d, indent=4))