run_exploits.py

#!/usr/bin/env python3


import os
import sys
import subprocess
import pexpect
import traceback
import time
import signal

#TODO:
# This script is a mess...It needs a serious cleanup and probably a more effective way to run the
# the exploits from ExploitDB and the TriforceAFL bugs

exploits_local = ["IPv4_TCP-CVE-2010-4165","Netfilter-CVE-2016-3135","Scsi-CVE-2017-14489","bcm-CVE-2010-2959","sctp-CVE-2013-1828","Ip_tables-CVE-2016-3134","Netfilter-CVE-2016-4997","tun-CVE-2009-1897","l2tp-CVE-2014-4943","vfat-CVE-2007-2878","Xfrm-CVE-2017-16939"]

exploits_remote = ["netusb-CVE-2015-3036","sctp-CVE-2010-1173"]

afl_bugs = {
            "art_1" : ["/dev/dk", "c 63 0"],
            "edinvram2_1" : ["/dev/edienvram", "c 247 0"],
            "edinvram2_2" : ["/dev/edienvram", "c 247 0"],
            "rt_rdm_1"    : ["/dev/rdm0", "c 253 0"],
            "acos_nat_1"  : ["/dev/acos_nat_cli", "c 100 0"],
            "acos_nat_2"  : ["/dev/acos_nat_cli", "c 100 0"],
            "acos_nat_3"  : ["/dev/acos_nat_cli", "c 100 0"],
            "gpio_1"      : ["/dev/gpio", "c 101 0"],
            "gpio_2"      : ["/dev/gpio-hermon", "c 101 0"],
            "smcdrv_1"    : ["/dev/smcdrv", "c 70 0"],
            "rt_rdm_2"    : ["/dev/rdm0", "c 253 0"],
            "ipv6_spi_1"   : ["/dev/spiv6", "c 101 0"],
            "ipv6_spi_2"   : ["/dev/spiv6", "c 101 0"],
            "IDP_1"       : ["/dev/detector", "c 190 0"],
            "IDP_2"       : ["/dev/detector", "c 190 0"],
            "IDP_3"       : ["/dev/detector", "c 190 0"],
            "ufilter_1" : ["/dev/ufilter", "c 254 0"],
            "ufilter_2" : ["/dev/ufilter", "c 254 0"],
            "ufilter_3" : ["/dev/ufilter", "c 254 0"],
            "ufilter_4" : ["/dev/ufilter", "c 254 0"],
            "ufilter_5" : ["/dev/ufilter", "c 254 0"],
        }

def clean_buffer(child):
    child.buffer = b''
    child.before = b''
    return child
    #if child.before:
        #child.expect (r'.+')a


class Image:
    def __init__(self,image,scratch_dir,arch,kernel):
        self.image = image
        self.scratch = scratch_dir
        self.arch = arch
        self.kernel = kernel
        self.segfaults = []
        self.log_err = "{0}/{1}/log_errs.out".format(self.scratch,self.image)
        self.create_dirs()
        self.prompt = ["~","#","$","/"]

    def write_error(self,error):
        with open(self.log_err,"a") as f:
            f.write(str(error)+"\n\n")
    
    def create_dirs(self):
        local_dir = "{0}/{1}/local".format(self.scratch,self.image)
        remote_dir = "{0}/{1}/remote".format(self.scratch,self.image)
        afl_dir = "{0}/{1}/afl".format(self.scratch,self.image)
        
        cmd1= "mkdir {0}".format(local_dir)
        cmd2= "mkdir {0}".format(remote_dir)
        cmd3= "mkdir {0}".format(afl_dir)
        try:
            res = subprocess.call(cmd1,shell=True)
        except:
            print("Local log dir for",self.image,"already exists")

        try:
            res = subprocess.call(cmd2,shell=True)
        except:
            print("Remote log dir for",self.image,"already exists")
        
        try:
            res = subprocess.call(cmd3,shell=True)
        except:
            print("AFL log dir for",self.image,"already exists")


    def stop_child(self,child,qemu):

        try:
            if child.isalive():
                print("Killing forcefully the child with pid",str(child.pid))
                cmd = "kill -9 "+ str(child.pid)
                os.system(cmd)
            else:
                print('Child exited gracefully.')
        except:
            pass
        
        try:
            os.killpg(os.getpgid(qemu.pid), signal.SIGTERM)
        except:
            print("Could not stop qemu for image",self.image)
	
        res = ""

        try:
            pid_to_kill = 'ps ax | grep "qemu-system-"'
            res = subprocess.check_output(pid_to_kill,shell=True).decode("utf-8")
        except Exception as e:
            print(e)
        
        pid = None
        
        results = res.split("\n")
        
        if results == ['']:
            return

        for rs in results:
            if rs == "":
                continue
            if "grep" not in rs: 
                pid = int(rs.split()[0])
                print("Killing pid",pid)
                break
        try:
            if pid != None:
                os.kill(pid, signal.SIGINT)
        except Exception as e:
            print(e)
        
        time.sleep(2)
        try:
            pid_to_kill = 'tunctl -d tap{}_0'.format(self.image)
            res = subprocess.check_output(pid_to_kill,shell=True).decode("utf-8")
        except Exception as e:
            print(e)
        
    
    def get_child_ip(self):
        run_file = self.scratch + "/" + self.image + "/run.sh"
        with open(run_file,"r") as f:
            lines = f.readlines()
        
        ip_addresses = []
        ip_addr = ""
        for line in lines:
            if "ip route add" in line:
                ip_addr = line.split()[4].split("/")[0]
                ip_addresses.append(ip_addr)
                print("The IP address of target {0} is".format(self.image),ip_addr)
        
        return ip_addresses
    
    def save_log_file(self,exploit_num,tp, *args):
        logfile = "{0}/{1}/qemu.final.serial.log".format(self.scratch,self.image)
        if tp == "local":
            dest = "{0}/{1}/local/qemu.final.serial.log{2}".format(self.scratch,self.image,exploit_num)
            cmd = "mv {0} {1}".format(logfile,dest)
        elif tp == "remote":
            dest = "{0}/{1}/remote/qemu.final.serial.log{2}".format(self.scratch,self.image,exploit_num)
            exploit_out = "{0}/{1}/remote/exploit_out{2}".format(self.scratch,self.image,exploit_num)
            cmd = "mv {0} {1}".format(logfile,dest)
        else:
            dest = "{0}/{1}/afl/qemu.final.serial.log{2}".format(self.scratch,self.image,exploit_num)
            cmd = "mv {0} {1}".format(logfile,dest)

        try:
            res = subprocess.call(cmd,shell=True)
        except:
            print("Could not copy the logfile for the exploit",tp,exploit_num)
        
        if args != [] and args != None:
            try:
                with open(exploit_out, "w") as f:
                    for line in args[0]:
                        f.write(line + "\n")
            except:
                pass
                #print(traceback.format_exc())
                #print("Could not copy the exploit output",tp,exploit_num)


    def truncate_output(self):
        logfile = "{0}/{1}/qemu.final.serial.log".format(self.scratch,self.image)
        #trunc = ":> {0}".format(logfile)
        trunc = "echo \" \" > {0}".format(logfile)
        try:
            res = subprocess.call(trunc,shell=True)
        except:
            print("Something happened when truncating the log file")

    def run_the_exploit(self,child,extension,exploit_num,qemu):
        explt_name = exploits_local[exploit_num] + "/exploit" + extension
        exploit = "/firmsolo_analysis/local/" + explt_name
        
        ### For kernel panics

        echo = 'echo "Running exploit {0} for image {1}"\n\n\n\n\n\n\n'.format(exploit,self.image)
        print(echo)
        try:
            try:
                child.read_nonblocking(16834, timeout = 10)
            except:
                pass
            child.sendline(echo)
            child.expect_exact(self.prompt,timeout=30)
            output = child.before.decode("utf-8")
        except:
            print("An error happened when running echo for exploit",explt_name)
            self.write_error(traceback.format_exc())
        
        if exploit_num == 7:
            try:
                try:
                    child.read_nonblocking(16834, timeout = 10)
                except:
                    pass
                child.sendline("mkdir /dev/net")
                child.expect_exact(self.prompt,timeout=30)
                output = child.before.decode("utf-8")
                child.sendline("mknod /dev/net/tun c 10 200")
                child.expect_exact(self.prompt,timeout=30)
                output = child.before.decode("utf-8")
            except:
                print("An error happened when running echo for exploit",explt_name)
                self.write_error(traceback.format_exc())

        try:
            try:
                child.read_nonblocking(16834, timeout = 10)
            except:
                pass
            child.sendline(exploit)
            child.expect_exact(self.prompt,timeout=30)
            output = child.before.decode("utf-8")
        except:
            print("An error happened when running exploit",explt_name)
            self.write_error(traceback.format_exc())
            self.stop_child(child,qemu)

        time.sleep(10)
        
        self.save_log_file(exploit_num,"local")
        return "Done"
    
    def check_local_exploits(self):
        try:
            child,qemu = self.spawn_image()
            if child == None:
                return
        except:
            print("Something went wrong when spawning the image for local exploits")
            self.stop_child(child,qemu)
            self.write_error(traceback.format_exc())
        
        extension = self.find_extension()

        for num in range(11):
            outcome = self.run_the_exploit(child,extension,num,qemu)
            time.sleep(30)

            if "Done" in outcome:
                self.stop_child(child,qemu)
                time.sleep(2)
                if num == 10:
                    continue
                child,qemu = self.spawn_image()
            else:
                continue
        
        self.stop_child(child,qemu)
        
    def run_the_remote_exploit(self,ip,exploit_num,child,qemu):
        cwd = os.getcwd()
        
        if self.arch == "mipsel":
            explt_name = exploits_remote[exploit_num] + "/netusb-2015-3036_2.py"
        elif self.arch == "mipseb":
            explt_name = exploits_remote[exploit_num] + "/netusb-2015-3036_2.py"
        else:
            return "Done"

        exploit = "python2.7 {}/firmsolo_analysis/remote/".format(cwd) + explt_name + " {0} {1}".format(ip,"20005") 
        
        print("Waiting till the image finishes loading the NetUSB module")
        time.sleep(230)
        echo = '"Running exploit {0} for image {1}"\n'.format(exploit,self.image)
        print(echo)
        
        try:
            res = subprocess.check_output(exploit,shell=True).decode("utf-8").split("\n")
            print(res)
        except:
            print(traceback.format_exc())
            self.stop_child(child,qemu)
        
        time.sleep(2)
        self.save_log_file(exploit_num,"remote", res)

        return "Done"


    def check_remote_exploits(self):
        child = None
        try:
            qemu = self.spawn_qemu()
            if qemu == None:
                return
        except:
            print("Something went wrong when spawning the image for remote exploits")
            self.stop_child(child,qemu)
        
        ip_addresses = self.get_child_ip()
        
        for ip in ip_addresses:
            outcome = self.run_the_remote_exploit(ip,0,child,qemu)
            self.stop_child(child,qemu)
            time.sleep(3)

        self.stop_child(child,qemu)

    def run_afl_bugs(self, child, mknod, cmd, rm, qemu, bug):

        echo = 'echo "Running afl_bug {0} for image {1}"\n\n\n\n\n\n\n'.format(cmd,self.image)
        print(echo)
        try:
            try:
                child.read_nonblocking(16834, timeout = 10)
            except:
                pass
            child.sendline(echo)
            child.expect_exact(self.prompt,timeout=30)
            output = child.before.decode("utf-8")
            print(output)
        except:
            print("An error happened when running echo for afl_test", cmd)
            self.write_error(traceback.format_exc())

        try:
            ### Clean the buffer
            try:
                child.read_nonblocking(16834, timeout = 10)
            except:
                print(traceback.format_exc())
                pass
            print("Removing existing char device", rm)
            child.sendline(rm)
            child.expect_exact(self.prompt,timeout=30)
            
            print("Creating char device", mknod)
            child.sendline(mknod)
            child.expect_exact(self.prompt,timeout=30)

            child.sendline(cmd)
            child.expect_exact(self.prompt,timeout=45)
            time.sleep(20)
            try:
                output = child.read_nonblocking(4096, timeout = 10)
            except:
                print(traceback.format_exc())
                pass

            print("Pexpect output", output)
        except:
            print("An error happened when running exploit", cmd)
            self.write_error(traceback.format_exc())
            self.stop_child(child,qemu)

        self.save_log_file("_" + bug,"afl")
        return "Done"


    def check_afl_exploits(self, modules):
        fuzz_test_bin = self.get_fuzz_test_path()
        mod_exists = False
        for indx,bug in enumerate(afl_bugs):
            tokens = bug.split("_")

            module = "_".join(tokens[1:-1])
            if module == "ufilter":
                module = "u_filter"

            if module not in modules:
                continue
            else:
                mod_exists = True
                break
        if not mod_exists:
            return

        print("Before running the AFL exploits")
        try:
            child,qemu = self.spawn_image()
            if child == None:
                print("The child in AFL experiments is None")
                return
        except:
            print("Something went wrong when spawning the image for local exploits")
            self.stop_child(child,qemu)
            self.write_error(traceback.format_exc())

        print("Before getting extension")
        extension = self.find_extension()

        for indx,bug in enumerate(afl_bugs):
            tokens = bug.split("_")
            module = "_".join(tokens[1:-1])

            if module == "ufilter":
                module = "u_filter"

            if module not in modules:
                continue
            
            print("Before sleeping...")
            time.sleep(60)
            rm = "rm {}".format(afl_bugs[bug][0])
            mknod = "mknod {} {}".format(afl_bugs[bug][0], afl_bugs[bug][1])
            cmd = "{} {} /firmsolo_analysis/triforceafl_bugs/{}".format(fuzz_test_bin, afl_bugs[bug][0], bug)
            print("Sending command", cmd)
            try:
                outcome = self.run_afl_bugs(child, mknod, cmd, rm, qemu, bug)
            except:
                self.stop_child(child,qemu)
                outcome = ""

            if "Done" in outcome:
                self.stop_child(child,qemu)
                time.sleep(2)
                child,qemu = self.spawn_image()
            else:
                continue
        
        self.stop_child(child,qemu)
        

    def find_extension(self):

        if self.kernel < "linux-2.6.32":
            if self.arch == "mipsel":
                vers = "_deb_mipsel"
            elif self.arch == "mipseb":
                vers = "_deb_mipseb"
            else:
                vers = "_armel_old"
        else:
            if self.arch == "mipsel":
                vers = "_ub_mipsel"
            elif self.arch == "mipseb":
                vers = "_ub_mipseb"
            else:
                vers = "_armel"
        
        return vers

    def get_fuzz_test_path(self):

        if self.kernel < "linux-2.6.32":
            if self.arch == "mipsel":
                path = "/firmsolo_analysis/fuzzer_binaries/mips_old/driver_mips_le"
            elif self.arch == "mipseb":
                path = "/firmsolo_analysis/fuzzer_binaries/mips_old/driver_mips_be"
            else:
                path = "/firmsolo_analysis/fuzzer_binaries/arm_old/driver_arm_old"
        else:
            if self.arch == "mipsel":
                path = "/firmsolo_analysis/fuzzer_binaries/mips_new/driver_mips_le"
            elif self.arch == "mipseb":
                path = "/firmsolo_analysis/fuzzer_binaries/mips_new/driver_mips_le"
            else:
                path = "/firmsolo_analysis/fuzzer_binaries/arm_new/driver_arm_new"

        return path

    def spawn_qemu(self):
        cmd = "{0}/{1}/run.sh".format(self.scratch, self.image)
        
        qemu = None
        try:
            qemu = subprocess.Popen(cmd,preexec_fn=os.setsid)
        except:
            print("Something went wrong when spawning qemu")
            self.write_error(traceback.format_exc())

        return qemu

    def spawn_image(self):
        qemu = self.spawn_qemu()
        
        time.sleep(20)
        child = None
        print("Running pexpect for image {0}".format(self.image))
        cmd = "socat - UNIX-CONNECT:/tmp/qemu.{}.S1".format(self.image)
        try:
            child = pexpect.spawn(cmd)
            child.expect(self.prompt,timeout=10)
        except:
            print("Image failed to spawn:",self.image)
            self.write_error(traceback.format_exc())
            self.stop_child(child,qemu)
            return None, qemu
            
        time.sleep(2)
    
        print("Sending enter to image",self.image)
        try:
            child.sendline("\n")
            child.expect_exact(self.prompt,timeout=30)
        except:
            print(traceback.format_exc())
            self.write_error(traceback.format_exc())
            self.stop_child(child,qemu)
            print("Wrong prompt for image",self.image)
            return None, qemu
        
        return child,qemu

def get_image_data(scratch_dir,image):

    with open(scratch_dir + "/{0}/image_data".format(image),"r") as f:
        vers = f.readline().strip("\n")
        arch = f.readline().strip("\n")
        kernel = f.readline().strip("\n")

    return vers,arch,kernel


def get_loaded_modules(image):
    pkl = f"{cu.loaded_mods_path}{image}/{image}_ups_subs.pkl"
    
    mods = []

    try:
        data = cu.read_pickle(pkl)
        loaded_modules = data[0]
        mods = list(map(lambda x:x.split("/")[-1],loaded_modules))
    except:
        print(traceback.format_exc())
        pass

    return mods


if __name__ == "__main__":
    
    image = sys.argv[1]
    scratch_dir = sys.argv[2]
    firmsolo_dir = sys.argv[3]
    
    sys.path.append(firmsolo_dir)
    globals()["cu"] = __import__("custom_utils")
    
    vers, arch, kernel = get_image_data(scratch_dir,image)
    obj = Image(image,scratch_dir,arch,kernel)

    modules = get_loaded_modules(image)

    if not modules:
        sys.exit(0)
    
    print("Remote Exploits")
    if "NetUSB" in modules:
        try:
            obj.check_remote_exploits()
        except:
            print("Something went wrong when running the remote exploits")
            print(traceback.format_exc())
    
    time.sleep(1)
    
    print("Local exploits")
    try:
        obj.check_local_exploits()
    except:
        print("Something went wrong when running the local exploits")
        print(traceback.format_exc())
    
    time.sleep(2)
    print("afl bugs")
    try:
        obj.check_afl_exploits(modules)
    except:
        print("Something went wrong when running the afl exploits")
        print(traceback.format_exc())