ExploitGenerator.cpp

// Copyright 2021-2022 Software Quality Laboratory, NYCU.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <s2e/Plugins/CRAX/CRAX.h>
#include <s2e/Plugins/CRAX/Expr/BinaryExprEval.h>
#include <s2e/Plugins/CRAX/Utils/StringUtil.h>

#include <cassert>
#include <fstream>

#include "ExploitGenerator.h"

namespace s2e::plugins::crax {

ExploitGenerator::ExploitGenerator()
    : m_state(),
      m_ropGadgetResolver(),
      m_ropPayloadBuilder(),
      m_coreGenerator() {}


void ExploitGenerator::run(S2EExecutionState *state) {
    std::vector<RopPayload> ropPayload;
    m_state = state;

    if (!checkRequirements()) {
        return;
    }

    initialize();

    if (g_crax->getExploitForm() == CRAX::ExploitForm::SCRIPT) {
        ropPayload = buildFullRopPayload();
        generateExploitScript(ropPayload);
    } else {
        ropPayload = buildStage1RopPayload();
        generateExploit(RopPayloadBuilder::getStage1Payload(ropPayload));
    }
}

bool ExploitGenerator::checkRequirements() const {
    for (auto m : g_crax->getModules()) {
        if (!m->checkRequirements()) {
            log<WARN>() << "Requirements unmet (Module: " << m->toString() << ")\n";
            return false;
        }
    }

    for (auto t : g_crax->getTechniques()) {
        if (!t->checkRequirements()) {
            log<WARN>() << "Requirements unmet (Technique: " << t->toString() << ")\n";
            return false;
        }
    }

    return true;
}

void ExploitGenerator::initialize() {
    // Maybe initialize the core generator.
    // XXX: What if more than one module has a core generator?
    if (!m_coreGenerator) {
        for (auto m : g_crax->getModules()) {
            if (auto coreGenerator = m->makeCoreGenerator()) {
                log<WARN>() << "Using core generator from module: " << m->toString() << '\n';
                m_coreGenerator = std::move(coreGenerator);
                break;
            }
        }
        // No module has a core generator, so we'll use the default one.
        if (!m_coreGenerator) {
            m_coreGenerator = std::make_unique<CoreGenerator>();
        }
    }

    g_crax->getExploit().reset();
    m_ropPayloadBuilder.reset();

    for (auto t : g_crax->getTechniques()) {
        log<INFO>() << "Initializing technique: " << t->toString() << '\n';
        t->initialize();
    }
}

std::vector<RopPayload> ExploitGenerator::buildFullRopPayload() {
    for (auto t : g_crax->getTechniques()) {
        if (!m_ropPayloadBuilder.chain(*t)) {
            return {};
        }
    }

    return m_ropPayloadBuilder.build();
}

std::vector<RopPayload> ExploitGenerator::buildStage1RopPayload() {
    for (auto t : g_crax->getTechniques()) {
        if (!m_ropPayloadBuilder.chain(*t)) {
            return {};
        }
        if (!m_ropPayloadBuilder.isSymbolicMode()) {
            break;
        }
    }

    return m_ropPayloadBuilder.build();
}

bool ExploitGenerator::generateExploitScript(const std::vector<RopPayload> &ropPayload) const {
    if (ropPayload.empty()) {
        return false;
    }

    Exploit &exploit = g_crax->getExploit();
    const ELF &elf = exploit.getElf();
    const ELF &libc = exploit.getLibc();
    const std::string &elfPrefix = elf.getVarPrefix();
    const std::string &libcPrefix = libc.getVarPrefix();

    // Write exploit shebang.
    exploit.writeline(Exploit::s_shebang);

    // Pwntools stuff.
    exploit.writelines({
        "from pwn import *",
        "context.update(arch = 'amd64', os = 'linux', log_level = 'info')",
        "",
        format("%s = ELF('%s', checksec=False)", elfPrefix.c_str(), elf.getFilename().c_str()),
        format("%s = ELF('%s', checksec=False)", libcPrefix.c_str(), libc.getFilename().c_str()),
        ""
    });

    exploit.registerSymbol("canary", 0);

    // XXX: add support for shared libraries other than libc.so.6
    exploit.registerSymbol(elfPrefix + "_base", 0);
    exploit.registerSymbol(libcPrefix + "_base", 0);

    // Declare symbols and values.
    for (const auto &[name, value] : exploit.getSymtab()) {
        exploit.writeline(format("%s = 0x%llx", name.c_str(), value));
    }

    exploit.writeline();

    // Define solve_stage1() function.
    if (elf.checksec.hasCanary || elf.checksec.hasPIE) {
        exploit.writeline("def solve_stage1(canary, elf_base, iostates) -> bytes:");
        exploit.setIndentLevel(4);
        exploit.writelines({
            "os.system('./launch-crax.sh -c {} -e {} -s {}'"
            ".format(hex(canary), hex(elf_base), iostates))",
            "with open('stage1.bin', 'rb') as f:",
            "    return f.read()",
        });
        exploit.setIndentLevel(0);
        exploit.writeline();
    }

    // Generate the main function.
    exploit.writeline("if __name__ == '__main__':");
    exploit.setIndentLevel(4);

    // Generate the main logic of the exploit script.
    m_coreGenerator->generateMainFunction(m_state, ropPayload);

    // Write exploit trailer.
    exploit.setIndentLevel(4);
    exploit.writeline("proc.interactive()");

    // Write the buffered content to the file.
    std::string filename = exploit.getFilename(m_state->getID());
    std::ofstream ofs(filename);
    ofs << exploit.getContent();

    log<WARN>() << "Generated exploit script: " << filename << '\n';
    return true;
}

bool ExploitGenerator::generateExploit(const std::vector<uint8_t> &stage1,
                                       std::string filename) const {
    if (stage1.empty()) {
        return false;
    }

    std::ofstream ofs(filename, std::ios::binary);
    ofs.write(reinterpret_cast<const char *>(stage1.data()), stage1.size());

    log<WARN>() << "Generated exploit: " << filename << '\n';
    return true;
}

}  // namespace s2e::plugins::crax