#builtin lambda imports
import json
import boto3
import pprint

#external lambda imports
from androguard.core.bytecodes.apk import APK

#connecting before lambda_handler to keep during warm boot
s3_client = boto3.client('s3')


def return_code(status, body):
    return {"statusCode": status,"body": body}


def getFileObject(Bucket,Key):
    print("getting file object from s3")
    try:
        file_object = s3_client.get_object(Bucket=Bucket,Key=Key)
        if file_object:
            return {'code':'SUCCESS','content':file_object['Body'].read()}
        else:
            return {'code':'ERROR','content':'Could not get file from s3.'}
    except Exception as E:
        print(E)
        if "NoSuchKey" in str(E):
            return {'code':'NOT_FOUND','content':'File not found.'}
        else:
            return {'code':'ERROR','content':'Unknown error'}


def getAllInfo(buff):
    try:
        apkInfo = {"package":"","androidversion":"","permissions":"","uses_permissions":"",
        "declared_permissions":"","valid_apk":"","_is_signed_v2":"","_v2_blocks":""}
        apkObject = APK(buff,raw=True)
        apkInfo.update({"package":apkObject.package})
        apkInfo.update({"androidversion":apkObject.androidversion})
        apkInfo.update({"permissions":apkObject.permissions})
        apkInfo.update({"uses_permissions":apkObject.uses_permissions})
        apkInfo.update({"declared_permissions":apkObject.declared_permissions})
        apkInfo.update({"valid_apk":apkObject.valid_apk})
        apkInfo.update({"_is_signed_v2":apkObject._is_signed_v2})
        apkInfo.update({"_v2_blocks":apkObject._v2_blocks})
        return {'code':'SUCCESS','content':apkInfo}
    except Exception as E:
        print(E)
        return {'code':'ERROR','content':'Unknown error.'}
        


def lambda_handler(event, context):
    try:
        if event['pathParameters'] and event['pathParameters']['proxy']:
            sha256 = event['pathParameters']['proxy']
            fileObject = getFileObject(Bucket='malscanbotfiles', Key=sha256)
            if fileObject['code'] == 'SUCCESS':
                apkInfo = getAllInfo(fileObject['content'])
                if apkInfo['code'] == 'SUCCESS':
                    return return_code(200, json.dumps(apkInfo['content']))
                else:
                    return return_code(501, 'androguard failed')
            else:
                return return_code(404,'Error, APK file not found.')
        else:
            return return_code(400, 'sha256 param must be present.')
                
    except Exception as E:
        print(E)
        return return_code(501, 'Internal server error')
        