#!/usr/bin/env python3

'''
A set of tools that can be used to help obfuscate the malware
obf_string is used interactive: python -i ./obfuscation.py

Given 2 arguments to the program, the output will generate code to be used with the malware
for function obfuscation:
    ex:
    python string kernel32.dll
    python function ProcessOpen
'''
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from os import urandom
from random import choice
import hashlib
import string
import sys

def obf_string():
    """
    Output 15 random ascii letters to replace the string
    """
    letters = string.ascii_letters 
    return ''.join(choice(letters) for i in range(15))

def encrypt_AES(plaintext, key):
    k = hashlib.sha256(key).digest()
    iv = b'\x00' * 16
    plaintext = pad(plaintext, AES.block_size)
    cipher = AES.new(k, AES.MODE_CBC, iv)
    return cipher.encrypt(plaintext)

def print_obf_str(string):
    # in c, strings are terminated with a null byte
    string_w_null = string.encode() + b'\x00' 
    key = urandom(16)
    enc_name = encrypt_AES(string_w_null, key)
    
    # make var names for c syntax
    string = string.replace(".","").capitalize()

    # print out the string
    out_1 = "unsigned char s" + string + "[] = "
    out_1_hex = '{ 0x' + ', 0x'.join(hex(ord(chr(x)))[2:] for x in enc_name) + ' };'
    print(out_1 + out_1_hex)

    # print out the key
    out_2 = "unsigned char k" + string + "[] = "
    out_2_hex = '{ 0x' + ', 0x'.join(hex(ord(chr(x)))[2:] for x in key) + ' };'
    print(out_2 + out_2_hex)

    # print out the AES decryption for the string
    out_3 = "AESDecrypt((char *) s" + string + \
            ", sizeof(s" + string + "), k" + \
            string + ", sizeof(k" + string + "));"
    print(out_3)

def print_obf_func(string):
    # in c, strings are terminated with a null byte
    string_w_null = string.encode() + b'\x00'
    key = urandom(16)
    enc_name = encrypt_AES(string_w_null, key)
    

    # print out the function name string
    out_1 = "unsigned char s" + string + "[] = "
    out_1_hex = '{ 0x' + ', 0x'.join(hex(ord(chr(x)))[2:] for x in enc_name) + ' };'
    print(out_1 + out_1_hex)

    # print out the key for the function name string
    out_2 = "unsigned char k" + string + "[] = "
    out_2_hex = '{ 0x' + ', 0x'.join(hex(ord(chr(x)))[2:] for x in key) + ' };'
    print(out_2 + out_2_hex)

    # print out the AES decryption for the function
    out_3 = "AESDecrypt((char *) s" + string + \
            ", sizeof(s" + string + "), k" + \
            string + ", sizeof(k" + string + "));"
    print(out_3)

    # print out the pointer to the function
    out_4 = "p" + string + " = GetProcAddress(GetModuleHandle(sKernel32dll), s" + \
        string + ");"
    print(out_4)



def main():
    if len(sys.argv) != 3:
        usage = "Usage: obfuscation.py <string/function> <function_name>"
        print(usage)
        exit()
    
    if sys.argv[1] == "string":
        print_obf_str(sys.argv[2])
    elif sys.argv[1] == "function":
        print_obf_func(sys.argv[2])

if __name__ == "__main__":
    main()