// 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/Utils/StringUtil.h>

#include "Exploit.h"

namespace s2e::plugins::crax {

void Script::reset() {
    m_indentLevel = 0;
    m_content.clear();
    //m_symtab.clear();
}

bool Script::hasSymbol(const std::string &name) const {
    return m_symtab.find(name) != m_symtab.end();
}

void Script::registerSymbol(const std::string &name, uint64_t addr) {
    //if (m_symtab.find(name) != m_symtab.end()) {
    //    throw Script::DuplicatedSymbolException(name);
    //}
    m_symtab[name] = addr;
}

uint64_t Script::getSymbolValue(const std::string &name) const {
    auto it = m_symtab.find(name);
    if (it == m_symtab.end()) {
        throw Script::UndefinedSymbolException(name);
    }
    return it->second;
}

std::string Script::getFilename(int stateID) const {
    return format("exploit_%d.py", stateID);
}



const std::string Exploit::s_shebang = "#!/usr/bin/env python3";

Exploit::Exploit(const std::string &elfFilename,
                 const std::string &libcFilename,
                 const std::string &ldFilename)
    : Script(),
      m_elf(elfFilename),
      m_libc(libcFilename),
      m_ld(ldFilename),
      m_process(ldFilename, elfFilename, libcFilename),
      m_ropPayloadLines() {}


uint64_t Exploit::resolveGadget(const ELF &elf, const std::string &gadgetAsm) const {
    return g_crax->getExploitGenerator().getRopGadgetResolver()
        .resolveGadget(elf, gadgetAsm);
}

void Exploit::appendRopPayload(const std::string &ropPayload) {
    if (m_ropPayloadLines.empty()) {
        m_ropPayloadLines.push_back(format("payload  = %s", ropPayload.c_str()));
    } else {
        m_ropPayloadLines.push_back(format("payload += %s", ropPayload.c_str()));
    }
}

void Exploit::flushRopPayload() {
    std::vector<std::string> lines;

    writelines(m_ropPayloadLines);
    m_ropPayloadLines.clear();

    // FIXME: Perhaps we should find a better way to do this...
    static const char *sendStr = m_elf.hasSymbol("gets") ? "sendline" : "send";
    lines.push_back(format("proc.%s(payload)", sendStr));
    lines.push_back("time.sleep(0.2)");

    lines.push_back("");
    writelines(lines);
}

void Exploit::writeLeakCanary() {
    writelines({
        "canary = u64(b'\\x00' + proc.recv(7))",
        "log.info('leaked canary: {}'.format(hex(canary)))",
    });
}

void Exploit::writeLeakElfBase(uint64_t offset) {
    const std::string &prefix = m_elf.getVarPrefix();

    writelines({
        "leaked = u64(proc.recv(6).ljust(8, b'\\x00'))",
        format("%s_base = leaked - 0x%x", prefix.c_str(), offset),
        format("log.info('leaked ELF base: {}'.format(hex(%s_base)))", prefix.c_str()),
    });
}

void Exploit::writeLeakLibcBase(uint64_t offset) {
    const std::string &prefix = m_libc.getVarPrefix();

    writelines({
        "leaked = u64(proc.recv(6).ljust(8, b'\\x00'))",
        format("%s_base = leaked - 0x%x", prefix.c_str(), offset),
        format("log.info('leaked libc base: {}'.format(hex(%s_base)))", prefix.c_str()),
    });
}

std::string Exploit::toVarName(const std::string &s) {
    std::string ret;

    for (auto c : s) {
        if (std::isalnum(c)) {
            ret += c;
        } else if (ret.size() && ret.back() != '_') {
            ret += '_';
        }
    }
    return ret;
}

std::string Exploit::toVarName(const ELF &elf, const std::string &gadgetAsm) {
    return toVarName(elf.getFilename() + '_' + gadgetAsm);
}

}  // namespace s2e::plugins::crax
