1- ###################################################################################################################### 
2- #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.                                           # 
3- #                                                                                                                    # 
4- #  Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance    # 
5- #  with the License. A copy of the License is located at                                                             # 
6- #                                                                                                                    # 
7- #      http://www.apache.org/licenses/LICENSE-2.0                                                                    # 
8- #                                                                                                                    # 
9- #  or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES # 
10- #  OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions    # 
11- #  and limitations under the License.                                                                                # 
12- ###################################################################################################################### 
1+ #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 
2+ #  SPDX-License-Identifier: Apache-2.0 
133
144import  os 
15- from  ipaddress  import  ip_address 
165from  ipaddress  import  IPv4Network 
176from  ipaddress  import  IPv6Network 
7+ from  ipaddress  import  ip_address 
188from  os  import  environ 
19- from  lib .waflibv2  import  WAFLIBv2 
20- from  lib .solution_metrics  import  send_metrics 
9+ 
10+ from  aws_lambda_powertools  import  Logger 
11+ 
2112from  lib .cw_metrics_util  import  WAFCloudWatchMetrics 
22- from  lib .logging_util  import  set_log_level 
13+ from  lib .solution_metrics  import  send_metrics 
14+ from  lib .waflibv2  import  WAFLIBv2 
15+ 
16+ logger  =  Logger (
17+     level = os .getenv ('LOG_LEVEL' )
18+ )
2319
2420waflib  =  WAFLIBv2 ()
2521CW_METRIC_PERIOD_SECONDS  =  12  *  3600     # Twelve hours in seconds 
@@ -38,13 +34,13 @@ def initialize_usage_data():
3834    return  usage_data 
3935
4036
41- def  get_bad_bot_usage_data (log ,  scope , cw , ipset_name_v4 , ipset_arn_v4 , ipset_name_v6 , ipset_arn_v6 , usage_data ):
42-     log .info ("[get_bad_bot_usage_data] Get bad bot data" )
37+ def  get_bad_bot_usage_data (scope , cw , ipset_name_v4 , ipset_arn_v4 , ipset_name_v6 , ipset_arn_v6 , usage_data ):
38+     logger .info ("[get_bad_bot_usage_data] Get bad bot data" )
4339
4440    if  'IP_SET_ID_BAD_BOTV4'  in  environ  or  'IP_SET_ID_BAD_BOTV6'  in  environ :
4541        # Get the count of ipv4 and ipv6 in bad bot ip sets 
46-         ipv4_count  =  waflib .get_ip_address_count (log , scope , ipset_name_v4 , ipset_arn_v4 )
47-         ipv6_count  =  waflib .get_ip_address_count (log , scope , ipset_name_v6 , ipset_arn_v6 )
42+         ipv4_count  =  waflib .get_ip_address_count (logger , scope , ipset_name_v4 , ipset_arn_v4 )
43+         ipv6_count  =  waflib .get_ip_address_count (logger , scope , ipset_name_v6 , ipset_arn_v6 )
4844        usage_data ['bad_bot_ip_set_size' ] =  str (ipv4_count  +  ipv6_count )
4945
5046        # Get the count of blocked requests for the bad bot rule from cloudwatch metrics 
@@ -59,14 +55,14 @@ def get_bad_bot_usage_data(log, scope, cw, ipset_name_v4, ipset_arn_v4, ipset_na
5955    return  usage_data 
6056
6157
62- def  send_anonymized_usage_data (log ,  scope , ipset_name_v4 , ipset_arn_v4 , ipset_name_v6 , ipset_arn_v6 ):
58+ def  send_anonymized_usage_data (scope , ipset_name_v4 , ipset_arn_v4 , ipset_name_v6 , ipset_arn_v6 ):
6359    try :
6460        if  'SEND_ANONYMIZED_USAGE_DATA'  not  in   environ  or  os .getenv ('SEND_ANONYMIZED_USAGE_DATA' ).lower () !=  'yes' :
6561            return 
6662
67-         log .info ("[send_anonymized_usage_data] Start" )
63+         logger .info ("[send_anonymized_usage_data] Start" )
6864
69-         cw  =  WAFCloudWatchMetrics (log )
65+         cw  =  WAFCloudWatchMetrics (logger )
7066        usage_data  =  initialize_usage_data ()
7167
7268        # Get the count of allowed requests for all the waf rules from cloudwatch metrics 
@@ -90,22 +86,22 @@ def send_anonymized_usage_data(log, scope, ipset_name_v4, ipset_arn_v4, ipset_na
9086        )
9187
9288        # Get bad bot specific usage data 
93-         usage_data  =  get_bad_bot_usage_data (log ,  scope , cw , ipset_name_v4 , ipset_arn_v4 ,
94-             ipset_name_v6 ,  ipset_arn_v6 ,  usage_data )
89+         usage_data  =  get_bad_bot_usage_data (scope , cw , ipset_name_v4 , ipset_arn_v4 ,  ipset_name_v6 ,  ipset_arn_v6 ,
90+                                              usage_data )
9591
9692        # Send usage data 
97-         log .info ('[send_anonymized_usage_data] Send usage data: \n {}' .format (usage_data ))
93+         logger .info ('[send_anonymized_usage_data] Send usage data: \n {}' .format (usage_data ))
9894        response  =  send_metrics (data = usage_data )
9995        response_code  =  response .status_code 
100-         log .info ('[send_anonymized_usage_data] Response Code: {}' .format (response_code ))
101-         log .info ("[send_anonymized_usage_data] End" )
96+         logger .info ('[send_anonymized_usage_data] Response Code: {}' .format (response_code ))
97+         logger .info ("[send_anonymized_usage_data] End" )
10298
10399    except  Exception  as  error :
104-         log .info ("[send_anonymized_usage_data] Failed to Send Data" )
105-         log .error (str (error ))
100+         logger .info ("[send_anonymized_usage_data] Failed to Send Data" )
101+         logger .error (str (error ))
106102
107103
108- def  add_ip_to_ip_set (log ,  scope , ip_type , source_ip , ipset_name , ipset_arn ):
104+ def  add_ip_to_ip_set (scope , ip_type , source_ip , ipset_name , ipset_arn ):
109105    new_address  =  []
110106    output  =  None 
111107
@@ -114,24 +110,24 @@ def add_ip_to_ip_set(log, scope, ip_type, source_ip, ipset_name, ipset_arn):
114110    elif  ip_type  ==  "IPV6" :
115111        new_address .append (IPv6Network (source_ip ).with_prefixlen )
116112
117-     ipset  =  waflib .get_ip_set (log , scope , ipset_name , ipset_arn )
113+     ipset  =  waflib .get_ip_set (logger , scope , ipset_name , ipset_arn )
118114    # merge old addresses with this one 
119-     log .info (ipset )
115+     logger .info (ipset )
120116    current_list  =  ipset ["IPSet" ]["Addresses" ]
121-     log .info (current_list )
117+     logger .info (current_list )
122118    new_list  =  list (set (current_list ) |  set (new_address ))
123-     log .info (new_list )
124-     output  =  waflib .update_ip_set (log , scope , ipset_name , ipset_arn , new_list )
119+     logger .info (new_list )
120+     output  =  waflib .update_ip_set (logger , scope , ipset_name , ipset_arn , new_list )
125121
126122    return  output 
127123
128124
129125# ====================================================================================================================== 
130126# Lambda Entry Point 
131127# ====================================================================================================================== 
128+ @logger .inject_lambda_context  
132129def  lambda_handler (event , _ ):
133-     log  =  set_log_level ()
134-     log .info ('[lambda_handler] Start' )
130+     logger .info ('[lambda_handler] Start' )
135131
136132    # ---------------------------------------------------------- 
137133    # Read inputs parameters 
@@ -144,30 +140,30 @@ def lambda_handler(event, _):
144140        ipset_arn_v6  =  os .getenv ('IP_SET_ID_BAD_BOTV6' )
145141
146142        # Fixed as old line had security exposure based on user supplied IP address 
147-         log .info ("Event->%s<-" , str (event ))
143+         logger .info ("Event->%s<-" , str (event ))
148144        if  event ['requestContext' ]['identity' ]['userAgent' ] ==  'Amazon CloudFront' :
149145            source_ip  =  str (event ['headers' ]['X-Forwarded-For' ].split (',' )[0 ].strip ())
150146        else :
151147            source_ip  =  str (event ['requestContext' ]['identity' ]['sourceIp' ])
152148
153-         log .info ("scope = %s" , scope )
154-         log .info ("ipset_name_v4 = %s" , ipset_name_v4 )
155-         log .info ("ipset_name_v6 = %s" , ipset_name_v6 )
156-         log .info ("IPARNV4 = %s" , ipset_arn_v4 )
157-         log .info ("IPARNV6 = %s" , ipset_arn_v6 )
158-         log .info ("source_ip = %s" , source_ip )
149+         logger .info ("scope = %s" , scope )
150+         logger .info ("ipset_name_v4 = %s" , ipset_name_v4 )
151+         logger .info ("ipset_name_v6 = %s" , ipset_name_v6 )
152+         logger .info ("IPARNV4 = %s" , ipset_arn_v4 )
153+         logger .info ("IPARNV6 = %s" , ipset_arn_v6 )
154+         logger .info ("source_ip = %s" , source_ip )
159155
160156        ip_type  =  "IPV%s"  %  ip_address (source_ip ).version 
161157        output  =  None 
162158        if  ip_type  ==  "IPV4" :
163-             output  =  add_ip_to_ip_set (log ,  scope , ip_type , source_ip , ipset_name_v4 , ipset_arn_v4 )
159+             output  =  add_ip_to_ip_set (scope , ip_type , source_ip , ipset_name_v4 , ipset_arn_v4 )
164160        elif  ip_type  ==  "IPV6" :
165-             output  =  add_ip_to_ip_set (log ,  scope , ip_type , source_ip , ipset_name_v6 , ipset_arn_v6 )
161+             output  =  add_ip_to_ip_set (scope , ip_type , source_ip , ipset_name_v6 , ipset_arn_v6 )
166162    except  Exception  as  e :
167-         log .error (e )
163+         logger .error (e )
168164        raise 
169165    finally :
170-         log .info ("Output->%s<-" , output )
166+         logger .info ("Output->%s<-" , output )
171167        message  =  "message: [%s] Thanks for the visit."  %  source_ip 
172168        response  =  {
173169            'statusCode' : 200 ,
@@ -176,7 +172,7 @@ def lambda_handler(event, _):
176172        }
177173
178174    if  output  is  not   None :
179-         send_anonymized_usage_data (log ,  scope , ipset_name_v4 , ipset_arn_v4 , ipset_name_v6 , ipset_arn_v6 )
180-     log .info ('[lambda_handler] End' )
175+         send_anonymized_usage_data (scope , ipset_name_v4 , ipset_arn_v4 , ipset_name_v6 , ipset_arn_v6 )
176+     logger .info ('[lambda_handler] End' )
181177
182178    return  response 
0 commit comments