#!/usr/bin/env python

import boto3
import numpy as np
import argparse
import ast
from sklearn.preprocessing import StandardScaler

### Change the following to the correct endpoint name ###
myEndpointName = 'sagemaker-tensorflow-serving-2021-10-13-22-34-33-929'
def main():
 
    import json
    import ember
    
    from sklearn.preprocessing import RobustScaler
    rs = RobustScaler()
       
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--featureversion", type=int, default=2, help="EMBER feature version")
    parser.add_argument("binaries", metavar="BINARIES", type=str, nargs="+", help="PE files to classify")
    args = parser.parse_args()
    #opening the downloaded PE file
    data = open(args.binaries[0],'rb').read()
    #Feature extractor class of the ember project 
    mms = StandardScaler()
    
    extractor = ember.PEFeatureExtractor()
    data = extractor.feature_vector(data)
    data = rs.fit_transform([data])
    data = np.reshape(data,(1,2381))
    data= data.tolist()

    client = boto3.client('runtime.sagemaker',
				region_name='us-east-1',
                              	#### Change the following to your AWS credentials ####
				aws_access_key_id='ASIAXVT2C4KGVVT6WYMD',
                aws_secret_access_key='mOpnAWCSWAoyXpv7EG+eOQYUxM7g9iRtENOKSGPF',
                aws_session_token='FwoGZXIvYXdzEM3//////////wEaDDjPcCyMgM5yQP9GSyK8Ab0BFEGLeDJEu0fDrYmvmNWgNTh6pQ1n5YAUJzaY5LcfksNfCwI3/60HSZvtIbSBuw7CFmbU4dhDWcFiXSsTLPtQCxS9WgV1lGstYOam+eXzakDYHrF4Fe90A6dIFYctO60R5sH5PV7ZzHFAEUz+uM63hfsRdTUKucYB5sPV3w9tx9u36Tjf4u4AVjaF79qVJZFlBmXkkfV2GQNHnfz7DH/CaI/BAk/6+jedg4/SMMYNr0ye4UFlSxi9vtjWKOitoIsGMi2pEv22YrbFtOgzD+ofVnT4JwAGxvLQJ2PPC4+vDkIl2B3aSiO7Z401FRboWeQ=')
    
    response = client.invoke_endpoint(EndpointName=myEndpointName, Body=json.dumps(data),ContentType='application/json')
    response_body = response['Body']
    out = response_body.read()
    astr = out.decode("UTF-8")
    out = ast.literal_eval(astr)
    print(out)

    #if out[0] >0.5:
    #    print("Malicious")
    #else:
    #    print("Benign")
		
if __name__ == "__main__":
	main()
