Merge pull request #8 from muffat/update2

Update lambda function

.gitignore

@@ -4,4 +4,5 @@ terraform.tfstate
 terraform.tfstate.backup
 sgscanner.zip
 .terraform.tfstate.lock.info
-sg.config
+sg.config
+tes.py

README.md

@@ -22,6 +22,7 @@ module "sgscanner" {
     SLACK_URL      = "https://hooks.slack.com/.."
     SLACK_USERNAME = "test"
     SLACK_CHANNEL  = "slack-channel-name"
+    TAG_EXCEPTION  = "tag-key,tag-value"
   }
 
   # Find the non-compliant IP address and port
@@ -43,7 +44,10 @@ module "sgscanner" {
 | function_name | Lambda function name | `string` | `` | yes |
 | description | Some descriptions | `string` | `` | no |
 | s3_bucket | S3 bucket name to store the lambda object | `string` | `` | yes |
-| environment_variables | Environment variables list to include the SLACK details. `SLACK_URL`, `SLACK_USERNAME`, and `SLACK_CHANNEL` | `map(string)` | `null` | yes |
+| environment_variables(`SLACK_URL`) | Environment variable for Slack URL webhook | `string` | `null` | no |
+| environment_variables(`SLACK_USERNAME`) | Environment variable for Slack webhook username | `string` | `null` | no |
+| environment_variables(`SLACK_CHANNEL`) | Environment variable for Slack webhook channel name | `string` | `null` | no |
+| environment_variables(`TAG_EXCEPTION`) | Environment variable to whitelist the security group using specific tag. Format: `tag-key,tag-value` Example: `set-public,true` | `string` | `null` | no |
 | schedule_expression | Cloudwatch event custom cron expression | `string` | `cron(0 0 * * ? *)` | no |
 | role | A custom IAM role arn | `string` | `null` | no |
 | finder | A map of IP address and port to find | `map(string)` | `{}` | no |

example/main.tf

@@ -12,6 +12,7 @@ module "sgscanner" {
     SLACK_URL      = "https://hooks.slack.com/.."
     SLACK_USERNAME = "test"
     SLACK_CHANNEL  = "slack-channel-name"
+    TAG_EXCEPTION  = "tag-key,tag-value"
   }
 
   finder = {

src/sgscanner.py

@@ -2,13 +2,15 @@
 import boto3
 import os
 from urllib import request
+import sys
 
 
 def lambda_handler(event, context):
 
     SLACK_URL = os.environ['SLACK_URL']
     SLACK_CHANNEL = os.environ['SLACK_CHANNEL']
     SLACK_USERNAME = os.environ['SLACK_USERNAME']
+    TAG_EXCEPTION = os.environ['TAG_EXCEPTION']
 
     client = boto3.client('ec2')
     response = client.describe_security_groups()
@@ -19,20 +21,48 @@ def lambda_handler(event, context):
 
     slack_content = []
     for i in response['SecurityGroups']:
-        SecurityGroupId = i['GroupId']
+        try:
+            SecurityGroupTags = i['Tags']
+            TAG_EXCEPTION
+            is_tag = True
+        except (KeyError, NameError):
+            is_tag = False
+            pass
+        SecurityGroupId = i['GroupId']        
         for j in i['IpPermissions']:
-            try:
-                this_ip = j['IpRanges'][0]['CidrIp']
-                this_port = j['FromPort']
+            for ip in j['IpRanges']:
+                try:
+                    this_port = j['FromPort']
+                except KeyError:
+                    continue
                 for key in config_load:
+                    this_ip = ip['CidrIp']
                     if key == this_ip:
-                        if this_port in config_load[this_ip]:
+                        cig = config_load[this_ip]
+                        try:
+                            if len(cig) > 0:
+                                pass
+                        except TypeError:
+                            cig = [cig]
+                        if this_port in cig:
                             sg_finding = "+ Found `%s:%s` in `%s`" % (
                                 this_ip, this_port, SecurityGroupId)
                             if sg_finding not in slack_content:
-                                slack_content.append(sg_finding)
-            except:
-                pass
+                                if is_tag == True:
+                                    ignore = False
+                                    for tag in i['Tags']:
+                                        SecurityGroupTagKey = tag['Key']
+                                        SecurityGroupTagValue = tag['Value']
+                                        exception = TAG_EXCEPTION.split(",")
+                                        exception_key = exception[0]
+                                        exception_value = exception[1]
+                                        if SecurityGroupTagKey == exception_key:
+                                            if SecurityGroupTagValue == exception_value:
+                                                ignore = True
+                                else:
+                                    ignore = False
+                                if ignore == False:
+                                    slack_content.append(sg_finding)
 
     if len(slack_content) > 0:
         slack_content = "".join(slack_content)
@@ -48,4 +78,4 @@ def lambda_handler(event, context):
                               headers={'Content-Type': 'application/json'})
         request.urlopen(req)
 
-    return {'statusCode': 200}
+    return {'statusCode': 200}