#everything is from here : https://www.youtube.com/watch?v=OQuRwpUTBpQ
#i added comments so i can understand the code better

import pefile   #working with a pe file
import re       #regular expression search
import struct   #convert data to numbers
import hashlib  #sha1 hash
import sys      #arg

def rc4crypt(data, key):
    #If the input is a string convert to byte arrays
    if type(data) == str:
        data = data.encode('utf-8')
    if type(key) == str:
        key = key.encode('utf-8')
    x = 0
    box = list(range(256))
    for i in range(256):
        x = (x + box[i] + key[i % len(key)]) % 256
        box[i], box[x] = box[x], box[i]
    x = 0
    y = 0
    out = []
    for c in data:
        x = (x + 1) % 256
        y = (y + box[x]) % 256
        box[x], box[y] = box[y], box[x]
        out.append(c ^ box[(box[x] + box[y]) % 256])
    return bytes(out)

if len(sys.argv) != 2 :
    print("Usage : python extract_config.py [sample.bin]")
    exit(0)

sample_path = sys.argv[1]
data = open(sample_path,"rb").read()

pe = pefile.PE(data=data)

"""
6A 08                                   push    8               ; dwDataLen
68 10 50 57 00                          push    offset pbData   ; pbData
68 00 20 00 00                          push    2000h           ; pdwDataLen
"""
sig = rb"\x6A(.)\x68(....)\x68\x00\x20\x00\x00" #'.' is a wildcard, same as '?', '()' is for a capture group, so we can catch their values
match = re.search(sig, data)

if match == None :
    """
    Signature got updated
    6A 08                                   push    8               ; dwDataLen
    68 10 40 00 10                          push    offset pbData   ; pbData
    56                                      push    esi             ; pdwDataLen
    """
    sig = rb"\x6A(.)\x68(....)\x56" #'.' is a wildcard, same as '?', '()' is for a capture group, so we can catch their values
    match = re.search(sig, data)

print("key length: %r" % match.group(1))
print("key address: %r" % match.group(2))

key_len = struct.unpack('b', match.group(1))[0]
key_addr = struct.unpack('<I', match.group(2))[0] #'<I' : little endian
key_rva = key_addr - pe.OPTIONAL_HEADER.ImageBase #RVA : Relative Virtual Address, we need that so we can translate it to a file offset
key_offset = pe.get_offset_from_rva(key_rva)
key_data = data[key_offset:key_offset+key_len]

print("key data: %r" % key_data)

config_data = data[key_offset+key_len:key_offset+key_len+0x2000]

m = hashlib.sha1()
m.update(key_data)      #sha1 the key data
key = m.digest()[:5]    #keeping only the first 5 bytes of the sha1 of key_data

config = rc4crypt(config_data, key)
build_id = config.split(b'\x00')[0].decode()

c2_string = ''
for s in config.split(b'\x00')[1:]:
    if s != b'':
        c2_string = s.decode()
        break

c2_list = c2_string.split("|")[:-1]
print("build id: %s" % build_id)
print("C2: %s" % c2_list)