''' IDAPython script that generates a YARA rule to match against the basic blocks of the current function. It masks out relocation bytes and ignores jump instructions (given that we're already trying to match compiler-specific bytes, this is of arguable benefit).

If python-yara is installed, the IDAPython script also validates that the generated rule matches at least one segment in the current file.

author: Willi Ballenthin william.ballenthin@fireeye.com '''

import logging from collections import namedtuple

import idc import idaapi import idautils

logger = logging.getLogger(name)

BasicBlock = namedtuple('BasicBlock', ['va', 'size'])

each rule must have at least this many non-masked bytes

MIN_BB_BYTE_COUNT = 4

def get_basic_blocks(fva): ''' return sequence of BasicBlock instances for given function. ''' ret = [] func = idaapi.get_func(fva) if func is None: return ret

for bb in idaapi.FlowChart(func):
    ret.append(BasicBlock(va=bb.startEA,
                          size=bb.endEA - bb.startEA))

return ret

def get_function(va): ''' return va for first instruction in function that contains given va. ''' return idaapi.get_func(va).startEA

Rule = namedtuple('Rule', ['name', 'bytes', 'masked_bytes'])

def is_jump(va): ''' return True if the instruction at the given address appears to be a jump. ''' return idc.GetMnem(va).startswith('j')

def get_basic_block_rule(bb): ''' create and format a YARA rule for a single basic block. mask relocation bytes into unknown bytes (like '??'). do not include final instructions if they are jumps. ''' # fetch the instruction start addresses insns = [] va = bb.va while va < bb.va + bb.size: insns.append(va) va = idc.NextHead(va)

# drop the last instruction if its a jump
if is_jump(insns[-1]):
    insns = insns[:-1]

bytes = []
# `masked_bytes` is the list of formatted bytes,
#   not yet join'd for performance.
masked_bytes = []
for va in insns:
    size = idc.ItemSize(va)
    if idaapi.contains_fixups(va, size):
        # fetch the fixup locations within this one instruction.
        fixups = []
        fixupva = idaapi.get_next_fixup_ea(va)
        fixups.append(fixupva)
        # TODO: assume the fixup size is four bytes, probably bad.
        fixupva += 4

        while fixupva < va + size:
            fixupva = idaapi.get_next_fixup_ea(fixupva)
            fixups.append(fixupva)
            # TODO: assume the fixup size is four bytes, probably bad.
            fixupva += 4

        # assume each fixup is four bytes (TODO!),
        #  and compute the addresses of each component byte.
        fixup_byte_addrs = set([])
        for fixup in fixups:
            for i in range(fixup, fixup+4):
                fixup_byte_addrs.add(i)
        
        # fetch and format each byte of the instruction,
        #  possibly masking it into an unknown byte if its a fixup.
        for i, byte in enumerate(idc.GetManyBytes(va, size)):
            byte_addr = i + va
            if byte_addr in fixup_byte_addrs:
                bytes.append(ord(byte))
                masked_bytes.append('??')
            else:
                bytes.append(ord(byte))
                masked_bytes.append('%02X' % (ord(byte)))
    elif 'call' in idc.GetMnem(va):
        for i, byte in enumerate(idc.GetManyBytes(va, size)):
            bytes.append(ord(byte))
            masked_bytes.append('??')
    else:
        for byte in idc.GetManyBytes(va, size):
            bytes.append(ord(byte))
            masked_bytes.append('%02X' % (ord(byte)))

return Rule('\$0x%x' % (bb.va), bytes, masked_bytes)

def format_rules(fva, rules): ''' given the address of a function, and the byte signatures for basic blocks in the function, format a complete YARA rule that matches all of the basic block signatures. ''' name = idc.GetFunctionName(fva)

# some characters aren't valid for YARA rule names
safe_name = name
BAD_CHARS = '@ /\\!@#$%^&*()[]{};:\'",./<>?'
for c in BAD_CHARS:
    safe_name = safe_name.replace(c, '')

md5 = idautils.GetInputFileMD5()
ret = []
ret.append('rule a_{hash:s}_{name:s} {{'.format(
    hash=md5,
    name=safe_name))
ret.append('  meta:')
ret.append('    sample_md5 = "{md5:s}"'.format(md5=md5))
ret.append('    function_address = "0x{fva:x}"'.format(fva=fva))
ret.append('    function_name = "{name:s}"'.format(name=name))
ret.append('  strings:')
for rule in rules:
    formatted_rule = ' '.join(rule.masked_bytes)
    ret.append('    {name:s} = {{ {hex:s} }}'.format(
        name=rule.name,
        hex=formatted_rule))
ret.append('  condition:')
ret.append('    all of them')
ret.append('}')
return '\n'.join(ret)

def create_yara_rule_for_function(fva): ''' given the address of a function, generate and format a complete YARA rule that matches the basic blocks. ''' rules = [] for bb in get_basic_blocks(fva): rule = get_basic_block_rule(bb)

    # ensure there at least MIN_BB_BYTE_COUNT
    #  non-masked bytes in the rule, or ignore it.
    # this will reduce the incidence of many very small matches.
    unmasked_count = len(filter(lambda b: b != '??', rule.masked_bytes))
    if unmasked_count < MIN_BB_BYTE_COUNT:
        continue

    rules.append(rule)

return format_rules(fva, rules)

def get_segment_buffer(segstart): ''' fetch the bytes of the section that starts at the given address. if the entire section cannot be accessed, try smaller regions until it works. ''' segend = idaapi.getseg(segstart).endEA buf = None segsize = segend - segstart while buf is None: buf = idc.GetManyBytes(segstart, segsize) if buf is None: segsize -= 0x1000 return buf

Segment = namedtuple('Segment', ['start', 'size', 'name', 'buf'])

def get_segments(): ''' fetch the segments in the current executable. ''' for segstart in idautils.Segments(): segend = idaapi.getseg(segstart).endEA segsize = segend - segstart segname = str(idc.SegName(segstart)).rstrip('\x00') segbuf = get_segment_buffer(segstart) yield Segment(segstart, segend, segname, segbuf)

class TestDidntRunError(Exception): pass

def test_yara_rule(rule): ''' try to match the given rule against each segment in the current exectuable. raise TestDidntRunError if its not possible to import the YARA library. return True if there's at least one match, False otherwise. ''' try: import yara except ImportError: logger.warning("can't test rule: failed to import python-yara") raise TestDidntRunError('python-yara not available')

r = yara.compile(source=rule)

for segment in get_segments():
    matches = r.match(data=segment.buf)
    if len(matches) > 0:
        logger.info('generated rule matches section: {:s}'.format(segment.name))
        return True
return False

def main(): va = idc.ScreenEA() fva = get_function(va) print('-' * 80) rule = create_yara_rule_for_function(fva) print(rule)

if test_yara_rule(rule):
    print('success: validated the generated rule')
else:
    print('error: failed to validate generated rule')

if name == 'main': logging.basicConfig(level=logging.DEBUG) logging.getLogger().setLevel(logging.DEBUG) main()