1414logger = logging .getLogger (__name__ )
1515logger .setLevel (logging .DEBUG )
1616# Configure your cluster name and region here
17- KUBE_FILEPATH = ' /tmp/kubeconfig'
17+ KUBE_FILEPATH = " /tmp/kubeconfig"
1818MIRROR_POD_ANNOTATION_KEY = "kubernetes.io/config.mirror"
1919CONTROLLER_KIND_DAEMON_SET = "DaemonSet"
2020
21+
2122def create_kube_config (eks , cluster_name ):
2223 """Creates the Kubernetes config file required when instantiating the API client."""
23- cluster_info = eks .describe_cluster (name = cluster_name )[' cluster' ]
24- certificate = cluster_info [' certificateAuthority' ][ ' data' ]
25- endpoint = cluster_info [' endpoint' ]
24+ cluster_info = eks .describe_cluster (name = cluster_name )[" cluster" ]
25+ certificate = cluster_info [" certificateAuthority" ][ " data" ]
26+ endpoint = cluster_info [" endpoint" ]
2627
2728 kube_config = {
28- 'apiVersion' : 'v1' ,
29- 'clusters' : [
30- {
31- 'cluster' :
32- {
33- 'server' : endpoint ,
34- 'certificate-authority-data' : certificate
35- },
36- 'name' : 'k8s'
37-
38- }],
39- 'contexts' : [
29+ "apiVersion" : "v1" ,
30+ "clusters" : [
4031 {
41- 'context' :
42- {
43- 'cluster' : 'k8s' ,
44- 'user' : 'aws'
45- },
46- 'name' : 'aws'
47- }],
48- 'current-context' : 'aws' ,
49- 'Kind' : 'config' ,
50- 'users' : [
51- {
52- 'name' : 'aws' ,
53- 'user' : 'lambda'
54- }]
32+ "cluster" : {
33+ "server" : endpoint ,
34+ "certificate-authority-data" : certificate ,
35+ },
36+ "name" : "k8s" ,
37+ }
38+ ],
39+ "contexts" : [{"context" : {"cluster" : "k8s" , "user" : "aws" }, "name" : "aws" }],
40+ "current-context" : "aws" ,
41+ "Kind" : "config" ,
42+ "users" : [{"name" : "aws" , "user" : "lambda" }],
5543 }
5644
57- with open (KUBE_FILEPATH , 'w' ) as kube_file_content :
45+ with open (KUBE_FILEPATH , "w" ) as kube_file_content :
5846 yaml .dump (kube_config , kube_file_content , default_flow_style = False )
5947
6048
@@ -68,81 +56,87 @@ def get_bearer_token(cluster, region):
6856 STS_TOKEN_EXPIRES_IN = 60
6957 session = boto3 .session .Session ()
7058
71- client = session .client (' sts' , region_name = region )
59+ client = session .client (" sts" , region_name = region )
7260 service_id = client .meta .service_model .service_id
7361
7462 signer = RequestSigner (
75- service_id ,
76- region ,
77- 'sts' ,
78- 'v4' ,
79- session .get_credentials (),
80- session .events
63+ service_id , region , "sts" , "v4" , session .get_credentials (), session .events
8164 )
8265
8366 params = {
84- ' method' : ' GET' ,
85- ' url' : ' https://sts.{}.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15' .format (region ),
86- 'body' : {},
87- 'headers' : {
88- 'x-k8s-aws-id' : cluster
89- },
90- ' context' : {}
67+ " method" : " GET" ,
68+ " url" : " https://sts.{}.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15" .format (
69+ region
70+ ),
71+ "body" : {},
72+ "headers" : { "x-k8s-aws-id" : cluster },
73+ " context" : {},
9174 }
9275
9376 signed_url = signer .generate_presigned_url (
94- params ,
95- region_name = region ,
96- expires_in = STS_TOKEN_EXPIRES_IN ,
97- operation_name = ''
77+ params , region_name = region , expires_in = STS_TOKEN_EXPIRES_IN , operation_name = ""
9878 )
9979
100- base64_url = base64 .urlsafe_b64encode (signed_url .encode (' utf-8' )).decode (' utf-8' )
80+ base64_url = base64 .urlsafe_b64encode (signed_url .encode (" utf-8" )).decode (" utf-8" )
10181
10282 # need to remove base64 encoding padding:
10383 # https://github.com/kubernetes-sigs/aws-iam-authenticator/issues/202
104- return 'k8s-aws-v1.' + re .sub (r'=*' , '' , base64_url )
84+ return "k8s-aws-v1." + re .sub (r"=*" , "" , base64_url )
85+
10586
106- def get_evictable_pods (api , node_name ,label_selector ):
107- '''
87+ def get_evictable_pods (api , node_name , label_selector ):
88+ """
10889 This method will ensure we are only waiting for pods that matters based on
10990 label_selector
110- '''
111- field_selector = 'spec.nodeName=' + node_name
112- pods = api .list_pod_for_all_namespaces (watch = False , field_selector = field_selector ,
113- label_selector = label_selector , include_uninitialized = True )
91+ """
92+ field_selector = "spec.nodeName=" + node_name
93+ pods = api .list_pod_for_all_namespaces (
94+ watch = False ,
95+ field_selector = field_selector ,
96+ label_selector = label_selector ,
97+ include_uninitialized = True ,
98+ )
11499 return [pod for pod in pods .items ]
115100
116- def count_running_pods (api , node_name ,label_selector ):
117- '''
101+
102+ def count_running_pods (api , node_name , label_selector ):
103+ """
118104 Report count for total running pods based on the label
119- '''
120- pods = get_evictable_pods (api , node_name ,label_selector )
105+ """
106+ pods = get_evictable_pods (api , node_name , label_selector )
121107 return len (pods )
122108
109+
123110def handler (event , context ):
124- '''
111+ """
125112 Lambda handler, this function will call the
126113 private functions to get the running pod count based on the label selector provided
127- '''
128- eks = boto3 .client (' eks' , region_name = event [' region' ])
129- #loading Kube Config
114+ """
115+ eks = boto3 .client (" eks" , region_name = event [" region" ])
116+ # loading Kube Config
130117 if not os .path .exists (KUBE_FILEPATH ):
131- create_kube_config (eks , event [' cluster_name' ])
118+ create_kube_config (eks , event [" cluster_name" ])
132119 k8s .config .load_kube_config (KUBE_FILEPATH )
133120 configuration = k8s .client .Configuration ()
134- #getting the auth token
135- token = get_bearer_token (event [' cluster_name' ], event [' region' ])
136- configuration .api_key [' authorization' ] = token
137- configuration .api_key_prefix [' authorization' ] = ' Bearer'
121+ # getting the auth token
122+ token = get_bearer_token (event [" cluster_name" ], event [" region" ])
123+ configuration .api_key [" authorization" ] = token
124+ configuration .api_key_prefix [" authorization" ] = " Bearer"
138125 # API
139126 api = k8s .client .ApiClient (configuration )
140127 core_v1_api = k8s .client .CoreV1Api (api )
141128
142129 # Get all the pods
143- running_pod_count = count_running_pods (core_v1_api ,node_name = event ['node_name' ],
144- label_selector = event ['label_selector' ])
145- output_json = {"region" : event ['region' ], "node_name" : event ['node_name' ] ,
146- "instance_id" : event ['instance_id' ], "cluster_name" : event ['cluster_name' ],
147- "activePodCount" : running_pod_count }
130+ running_pod_count = count_running_pods (
131+ core_v1_api ,
132+ node_name = event ["node_name" ],
133+ label_selector = event ["label_selector" ],
134+ )
135+ output_json = {
136+ "region" : event ["region" ],
137+ "node_name" : event ["node_name" ],
138+ "instance_id" : event ["instance_id" ],
139+ "cluster_name" : event ["cluster_name" ],
140+ "activePodCount" : running_pod_count ,
141+ }
148142 return output_json
0 commit comments