run_yextend

#!/usr/bin/env python

"""
    run statement:

    ./run_yextend -r rule_entity -t target_file_entity [-j]

    I label them as 'entity' because they can each be
    either a directory or a file
"""
import subprocess
import sys
import os
import fnmatch
import time
import getopt
import json
from types import *
from collections import OrderedDict
########################################################
def usage():
    print "\n{}\n".format("Usage:")

    print "{} {} {} {}\n".format("./run_yextend", "-r rule_entity", "-t target_file_entity", "[-j]")

    print "    {}".format("-r RULES_FILE = Yara ruleset file [*required]")
    print "    {}".format("-t TARGET_ENTITY = file or directory [*required]")
    print "    {}".format("-j output in JSON format and nothing more [optional]")

    print "".format("\n\n")


def get_file_list(thedir='', thepattern='*'):

    if not thedir.endswith('/'):
        thedir = thedir + '/'
    f_list = []
    if thedir:
        for root, dirnames, filenames in os.walk(thedir):
            for name in fnmatch.filter(filenames, thepattern):
                complete_fname = "{}{}".format(root, name)
                if os.path.isfile(complete_fname):
                    if complete_fname not in f_list:
                        f_list.append(complete_fname)
    return f_list


def my_print(thestr='', thecolor=''):
    if thestr:
        format_str = '\033[1;{}m{}\033[1;m'
        if thecolor == 'gray':
            return format_str.format('30', thestr)
        elif thecolor == 'red':
            return format_str.format('31', thestr)
        elif thecolor == 'green':
            return format_str.format('32', thestr)
        elif thecolor == 'yellow':
            return format_str.format('33', thestr)
        elif thecolor == 'blue':
            return format_str.format('34', thestr)
        elif thecolor == 'magenta':
            return format_str.format('35', thestr)
        elif thecolor == 'cyan':
            return format_str.format('36', thestr)
        elif thecolor == 'white':
            return format_str.format('37', thestr)
        elif thecolor == 'crimson':
            return format_str.format('38', thestr)
        elif thecolor == 'highlighted_red':
            return format_str.format('41', thestr)
        elif thecolor == 'highlighted_red':
            return format_str.format('42', thestr)
        elif thecolor == 'highlighted_brown':
            return format_str.format('43', thestr)
        elif thecolor == 'highlighted_blue':
            return format_str.format('44', thestr)
        elif thecolor == 'highlighted_magenta':
            return format_str.format('45', thestr)
        elif thecolor == 'highlighted_cyan':
            return format_str.format('46', thestr)
        elif thecolor == 'highlighted_gray':
            return format_str.format('47', thestr)
        elif thecolor == 'highlighted_crimson':
            return format_str.format('48', thestr)


def which(program=""):
    def is_exe(fpath):
        return os.path.exists(fpath) and os.access(fpath, os.X_OK)

    def ext_candidates(fpath):
        yield fpath
        for ext in os.environ.get("PATHEXT", "").split(os.pathsep):
            yield fpath + ext

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        tarr = os.environ["PATH"].split(os.pathsep)
        tarr.append("/sbin")
        tarr.append("./")
        for path in tarr:
            exe_file = os.path.join(path, program)
            for candidate in ext_candidates(exe_file):
                if is_exe(candidate):
                    return candidate
    return None


def print_yextend_output(out=''):
    if out:
        for the_out in out.split('\n'):
            if the_out.startswith('='):
                print my_print(thestr=the_out, thecolor='white')
            elif the_out.startswith("Yara "):
                # print the_out
                t_o = the_out.split(': ')
                print "{}: ".format(my_print(thestr=t_o[0], thecolor='white'))
                for rule_data in t_o[1].split(', '):
                    if ':[' in rule_data:
                        rule_id, rule_meta = rule_data.split(':[')
                        rule_meta = rule_meta[:-1]
                        print "\tRule ID: {}".format(my_print(thestr=rule_id, thecolor='red'))
                        print "\tRule META:"
                        for rmeta in rule_meta.split(','):
                            rmeta = rmeta.replace('=', ' = ')
                            if rmeta.startswith('detected'):
                                rm_label, rm_dat = rmeta.split(' = ')
                                print "\t\t{}".format(my_print(thestr=rm_label + ':', thecolor='blue'))
                                rm_dat_spl = rm_dat.split('-')

                                for rmdspl in rm_dat_spl:
                                    try:
                                        rmdspl_offset, rmdspl_label = rmdspl.strip().split(':')
                                    except ValueError:
                                        pass
                                    rmdspl_out = "{} at {}".format(my_print(thestr=rmdspl_label, thecolor='red'), my_print(rmdspl_offset, thecolor='red'))
                                    print "\t\t\t{}".format(my_print(thestr=rmdspl_out, thecolor='blue'))
                            else:
                                print "\t\t{}".format(my_print(thestr=rmeta, thecolor='blue'))
                    else:
                        if rule_data == "No hit from Yara ruleset":
                            print "\t{}".format(my_print(thestr=rule_data, thecolor='red'))
                        else:
                            print "\tRule ID: {}".format(my_print(thestr=rule_data, thecolor='red'))
                    print
            elif the_out == '\n':
                print the_out
            else:
                try:
                    t_o = the_out.split(': ')
                    print "{}: {}".format(my_print(thestr=t_o[0], thecolor='white'), my_print(thestr=t_o[1], thecolor='green'))
                except:
                    def_out = my_print(thestr=the_out, thecolor='white')
                    if def_out:
                        print my_print(thestr=the_out, thecolor='white')
                    else:
                        print the_out


def print_yextend_json_output(out=''):

    if out:

        try:
            string_to_json = json.loads(out)

            for rule_file in string_to_json:

                print '\n{}'.format(my_print(thestr='===============================ALPHA===================================', thecolor='white'))
                print '{} : {}'.format(my_print(thestr="File Name", thecolor='white'), my_print(thestr=rule_file['file_name'], thecolor='green'))
                print '{} : {}'.format(my_print(thestr="File Size", thecolor='white'), my_print(thestr=rule_file['file_size'], thecolor='green'))
                print '{} : {}'.format(my_print(thestr="File Signature (MD5)", thecolor='white'), my_print(thestr=rule_file['file_signature_MD5'], thecolor='green'))

                if 'yara_ruleset_file_name' in rule_file.keys():
                    print '{} : {}'.format(my_print(thestr="Yara Ruleset File Name", thecolor='white'), my_print(thestr=rule_file['yara_ruleset_file_name'], thecolor='green'))

                print '\n{}\n'.format(my_print(thestr='=======================================================================', thecolor='white'))
                print my_print(thestr='Yara Result(s):', thecolor='white')

                if ('scan_results' in rule_file) and (rule_file['scan_results'] is not None):

                    for result in rule_file['scan_results']:
                        res_keys = result.keys()
                        if result['yara_matches_found']:

                            print '\tRule ID: {}'.format(my_print(thestr=result['yara_rule_id'], thecolor='red'))
                            print '\tRule META:'

                            for key, value in result.iteritems():
                                if key == 'detected offsets':
                                    constructed_out = key + ':'
                                    print '\t\t{}'.format(my_print(thestr=constructed_out, thecolor='blue'))
                                    for detect in value:
                                        if (":") in detect:
                                            offset,label = detect.split(':')
                                        else:
                                            offset = detect
                                            label = "UNK"
                                        constructed_out = '{} at {}'.format(my_print(thestr=label, thecolor='red'), my_print(offset, thecolor='red'))
                                        print '\t\t\t{}'.format(my_print(thestr=constructed_out, thecolor='blue'))
                                else:
                                    if key not in ['scan_type', 'file_name', 'file_signature_MD5', 'yara_rule_id', 'yara_matches_found']:
                                        constructed_pair = key + ' = ' + str(value)
                                        print '\t\t{}'.format(my_print(thestr=constructed_pair, thecolor='blue'))

                            if len(result) <= 2:
                                print '\t\t{}\n'.format('None')

                            if 'scan_type' in res_keys:
                                print '\n{} : {}'.format(my_print(thestr="Scan Type", thecolor='white'), my_print(thestr=result['scan_type'], thecolor='green'))

                            if 'file_name' in res_keys:
                                print '{} : {}'.format(my_print(thestr="File Name", thecolor='white'), my_print(thestr=result['file_name'], thecolor='green'))

                            if 'child_file_name' in res_keys:
                                print '{} : {}'.format(my_print(thestr="Child File Name", thecolor='white'), my_print(thestr=result['child_file_name'], thecolor='green'))

                            if 'parent_file_name' in res_keys:
                                print '{} : {}'.format(my_print(thestr="Parent File Name", thecolor='white'), my_print(thestr=result['parent_file_name'], thecolor='green'))

                            if 'file_size' in res_keys:
                                print '{} : {}'.format(my_print(thestr="File Size", thecolor='white'), my_print(thestr=result['file_size'], thecolor='green'))

                            if 'file_signature_MD5' in res_keys:
                                print '{} : {}\n'.format(my_print(thestr="File Signature (MD5)", thecolor='white'), my_print(thestr=result['file_signature_MD5'], thecolor='green'))

                        else:

                            print '\tRule ID: {}\n'.format('None')

                            if 'scan_type' in res_keys:
                                print '{} : {}'.format(my_print(thestr="Scan Type", thecolor='white'), my_print(thestr=result['scan_type'], thecolor='green'))

                            if 'non_archive_file_name' in res_keys:
                                print '{} : {}'.format(my_print(thestr="Non-Archive File Name", thecolor='white'), my_print(thestr=result['non_archive_file_name'], thecolor='green'))

                            if 'file_name' in res_keys:
                                print '{} : {}'.format(my_print(thestr="File Name", thecolor='white'), my_print(thestr=result['file_name'], thecolor='green'))

                            if 'parent_file_name' in res_keys:
                                print '{} : {}'.format(my_print(thestr="Parent File Name", thecolor='white'), my_print(thestr=result['parent_file_name'], thecolor='green'))

                            if 'child_file_name' in res_keys:
                                print '{} : {}'.format(my_print(thestr="Child File Name", thecolor='white'), my_print(thestr=result['child_file_name'], thecolor='green'))

                            if 'file_size' in res_keys:
                                print '{} : {}'.format(my_print(thestr="File Size", thecolor='white'), my_print(thestr=result['file_size'], thecolor='green'))

                            if 'file_signature_MD5' in res_keys:
                                print '{} : {}\n'.format(my_print(thestr="File Signature (MD5)", thecolor='white'), my_print(thestr=result['file_signature_MD5'], thecolor='green'))

                print '{}\n'.format(my_print(thestr='===============================OMEGA===================================', thecolor='white'))

        except:

            # when file can't be read
            t_o = out.split(': ')
            print '{}: {}'.format(my_print(thestr=t_o[0], thecolor='white'), my_print(thestr=t_o[1].strip("null"), thecolor='green'))


def format_child(child):
    
    if "yara_rule_id" in child:
        rule_id = child["yara_rule_id"]
        try:
            yara_results = {rule_id: {"description":""}}
            if "description" in child:
                yara_results[rule_id]["description"] = child["description"]
                del child["description"]
            if "detected offsets" in child:
                yara_results[rule_id]["offsets"] = child["detected offsets"]
                del child["detected offsets"]
            if "hit_count" in child:
                yara_results[rule_id]["hit_count"] = child["hit_count"]
                del child["hit_count"]
            del child["yara_rule_id"]
            child["yara_results"] = yara_results
        except Exception, e:
            print e
    else:
        # clean out empty k/v pairs
        for k,v in child.items():
            if v == "":
                del child[k]

    return child


def advanced_print(out="", ruleset_name=""):

    scan_results_key = 'scan_results'
    out = json.loads(out)
    r_final = []
    j_final = []
    files_level1 = {}
    
    for file_dict in out:
        
        files = []
        new_files = dict(file_dict)
        del new_files[scan_results_key]
        files.append(format_child(new_files))
        
        for child in file_dict[scan_results_key]:
            
            child_keys = child.keys()
            new_child = dict(child)
            if "child_file_name" in child_keys:
            
                new_child["file_name"] = new_child.pop("child_file_name")
                child["file_name"] = child.pop("child_file_name")

                for f in reversed(files):
                    f_keys = f.keys()
                    if "file_name" in f_keys:
                        if f["file_name"] == child["parent_file_name"]:
                            try:
                                f["children"].append(format_child(new_child))
                            except:
                                f["children"] = [format_child(new_child)]
                

                del child["parent_file_name"]
                del new_child["parent_file_name"]

                files.append(format_child(child))

            elif "non_archive_file_name" in child_keys:
            
                new_child["file_name"] = new_child.pop("non_archive_file_name")
                try:
                    files[0]["children"].append(format_child(new_child))
                except:
                    files[0]["children"] = [format_child(new_child)]
                    
            else:
            
                try:
                    files[0]["children"].append(format_child(new_child))
                except:
                    files[0]["children"] = [format_child(new_child)]

        for x in files:
            if "children" in x:
                if x not in j_final:
                    j_final.append(x)
                    #r_final.append(x)
        
        if len(j_final) > 0:
            files_level1.update(j_final[0])
            files_level1[scan_results_key] = []
            cnt = 0
            for i in j_final:
                if cnt > 0:
                    files_level1[scan_results_key].append(j_final[cnt])
                cnt = cnt + 1

            r_final.append(files_level1)

    return r_final
########################################################

RULES_ENTITY = ''
FILE_ENTITY = ''
RULES_ENTITY_DIR = False
YEXTEND = which(program="yextend")

USRLOCLIB = "{}".format("\/usr\/local\/lib")
try:
    proc = subprocess.Popen([YEXTEND],
                            env={LD_LIBRARY_PATH:USRLOCLIB},
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            )
    out, err = proc.communicate()
except:
    USRLOCLIB = "{}".format("/usr/local/lib")
    
LD_LIBRARY_PATH = "{}".format("LD_LIBRARY_PATH")

JSON_OUT = False
ADVANCED = False
try:
    options, rem = getopt.getopt(sys.argv[1:], 'r:t:ja', ['ruleset=','target='])
except getopt.GetoptError as err:
    print str(err)
    usage()
    sys.exit()

for opt, arg in options:
    if opt in ('-r', '--ruleset'):
        RULES_ENTITY = "{}".format(arg)
    elif opt in ('-t', '--target'):
        FILE_ENTITY = "{}".format(arg)
    elif opt in ('-j'):
        JSON_OUT = True
    elif opt in ('-a'):
        ADVANCED = True
    else:
        assert False, "unhandled option"

if os.path.isdir(RULES_ENTITY):
    RULES_ENTITY_DIR = True

if not RULES_ENTITY or not FILE_ENTITY:
    usage()
    sys.exit()

'''
    we will force the JSON option to yextend
    as the default run mode going forward

    function 'print_yextend_json_output' presents
    the results of a run to stdout the same way
    we always have except that under the hood it
    is transforming the JSON output from yextend
    instead of transforming the legacy output
    format (i.e. non-JSON)
'''
JSON_OUT_FORCE = "-j"
if JSON_OUT:
    JSON_OUT = "-j"
else:
    JSON_OUT = ""


ruleset_pool = []    
''' ruleset = file '''
if not RULES_ENTITY_DIR:
    ruleset_pool.append(RULES_ENTITY)
else:
    for r in get_file_list(thedir=RULES_ENTITY):
        ruleset_pool.append(r)

    
for lpol in ruleset_pool:

    out = ""
    err = ""
    proc = subprocess.Popen([YEXTEND, "-r", lpol, "-t", FILE_ENTITY, JSON_OUT_FORCE],
                            env={LD_LIBRARY_PATH:USRLOCLIB},
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            )
    out, err = proc.communicate()

    if out:
        if JSON_OUT:
            if ADVANCED:
                print json.dumps(advanced_print(out=out, ruleset_name=lpol), indent=4)
            else:
                #print_yextend_json_output(out=out)
                print out
        else:
            #print_yextend_output(out=out)
            print_yextend_json_output(out=out)
    
    if err:
        print "ERR: %s" % err