""" Extract the required information for a given test file & feed it
into the trained model to classify the file as Benign or Malicious """

import pefile
import os
import sys
import pickle
import array
import math
import csv, json
import numpy as np
from helper_functions import *

def get_entropy(data):
	""" learn Entropy calculated across various sections & resources """
	if len(data) == 0:
		return 0.0

	occurences = array.array('L', [0]*256)
	for x in data:
		occurences[x if isinstance(x, int) else ord(x)] += 1

	entropy = 0
	for x in occurences:
		if x:
			p_x = float(x) / len(data)
			entropy -= p_x*math.log(p_x, 2)

	return entropy

def get_resources(pe):
	""" Extract resources : [entropy, size] """
	resources = []

	if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
		try:
			for resource_type in pe.DIRECTORY_ENTRY_RESOURCE.entries:
				if hasattr(resource_type, 'directory'):
					for resource_id in resource_type.directory.entries:
						if hasattr(resource_id, 'directory'):
							for resource_lang in resource_id.directory.entries:
								data = pe.get_data(resource_lang.data.struct.OffsetToData, resource_lang.data.struct.Size)
								size = resource_lang.data.struct.Size
								entropy = get_entropy(data)

								resources.append([entropy, size])
		except Exception as e:
			return resources

	return resources

def get_version_info(pe):
	""" Return version info """
	res = {}

	for fileinfo in pe.FileInfo:
		if fileinfo.Key == 'StringFileInfo':
			for st in fileinfo.StringTable:
				for entry in st.entries.items():
					res[entry[0]] = entry[1]
		if fileinfo.Key == 'VarFileInfo':
			for var in fileinfo.Var:
				res[var.entry.items()[0][0]] = var.entry.items()[0][1]

	if hasattr(pe, 'VS_FIXEDFILEINFO'):
		  res['flags'] = pe.VS_FIXEDFILEINFO.FileFlags
		  res['os'] = pe.VS_FIXEDFILEINFO.FileOS
		  res['type'] = pe.VS_FIXEDFILEINFO.FileType
		  res['file_version'] = pe.VS_FIXEDFILEINFO.FileVersionLS
		  res['product_version'] = pe.VS_FIXEDFILEINFO.ProductVersionLS
		  res['signature'] = pe.VS_FIXEDFILEINFO.Signature
		  res['struct_version'] = pe.VS_FIXEDFILEINFO.StrucVersion

	return res

def extract_infos(fpath):
	""" Extract the attributes for a given PEfile """
	res = {}

	try:
		#check if file is a PEfile or not
		pe = pefile.PE(fpath)
	except:
		return -1

	res['Machine'] = pe.FILE_HEADER.Machine
	res['SizeOfOptionalHeader'] = pe.FILE_HEADER.SizeOfOptionalHeader
	res['Characteristics'] = pe.FILE_HEADER.Characteristics
	res['MajorLinkerVersion'] = pe.OPTIONAL_HEADER.MajorLinkerVersion
	res['MinorLinkerVersion'] = pe.OPTIONAL_HEADER.MinorLinkerVersion
	res['SizeOfCode'] = pe.OPTIONAL_HEADER.SizeOfCode
	res['SizeOfInitializedData'] = pe.OPTIONAL_HEADER.SizeOfInitializedData
	res['SizeOfUninitializedData'] = pe.OPTIONAL_HEADER.SizeOfUninitializedData
	res['AddressOfEntryPoint'] = pe.OPTIONAL_HEADER.AddressOfEntryPoint
	res['BaseOfCode'] = pe.OPTIONAL_HEADER.BaseOfCode
	try:
		res['BaseOfData'] = pe.OPTIONAL_HEADER.BaseOfData
	except AttributeError:
		res['BaseOfData'] = 0
	res['ImageBase'] = pe.OPTIONAL_HEADER.ImageBase
	res['SectionAlignment'] = pe.OPTIONAL_HEADER.SectionAlignment
	res['FileAlignment'] = pe.OPTIONAL_HEADER.FileAlignment
	res['MajorOperatingSystemVersion'] = pe.OPTIONAL_HEADER.MajorOperatingSystemVersion
	res['MinorOperatingSystemVersion'] = pe.OPTIONAL_HEADER.MinorOperatingSystemVersion
	res['MajorImageVersion'] = pe.OPTIONAL_HEADER.MajorImageVersion
	res['MinorImageVersion'] = pe.OPTIONAL_HEADER.MinorImageVersion
	res['MajorSubsystemVersion'] = pe.OPTIONAL_HEADER.MajorSubsystemVersion
	res['MinorSubsystemVersion'] = pe.OPTIONAL_HEADER.MinorSubsystemVersion
	res['SizeOfImage'] = pe.OPTIONAL_HEADER.SizeOfImage
	res['SizeOfHeaders'] = pe.OPTIONAL_HEADER.SizeOfHeaders
	res['CheckSum'] = pe.OPTIONAL_HEADER.CheckSum
	res['Subsystem'] = pe.OPTIONAL_HEADER.Subsystem
	res['DllCharacteristics'] = pe.OPTIONAL_HEADER.DllCharacteristics
	res['SizeOfStackReserve'] = pe.OPTIONAL_HEADER.SizeOfStackReserve
	res['SizeOfStackCommit'] = pe.OPTIONAL_HEADER.SizeOfStackCommit
	res['SizeOfHeapReserve'] = pe.OPTIONAL_HEADER.SizeOfHeapReserve
	res['SizeOfHeapCommit'] = pe.OPTIONAL_HEADER.SizeOfHeapCommit
	res['LoaderFlags'] = pe.OPTIONAL_HEADER.LoaderFlags
	res['NumberOfRvaAndSizes'] = pe.OPTIONAL_HEADER.NumberOfRvaAndSizes

	# Sections
	res['SectionsNb'] = len(pe.sections)
	entropy = list(map(lambda x:x.get_entropy(), pe.sections))
	res['SectionsMeanEntropy'] = sum(entropy)/float(len(entropy))
	res['SectionsMinEntropy'] = min(entropy)
	res['SectionsMaxEntropy'] = max(entropy)
	raw_sizes = list(map(lambda x:x.SizeOfRawData, pe.sections))
	res['SectionsMeanRawsize'] = sum(raw_sizes)/float(len(raw_sizes))
	res['SectionsMinRawsize'] = min(raw_sizes)
	res['SectionsMaxRawsize'] = max(raw_sizes)
	virtual_sizes = list(map(lambda x:x.Misc_VirtualSize, pe.sections))
	res['SectionsMeanVirtualsize'] = sum(virtual_sizes)/float(len(virtual_sizes))
	res['SectionsMinVirtualsize'] = min(virtual_sizes)
	res['SectionMaxVirtualsize'] = max(virtual_sizes)

	#Imports
	try:
		res['ImportsNbDLL'] = len(pe.DIRECTORY_ENTRY_IMPORT)
		imports = sum([x.imports for x in pe.DIRECTORY_ENTRY_IMPORT], [])
		res['ImportsNb'] = len(imports)
		res['ImportsNbOrdinal'] = len(list(filter(lambda x:x.name is None, imports)))
	except AttributeError:
		# No import
		res['ImportsNbDLL'] = 0
		res['ImportsNb'] = 0
		res['ImportsNbOrdinal'] = 0

	#Exports
	try:
		res['ExportNb'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
	except AttributeError:
		# No export
		res['ExportNb'] = 0

	#Resources
	resources= get_resources(pe)
	res['ResourcesNb'] = len(resources)

	if len(resources)> 0:
		entropy = list(map(lambda x:x[0], resources))
		res['ResourcesMeanEntropy'] = sum(entropy)/float(len(entropy))
		res['ResourcesMinEntropy'] = min(entropy)
		res['ResourcesMaxEntropy'] = max(entropy)
		sizes = list(map(lambda x:x[1], resources))
		res['ResourcesMeanSize'] = sum(sizes)/float(len(sizes))
		res['ResourcesMinSize'] = min(sizes)
		res['ResourcesMaxSize'] = max(sizes)
	else:
		res['ResourcesNb'] = 0
		res['ResourcesMeanEntropy'] = 0
		res['ResourcesMinEntropy'] = 0
		res['ResourcesMaxEntropy'] = 0
		res['ResourcesMeanSize'] = 0
		res['ResourcesMinSize'] = 0
		res['ResourcesMaxSize'] = 0

	# Load configuration size
	try:
		res['LoadConfigurationSize'] = pe.DIRECTORY_ENTRY_LOAD_CONFIG.struct.Size
	except AttributeError:
		res['LoadConfigurationSize'] = 0

	# Version configuration size
	try:
		version_infos = get_version_info(pe)
		res['VersionInformationSize'] = len(version_infos.keys())
	except AttributeError:
		res['VersionInformationSize'] = 0

	return res

if __name__ == '__main__':
	b=0
	m=0
	lis = []
	lis = os.listdir('Benign/')

	parameters = pickle.loads(open('classifier.pkl','rb').read())
	features = pickle.loads(open('features.pkl','rb').read())
	scaler = pickle.loads(open('scaler.pkl','rb').read())
	for i in range(len(lis)):
		#data = extract_infos(sys.argv[1])
		path = 'Benign/' + lis[i]
		data = extract_infos(path)
		if(data!=-1):
			pe_features = list(map(lambda x:data[x], features))
			pe_features = np.array(pe_features)
			pe_features = pe_features.reshape(1, -1)
			X = scaler.transform(pe_features)
			res = predict_binaries(X.T, parameters)
			print ('The file is %s' % (res))
			if(res == 'Benign'):
				b += 1
			else:
				m += 1
	print("No. of Benign files: " + str(b))
	print("No. of malicious files: " + str(m))


#!/bin/python3.6
# fork of previous ones forked from Apache MXNet examples
# https://github.com/tspannhw/mxnet_rpi/blob/master/analyze.py
import time
from hdfs import InsecureClient
import sys
import datetime
import subprocess
import sys
import os
import datetime
import traceback
import math
import random, string
import base64
import json
from time import gmtime, strftime
import mxnet as mx
import inception_predict
import numpy as np
import math
import random, string
import time

from time import gmtime, strftime
start = time.time()

packet_size=3000

# Create unique image name
uniqueid = 'mxnet_uuid_{0}_{1}'.format('img',strftime("%Y%m%d%H%M%S",gmtime()))

filename = sys.argv[1]
topn = []
# Run inception prediction on image
try:
     topn = inception_predict.predict_from_local_file(filename, N=5)
except:
     print("Error")
     errorcondition = "true"

try:
     # 5 MXNET Analysis
     top1 = str(topn[0][1])
     top1pct = str(round(topn[0][0],3) * 100)

     top2 = str(topn[1][1])
     top2pct = str(round(topn[1][0],3) * 100)

     top3 = str(topn[2][1])
     top3pct = str(round(topn[2][0],3) * 100)

     top4 = str(topn[3][1])
     top4pct = str(round(topn[3][0],3) * 100)

     top5 = str(topn[4][1])
     top5pct = str(round(topn[4][0],3) * 100)

     end = time.time()

     print('UUID:' + uniqueid + ' top1:' + top1 + ' ' + top1pct + '%')

     row = { 'uuid': uniqueid,  'top1pct': top1pct, 'top1': top1, 'top2pct': top2pct, 'top2': top2,'top3pct': top3pct, 'top3': top3,'top4pct': top4pct,'top4': top4, 'top5pct': top5pct,'top5': top5, 'imagefilename': filename, 'runtime': str(round(end - start)) }

     client = InsecureClient('http://princeton0.field.hortonworks.com:50070', user='root')
     from json import dumps
     client.write('/mxnetyarn/' + uniqueid + '.json', dumps(row))

except:
     print("{\"message\": \"Failed to run\"}")


import argparse
import json
import numpy as np
import pickle
import random
import torch
import torch.backends.cudnn as cudnn
from torch.utils.data import ConcatDataset, DataLoader, Subset
from base_model import BaseModel
from datasets import get_dataset
from attackers import MiaAttack

parser = argparse.ArgumentParser(description='Membership inference Attacks on Network Pruning')
parser.add_argument('device', default=0, type=int, help="GPU id to use")
parser.add_argument('config_path', default=0, type=str, help="config file path")
parser.add_argument('--dataset_name', default='mnist', type=str)
parser.add_argument('--model_name', default='mnist', type=str)
parser.add_argument('--num_cls', default=10, type=int)
parser.add_argument('--input_dim', default=1, type=int)
parser.add_argument('--image_size', default=28, type=int)
parser.add_argument('--hidden_size', default=128, type=int)
parser.add_argument('--seed', default=7, type=int)
parser.add_argument('--early_stop', default=5, type=int)
parser.add_argument('--batch_size', default=128, type=int)
parser.add_argument('--prune_epochs', default=50, type=int)
parser.add_argument('--pruner_name', default='l1unstructure', type=str, help="prune method for victim model")
parser.add_argument('--prune_sparsity', default=0.7, type=float, help="prune sparsity for victim model")
parser.add_argument('--adaptive', action='store_true', help="use adaptive attack")
parser.add_argument('--shadow_num', default=5, type=int)
parser.add_argument('--defend', default='', type=str)
parser.add_argument('--defend_arg', default=4, type=float)
parser.add_argument('--attacks', default="samia", type=str)
parser.add_argument('--original', action='store_true', help="original=true, then launch attack against original model")

def main(args):
    # import torch.multiprocessing
    # torch.multiprocessing.set_sharing_strategy('file_system')

    torch.manual_seed(args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)
    device = f"cuda:{args.device}"
    cudnn.benchmark = True
    prune_prefix = f"{args.pruner_name}_{args.prune_sparsity}" \
                   f"{'_' + args.defend if args.defend else ''}{'_' + str(args.defend_arg) if args.defend else ''}"
    prune_prefix2 = f"{args.pruner_name}_{args.prune_sparsity}" \
                    f"{'_' + args.defend if args.adaptive else ''}{'_' + str(args.defend_arg) if args.adaptive else ''}"

    save_folder = f"results/{args.dataset_name}_{args.model_name}"

    print(f"Save Folder: {save_folder}")

    # Load datasets
    trainset = get_dataset(args.dataset_name, train=True)
    testset = get_dataset(args.dataset_name, train=False)
    if testset is None:
        total_dataset = trainset
    else:
        total_dataset = ConcatDataset([trainset, testset])
    total_size = len(total_dataset)
    data_path = f"{save_folder}/data_index.pkl"
    with open(data_path, 'rb') as f:
        victim_train_list, victim_dev_list, victim_test_list, attack_split_list \
            = pickle.load(f)
    victim_train_dataset = Subset(total_dataset, victim_train_list)
    victim_test_dataset = Subset(total_dataset, victim_test_list)
    print(f"Total Data Size: {total_size}, "
          f"Victim Train Size: {len(victim_train_list)}, "
          f"Victim Test Size: {len(victim_test_list)}")
    victim_train_loader = DataLoader(victim_train_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4,
                                     pin_memory=False)
    victim_test_loader = DataLoader(victim_test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4,
                                    pin_memory=False)

    # Load pruned victim model
    victim_model_save_folder = save_folder + "/victim_model"
    victim_model_path = f"{victim_model_save_folder}/best.pth"
    victim_model = BaseModel(args.model_name, num_cls=args.num_cls, input_dim=args.input_dim, device=device)
    victim_model.load(victim_model_path)

    pruned_model_save_folder = f"{save_folder}/{prune_prefix}_model"
    print(f"Load Pruned Model from {pruned_model_save_folder}")
    victim_pruned_model = BaseModel(
        args.model_name, num_cls=args.num_cls, input_dim=args.input_dim, save_folder=pruned_model_save_folder,
        device=device)
    victim_pruned_model.model.load_state_dict(torch.load(f"{pruned_model_save_folder}/best.pth"))
    victim_pruned_model.test(victim_train_loader, "Victim Pruned Model Train")
    victim_pruned_model.test(victim_test_loader, "Victim Pruned Model Test")

    # Load pruned shadow models
    shadow_model_list, shadow_prune_model_list, shadow_train_loader_list, shadow_test_loader_list = [], [], [], []
    for shadow_ind in range(args.shadow_num):
        attack_train_list, attack_dev_list, attack_test_list = attack_split_list[shadow_ind]
        shadow_train_dataset = Subset(total_dataset, attack_train_list)
        shadow_dev_dataset = Subset(total_dataset, attack_dev_list)
        shadow_test_dataset = Subset(total_dataset, attack_test_list)
        shadow_train_loader = DataLoader(shadow_train_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4,
                                         pin_memory=False)
        shadow_dev_loader = DataLoader(shadow_dev_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4,
                                       pin_memory=False)
        shadow_test_loader = DataLoader(shadow_test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4,
                                        pin_memory=False)

        shadow_model_path = f"{save_folder}/shadow_model_{shadow_ind}/best.pth"
        shadow_model = BaseModel(args.model_name, num_cls=args.num_cls, input_dim=args.input_dim, device=device)
        shadow_model.load(shadow_model_path)
        pruned_shadow_model_save_folder = f"{save_folder}/shadow_{prune_prefix2}_model_{shadow_ind}"
        print(f"Load Pruned Shadow Model From {pruned_shadow_model_save_folder}")
        shadow_pruned_model = BaseModel(
            args.model_name, num_cls=args.num_cls, input_dim=args.input_dim,
            save_folder=pruned_shadow_model_save_folder, device=device)
        shadow_pruned_model.model.load_state_dict(torch.load(f"{pruned_shadow_model_save_folder}/best.pth"))
        shadow_pruned_model.test(shadow_train_loader, "Shadow Pruned Model Train")
        shadow_pruned_model.test(shadow_test_loader, "Shadow Pruned Model Test")

        shadow_model_list.append(shadow_model)
        shadow_prune_model_list.append(shadow_pruned_model)
        shadow_train_loader_list.append(shadow_train_loader)
        shadow_test_loader_list.append(shadow_test_loader)

    print("Start Membership Inference Attacks")

    if args.original:
        attack_original = True
    else:
        attack_original = False
    attacker = MiaAttack(
        victim_model, victim_pruned_model, victim_train_loader, victim_test_loader,
        shadow_model_list, shadow_prune_model_list, shadow_train_loader_list, shadow_test_loader_list,
        num_cls=args.num_cls, device=device, batch_size=args.batch_size,
        attack_original=attack_original)

    attacks = args.attacks.split(',')

    if "samia" in attacks:
        nn_trans_acc = attacker.nn_attack("nn_sens_cls", model_name="transformer")
        print(f"SAMIA attack accuracy {nn_trans_acc:.3f}")

    if "threshold" in attacks:
        conf, xent, mentr, top1_conf = attacker.threshold_attack()
        print(f"Ground-truth class confidence-based threshold attack (Conf) accuracy: {conf:.3f}")
        print(f"Cross-entropy-based threshold attack (Xent) accuracy: {xent:.3f}")
        print(f"Modified-entropy-based threshold attack (Mentr) accuracy: {mentr:.3f}")
        print(f"Top1 Confidence-based threshold attack (Top1-conf) accuracy: {top1_conf:.3f}")

    if "nn" in attacks:
        nn_acc = attacker.nn_attack("nn")
        print(f"NN attack accuracy {nn_acc:.3f}")

    if "nn_top3" in attacks:
        nn_top3_acc = attacker.nn_attack("nn_top3")
        print(f"Top3-NN Attack Accuracy {nn_top3_acc}")

    if "nn_cls" in attacks:
        nn_cls_acc = attacker.nn_attack("nn_cls")
        print(f"NNCls Attack Accuracy {nn_cls_acc}")

if __name__ == '__main__':
    args = parser.parse_args()
    with open(args.config_path) as f:
        t_args = argparse.Namespace()
        t_args.__dict__.update(json.load(f))
        args = parser.parse_args(namespace=t_args)

    print(args)
    main(args)