code added to register and validate daemon

Author: prashantramangupta

auth_token.py

@@ -0,0 +1,51 @@
+import calendar
+import secrets
+import time
+from datetime import datetime as dt
+
+from repository import Repository
+
+
+class Token:
+    def __init__(self, net_id):
+        self.net_id = net_id
+        self.repo = Repository(net_id)
+
+    def process_token(self, daemon_id):
+        print("process_token::daemon_id: ", daemon_id);
+        result = {}
+        token = secrets.token_urlsafe(64)
+        exist_info = self.daemon_id_exist(daemon_id=daemon_id)
+        if exist_info.get('status', False):
+            qry = "UPDATE daemon_token set token = %s WHERE daemon_id = %s "
+            updt_info = self.repo.execute(qry, [token, daemon_id])
+            print(updt_info)
+            if updt_info[0] > 0:
+                return {"token": token}
+            else:
+                return {"error": "unable to generate token"}
+        ctime_epoch = calendar.timegm(time.gmtime())
+        expiration = ctime_epoch + (6 * 60 * 60)
+        qry = "INSERT INTO daemon_token (daemon_id, token, expiration, row_updated, row_created) " \
+              "VALUES(%s, %s, %s, %s, %s)"
+        res = self.repo.execute(qry, [daemon_id, token, expiration, dt.utcnow(), dt.utcnow()])
+        if len(res) > 0 and res[0] > 0:
+            result["token"] = token
+        return result
+
+    def validate_token(self, daemon_id, token):
+        print("validate_token::daemon_id: ", daemon_id);
+        qry = "SELECT * FROM daemon_token WHERE daemon_id = %s and token = %s "
+        res = self.repo.execute(qry, [daemon_id, token])
+        if len(res) > 0:
+            return {'validated': True}
+        return {'validated': False}
+
+
+    def daemon_id_exist(self, daemon_id):
+        print("daemon_id_exist::daemon_id: ", daemon_id);
+        qry = "SELECT * FROM daemon_token WHERE daemon_id = %s"
+        res = self.repo.execute(qry, [daemon_id])
+        if len(res) > 0:
+            return {'status': True, "token": res[0]['token']}
+        return {'status': False}

build.sh

@@ -0,0 +1,2 @@
+rm daemon_authenticate.zip
+zip -r daemon_authenticate.zip * -x \*venv\*

lambda_handler.py

@@ -0,0 +1,92 @@
+import json
+import os
+
+from auth_token import Token
+from constant import DEFAULT_NETWORK_ID
+
+def request_handler(event, context):
+    print(event)
+    if 'path' not in event:
+        return get_response("400", {"status": "failed", "error": "Bad Request"})
+    try:
+        path = event['path'].lower()
+        data = None
+        if "/register" == path:
+            payload = event['body']
+            payload_dict = payload_check(payload=payload, path=path)
+            net_id = DEFAULT_NETWORK_ID
+            token_instance = Token(net_id)
+            data = token_instance.process_token(daemon_id=payload_dict['daemonId'])
+            if data is None:
+                response = get_response("400", {"status": "failed", "error": "Bad Request"})
+            else:
+                if data.get('error', '') == '':
+                    print(data)
+                    response = get_response("200", {"status": "success", "data": data})
+                else:
+                    error = data['error']
+                    data.pop('error')
+                    response = get_response("200", {"status": "failed", "data": data, "error": error})
+        elif "/event" == path:
+            try:
+                payload_dict = event['headers']
+                print("Processing [" + str(path) + "] with body [" + str(payload_dict) + "]")
+                net_id = 42
+                token_instance = Token(net_id)
+                data = token_instance.validate_token(daemon_id=payload_dict['x-daemonid'], token=payload_dict['x-token'])
+                response = get_lambda_authorizer_response_format(event=event, allow=data['validated'])
+            except Exception as e:
+                print(repr(e))
+                response = get_lambda_authorizer_response_format(event=event, allow=False)
+
+    except Exception as e:
+        response = get_response(500, {"status": "failed",
+                                      "error": repr(e)})
+
+    return response
+
+
+def payload_check(payload, path):
+    payload_dict = None
+    if payload is not None and len(payload) > 0:
+        payload_dict = json.loads(payload)
+    print("Processing [" + str(path) + "] with body [" + str(payload) + "]")
+    return payload_dict
+
+
+# Generate response JSON that API gateway expects from the lambda function
+def get_response(status_code, message):
+    return {
+        'statusCode': status_code,
+        'body': json.dumps(message),
+        'headers': {
+            'Content-Type': 'application/json',
+            "X-Requested-With": '*',
+            "Access-Control-Allow-Headers": 'Access-Control-Allow-Origin, Content-Type,X-Amz-Date,Authorization,X-Api-Key,x-requested-with',
+            "Access-Control-Allow-Origin": '*',
+            "Access-Control-Allow-Methods": 'GET,OPTIONS,POST'
+        }
+    }
+
+
+def get_lambda_authorizer_response_format(event, allow):
+    print(allow)
+    response = {
+        "principalId": os.environ['principalId'],
+        "policyDocument": {
+            "Version": '2012-10-17',
+            "Statement": [
+                {
+                    "Action": 'execute-api:Invoke',
+                    "Resource": event['methodArn']
+                }
+            ]
+        }
+    }
+    if allow:
+        response["policyDocument"]["Statement"][0]["Effect"] = 'Allow'
+    else:
+        response["policyDocument"]["Statement"][0]["Effect"] = 'Deny'
+
+    print(response)
+    return response

repository.py

@@ -0,0 +1,63 @@
+import pymysql
+from constant import NETWORKS
+
+
+class Repository:
+    connection = None
+
+    def __init__(self, netId):
+        self.DB_HOST = NETWORKS[netId]['db']['DB_HOST']
+        self.DB_USER = NETWORKS[netId]['db']['DB_USER']
+        self.DB_PASSWORD = NETWORKS[netId]['db']['DB_PASSWORD']
+        self.DB_NAME = NETWORKS[netId]['db']['DB_NAME']
+        self.DB_PORT = 3306
+        self.connection = self.__get_connection()
+        self.auto_commit = True
+
+    def execute(self, query, params=None):
+        return self.__execute_query(query, params)
+
+    def __get_connection(self):
+        open = True
+        if self.connection is not None:
+            try:
+                self.execute("select 1")
+                open = False
+            except Exception as e:
+                open = True
+
+        if open:
+            self.connection = pymysql.connect(self.DB_HOST, user=self.DB_USER,
+                                              passwd=self.DB_PASSWORD, db=self.DB_NAME, port=self.DB_PORT)
+        return self.connection
+
+    def __execute_query(self, query, params=None):
+        result = list()
+        try:
+            with self.connection.cursor() as cursor:
+                qry_resp = cursor.execute(query, params)
+                db_rows = cursor.fetchall()
+                if cursor.description is not None:
+                    field_name = [field[0] for field in cursor.description]
+                    for values in db_rows:
+                        row = dict(zip(field_name, values))
+                        result.append(row)
+                else:
+                    result.append(qry_resp)
+                    result.append({'last_row_id': cursor.lastrowid})
+                if self.auto_commit:
+                    self.connection.commit()
+        except Exception as e:
+            self.connection.rollback()
+            print("DB Error in %s, error: %s" % (str(query), repr(e)))
+        return result
+
+    def bulk_query(self, query, params=None):
+        try:
+            with self.connection.cursor() as cursor:
+                result = cursor.executemany(query, params)
+                self.connection.commit()
+                return result
+        except Exception as err:
+            self.connection.rollback()
+            print("DB Error in %s, error: %s" % (str(query), repr(err)))

requirement.txt

@@ -0,0 +1 @@
+pymysql==0.9.2