import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))

import binascii
import numpy as np
from PIL import Image
from collections import defaultdict
from androguard.core.bytecodes.dvm import DalvikVMFormat
import utils.utils as utils
import subprocess

def get_width(size_kb):
    if size_kb < 10:
        width = 32
    elif 10 <= size_kb < 30:
        width = 64
    elif 30 <= size_kb < 60:
        width = 128
    elif 60 <= size_kb < 100:
        width = 256
    elif 100 <= size_kb < 200:
        width = 384
    elif 200 <= size_kb < 500:
        width = 512
    elif 500 <= size_kb < 1000:
        width = 768
    else:
        width = 1024

    return width

def load_1d_dex(dex_file, max_size=80):
    # Convert a binary file to a 1d array by reading byte values.
    file_size = os.stat(dex_file).st_size >> 20  # in MB
    if file_size < max_size:

        # Read the binary file in hex
        with open(dex_file, 'rb') as fp:
            hexstring = binascii.hexlify(fp.read())

            # Convert hex string to byte array
            # 2 hex numbers give up to 256 (1 byte) in decimal
            byte_list = [int(hexstring[i: i + 2], 16) for i in
                         range(0, len(hexstring), 2)]

            return np.array(byte_list, dtype=np.uint8)


def get_dex_info(dex_file):
    dex_info = defaultdict(int)

    with open(dex_file, 'rb') as f:
        dex = DalvikVMFormat(f.read())
        print(dex.get_header_item())
        dex_info['file_size'] = dex.get_header_item().file_size

        # DEX HEADER
        dex_info['header_size'] = dex.get_header_item().header_size

        # INDEX
        dex_info['string_ids_size'] = dex.get_header_item().string_ids_size
        dex_info['string_ids_off'] = dex.get_header_item().string_ids_off

        dex_info['type_ids_size'] = dex.get_header_item().type_ids_size
        dex_info['type_ids_off'] = dex.get_header_item().type_ids_off

        dex_info['proto_ids_size'] = dex.get_header_item().proto_ids_size
        dex_info['proto_ids_off'] = dex.get_header_item().proto_ids_off

        dex_info['field_ids_size'] = dex.get_header_item().field_ids_size
        dex_info['field_ids_off'] = dex.get_header_item().field_ids_off

        dex_info['method_ids_size'] = dex.get_header_item().method_ids_size
        dex_info['method_ids_off'] = dex.get_header_item().method_ids_off

        dex_info['class_defs_size'] = dex.get_header_item().class_defs_size
        dex_info['class_defs_off'] = dex.get_header_item().class_defs_off

        # DATA
        dex_info['data_size'] = dex.get_header_item().data_size
        dex_info['data_off'] = dex.get_header_item().data_off

        # LINK DATA
        dex_info['link_size'] = dex.get_header_item().link_size
        dex_info['link_off'] = dex.get_header_item().link_off

    return dex_info


def set_color(dex_info, dex_b, dex_r, dex_g):
    # HEADER -> red
    s0_end = dex_info['header_size']
    dex_r[0: s0_end] = dex_b[0:s0_end]
    dex_b[0: s0_end] = 0

    # DATA -> green
    s7_idx = dex_info['data_off']
    e7_idx = s7_idx + dex_info['data_size']
    dex_g[s7_idx:e7_idx] = dex_b[s7_idx:e7_idx]
    dex_b[s7_idx:e7_idx] = 0

    return dex_r, dex_g


def get_linear(dex_array, size_kb):
    width = get_width(size_kb)
    height = int(len(dex_array) / width)

    dex_array = dex_array[:height * width]  # trim array to 2-D shape
    dex_array_2d = np.reshape(dex_array, (height, width))

    return dex_array_2d


def create_linear_image(file_name):
    '''
    @format : {hash_value}.exe
    '''
    params = utils.config_parse('PRODUCTION')
    
    name = file_name.split('.')[0]
    dex_path = params['temp_data_path'] + 'malware/' + file_name
    save_img = params['temp_data_path'] + 'malimg/' + name + ".png"
    print(dex_path)
    dex_info = get_dex_info(dex_path)

    dex_b = load_1d_dex(dex_path)
    dex_r = np.zeros(len(dex_b), dtype=np.uint8)
    dex_g = np.zeros(len(dex_b), dtype=np.uint8)
    dex_r, dex_g = set_color(dex_info, dex_b, dex_r, dex_g)

    file_kb = int(os.path.getsize(dex_path) / 1000)
    b_c = get_linear(dex_b, file_kb)
    r_c = get_linear(dex_r, file_kb)
    g_c = get_linear(dex_g, file_kb)

    im = np.stack((r_c, g_c, b_c), axis=-1)
    im = Image.fromarray(im)

    image_size = 256
    
    im_resized = im.resize(size=(image_size, image_size), resample=Image.ANTIALIAS)  # ANTIALIAS provides higher quality image
    im_resized.save(save_img)
    return True

def apk2tdex(apk_path, dex_path):
    # Convert apk to dex
    try:
        subprocess.check_output(['java', '-jar', 'apktool.jar', 'd', apk_path, '-o', dex_path])
        return True
    except Exception as e:
        print(e)
        return False

if __name__ == '__main__':
    create_linear_image("putty.dex")