import json
from botocore.vendored import requests
import boto3
import base64
from botocore.exceptions import ClientError

import magic

malscanbot_s3_bucket_name = "malscanbotfiles"
malscanbotfiles_table_name = 'malscanbotfiles'

try:
    dynamodb = boto3.resource('dynamodb')
    table = dynamodb.Table(malscanbotfiles_table_name)
except Exception as E:
    print(E)  
    
    
try:
    s3_client = boto3.client('s3')
except Exception as e:
    print(e)

def return_code(status, body):
    return {"statusCode": status,"body":json.dumps(body)}
    
    
def getFileObject(Bucket,Key):
    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:
        if "NoSuchKey" in str(E):
            return {'code':'NOT_FOUND','content':'File not found.'}
        else:
            return {'code':'ERROR','content':'Unknown error'}
    
def get_mime_type_from_dynamodb(sha256):
    try:
        sha256 = sha256.lower()
        result = table.get_item(Key={"sha256":sha256})
        if "Item" in result.keys() and "mime_type" in result['Item'].keys():
            return {"code":"SUCCESS","content": result['Item']['mime_type']}
        else:
            fileObject = getFileObject(Bucket=malscanbot_s3_bucket_name, Key=sha256)
            if fileObject['code'] == 'SUCCESS':
                file_magic = magic.from_buffer(fileObject['content'])
                UpdateExpression="set mime_type = :r"
                ExpressionAttributeValues={':r': file_magic}
                response = table.update_item(Key={"sha256":sha256.lower()},UpdateExpression=UpdateExpression,ExpressionAttributeValues=ExpressionAttributeValues)
                return {"code":"SUCCESS","content": file_magic}
    except Exception as E:
        print(E)
        return {'code':'UNKNOWN','content':'Unknown error.'}


def lambda_handler(event, context):
    try:
        if event['pathParameters'] and event['pathParameters']['proxy']:
            sha256 = event['pathParameters']['proxy']
            sha256 = sha256.lower()
            result = get_mime_type_from_dynamodb(sha256)
            if result['code'] == 'SUCCESS':
                return return_code(200,result['content'])
            if result['code'] == 'NOT_FOUND':
                return return_code(404,"mime_type value not found")
            else:
                return return_code(501, 'Internal server error')
        else:
            return return_code(400, 'sha256 param must be present.')
    except Exception as E:
        print(E)
        return return_code(501, 'Internal server error')