def test_dynamo_upsert(self): vpc_id, acl_id = self.create_resources() ip_hash_map = {'244.244.244.244': 1} ec2_client = EC2Allocator().client('eu-west-1') dynamo_allocator = DynamoAllocator() dynamo_resource = dynamo_allocator.resource() table = dynamo_resource.Table(get_dynamo_table()) struct_time = datetime.strptime("21/11/06 16:30", "%d/%m/%y %H:%M") epoch = UTC_time_to_epoch(struct_time) table.put_item( Item={ 'ip': '244.244.244.244', 'last_seen': epoch, 'times_seen': 1, 'banned': True }) DynamoUpserter().upsert_ip_in_table(dynamo_resource, '244.244.244.244') response = table.query( KeyConditionExpression=Key('ip').eq('244.244.244.244')) print(response) assert len(response['Items']) == 1 assert response['Items'][0]['banned'] == False
def upsert_ip_in_table(dynamo_resource, ip): table = dynamo_resource.Table(get_dynamo_table()) times_seen = 1 last_seen= UTC_time_to_epoch(datetime.utcnow()) response = table.query( KeyConditionExpression=Key('ip').eq(ip) ) banned = False if len(response['Items']) == 1: times_seen = response['Items'][0]['times_seen'] + 1 banned = bool(response['Items'][0]['banned']) table.put_item( Item={ 'ip': ip, 'last_seen': last_seen, 'times_seen': times_seen, 'banned': banned } ) logging.getLogger().info('Scheduled for banning ip {}'.format(ip))
def test_ipbanner(self): vpc_id, acl_id = self.create_resources() datadog_client = DatadogClient() ip_hash_map = {'244.244.244.244': 1} ec2_client = EC2Allocator().client('eu-west-1') dynamo_allocator = DynamoAllocator() dynamo_resource = dynamo_allocator.resource() table = dynamo_resource.Table(get_dynamo_table()) struct_time = datetime.strptime("21/11/06 16:30", "%d/%m/%y %H:%M") epoch = UTC_time_to_epoch(struct_time) table.put_item( Item={ 'ip': '244.244.244.244', 'last_seen': epoch, 'times_seen': 1, 'banned': False }) ip_banner = IpBanner(ec2_client, datadog_client) ip_banner.ban_ip('244.244.244.244', acl_id, 90) response = table.query( KeyConditionExpression=Key('ip').eq('244.244.244.244')) assert response['Items'][0]['banned'] == True
def test_is_recent(self): struct_time = datetime.strptime("21/11/06 16:30", "%d/%m/%y %H:%M") epoch = UTC_time_to_epoch(struct_time) event = { "eventID": "1", "eventVersion": "1.0", "dynamodb": { "Keys": { "ip": { "S": "244.244.244.244" }, "times_seen": { "S": "244.244.244.244" }, "last_seen": { "N": "{}".format(epoch) }, "banned": { "BOOL": False } }, "NewImage": { "ip": { "S": "244.244.244.244" }, "times_seen": { "S": "244.244.244.244" }, "last_seen": { "N": "{}".format(epoch) }, "banned": { "BOOL": False } }, "StreamViewType": "NEW_AND_OLD_IMAGES", "SequenceNumber": "111", "SizeBytes": 26 }, "awsRegion": "eu-west-1", "eventName": "INSERT", "eventSourceARN": 'testarn', "eventSource": "aws:dynamodb" } dynamo_db_event = DynamoDBEvent(event) print(TimeDecision().is_recent(dynamo_db_event)) assert False == TimeDecision().is_recent(dynamo_db_event)
def create_resources(self): # Create Dynamo and user dynamo_allocator = DynamoAllocator() dynamo_resource = dynamo_allocator.resource() self.create_table() dynamo_table = dynamo_allocator.table('Autobahn') ec2_client = EC2Allocator().client('eu-west-1') vpc_response = ec2_client.create_vpc(CidrBlock='10.1.0.0/24', AmazonProvidedIpv6CidrBlock=False, DryRun=False, InstanceTenancy='default') self.vpc_id = vpc_response['Vpc']['VpcId'] ec2_client.create_tags(DryRun=False, Resources=[self.vpc_id], Tags=[{ 'Key': 'Name', 'Value': "vpc-test" }]) acl_response = ec2_client.create_network_acl(DryRun=False, VpcId=self.vpc_id) self.acl_id = acl_response['NetworkAcl']['NetworkAclId'] entry_response = ec2_client.create_network_acl_entry( CidrBlock='244.244.244.244/32', DryRun=False, Egress=False, NetworkAclId=self.acl_id, Protocol='-1', RuleAction='deny', RuleNumber=77) struct_time = datetime.strptime("21/11/06 16:30", "%d/%m/%y %H:%M") epoch = UTC_time_to_epoch(struct_time) table = dynamo_resource.Table('Autobahn') table.put_item( Item={ 'ip': '244.244.244.244', 'last_seen': epoch - 99100, 'times_seen': 1, 'banned': True }) return self.vpc_id, self.acl_id
def upsert_ip_in_table(dynamo_resource, ip): table = dynamo_resource.Table(get_dynamo_table()) times_seen = 0 response = table.query(KeyConditionExpression=Key('ip').eq(ip)) if len(response['Items']) == 1: times_seen = response['Items'][0]['times_seen'] table.put_item( Item={ 'ip': ip, 'last_seen': UTC_time_to_epoch(datetime.utcnow()), 'times_seen': times_seen + 1, 'banned': True }) logging.getLogger().info('Banned ip {}'.format(ip))
def test_get_unbannables(self): vpc_id, acl_id = self.create_resources() ip_hash_map = {'244.244.244.244': 1} ec2_client = EC2Allocator().client('eu-west-1') dynamo_allocator = DynamoAllocator() dynamo_resource = dynamo_allocator.resource() dynamo_client = dynamo_allocator.client() table = dynamo_resource.Table(get_dynamo_table()) struct_time = datetime.strptime("21/11/06 16:30", "%d/%m/%y %H:%M") epoch = UTC_time_to_epoch(struct_time) table.put_item( Item={ 'ip': '244.244.244.244', 'last_seen': epoch, 'times_seen': 1, 'banned': True }) unbannables = DynamoGetUnbannable().get_unbannables( dynamo_client, dynamo_resource) assert '244.244.244.244' in unbannables
def test_time_decision(self): struct_time = datetime.strptime("21/11/06 16:30", "%d/%m/%y %H:%M") epoch = UTC_time_to_epoch(struct_time) assert True == TimeDecision().should_unban(epoch)
def test_utc_convert(self): struct_time = datetime.strptime("21/11/06 16:30", "%d/%m/%y %H:%M") epoch = UTC_time_to_epoch(struct_time) assert epoch == 1164126600
def test_event_processer(self): print(__name__) vpc_id, acl_id = self.create_resources() struct_time = datetime.utcnow() epoch = UTC_time_to_epoch(struct_time) record = {'country': 'dev', 'project': 'test'} ip_hash_map = {'244.244.244.244': 1} event = { "eventID": "1", "eventVersion": "1.0", "dynamodb": { "Keys": { "ip": { "S": "244.244.244.244" }, "times_seen": { "S": "244.244.244.244" }, "last_seen": { "N": "{}".format(epoch) }, "banned": { "BOOL": False } }, "NewImage": { "ip": { "S": "244.244.244.244" }, "times_seen": { "S": "244.244.244.244" }, "last_seen": { "N": "{}".format(epoch) }, "banned": { "BOOL": False } }, "StreamViewType": "NEW_AND_OLD_IMAGES", "SequenceNumber": "111", "SizeBytes": 26 }, "awsRegion": "eu-west-1", "eventName": "INSERT", "eventSourceARN": 'testarn', "eventSource": "aws:dynamodb" } ec2_client = EC2Allocator().client('eu-west-1') dynamo_allocator = DynamoAllocator() dynamo_resource = dynamo_allocator.resource() table = dynamo_resource.Table(get_dynamo_table()) table.put_item( Item={ 'ip': '244.244.244.244', 'last_seen': epoch, 'times_seen': 1, 'banned': False }) EventProcesser().process(event) response = table.query( KeyConditionExpression=Key('ip').eq('244.244.244.244')) datadog_client = DatadogClient() ip_banner = IpBanner(ec2_client, datadog_client) assert response['Items'][0]['banned'] == True assert ip_banner.not_already_banned('244.244.244.244', vpc_id) == False