// 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.

#ifndef S2E_PLUGINS_CRAX_EXPLOIT_H
#define S2E_PLUGINS_CRAX_EXPLOIT_H

#include <s2e/S2E.h>
#include <s2e/S2EExecutionState.h>
#include <s2e/Plugins/CRAX/Pwnlib/ELF.h>
#include <s2e/Plugins/CRAX/Pwnlib/Process.h>

#include <exception>
#include <map>
#include <string>
#include <vector>

namespace s2e::plugins::crax {

class Script {
public:
    class DuplicatedSymbolException : public std::exception {
    public:
        DuplicatedSymbolException(const std::string &name) : m_name(name) {}
        virtual ~DuplicatedSymbolException() override = default;

        virtual const char *what() const throw() override {
            return "Duplicated symbol found: \n";
        }

    private:
        std::string m_name;
    };


    class UndefinedSymbolException : public std::exception {
    public:
        UndefinedSymbolException(const std::string &name) : m_name(name) {}
        virtual ~UndefinedSymbolException() override = default;

        virtual const char *what() const throw() override {
            return "Undefined symbol found.\n";
        }

    private:
        std::string m_name;
    };


    Script() : m_indentLevel(), m_content(), m_symtab() {}
    virtual ~Script() = default;

    void reset();
    bool hasSymbol(const std::string &name) const;
    void registerSymbol(const std::string &name, uint64_t value);
    uint64_t getSymbolValue(const std::string &name) const;

    void writeline() {
        m_content += '\n';
    }

    void writeline(const std::string &line) {
        if (line.size()) {
            indent();
            m_content += line;
        }
        writeline();
    }

    void writelines(const std::vector<std::string> &lines) {
        for (const auto &line : lines) {
            writeline(line);
        }
    }

    std::string getFilename(int stateID) const;
    unsigned getIndentLevel() const { return m_indentLevel; }
    const std::string &getContent() const { return m_content; }
    const std::map<std::string, uint64_t> &getSymtab() const { return m_symtab; }

    void setIndentLevel(unsigned indentLevel) { m_indentLevel = indentLevel; }

protected:
    void indent() {
        for (unsigned i = 0; i < m_indentLevel; i++) {
            m_content += ' ';
        }
    }

    unsigned m_indentLevel;
    std::string m_content;
    std::map<std::string, uint64_t> m_symtab;
};


class Exploit : public Script {
public:
    Exploit(const std::string &elfFilename,
            const std::string &libcFilename,
            const std::string &ldFilename);
    virtual ~Exploit() override = default;

    // Look for an exact match of the gadget specified by `gadgetAsm` within `elf`.
    // If found, then the offset of the gadget will be returned, and zero otherwise.
    uint64_t resolveGadget(const ELF &elf, const std::string &gadgetAsm) const;

    void appendRopPayload(const std::string &ropPayload);
    void flushRopPayload();

    void writeLeakCanary();
    void writeLeakElfBase(uint64_t offset);
    void writeLeakLibcBase(uint64_t offset);

    const ELF &getElf() const { return m_elf; }
    const ELF &getLibc() const { return m_libc; }
    const ELF &getLd() const { return m_ld; }
    const Process &getProcess() const { return m_process; }

    ELF &getElf() { return m_elf; }
    ELF &getLibc() { return m_libc; }
    ELF &getLd() { return m_ld; }
    Process &getProcess() { return m_process; }

    static std::string toVarName(const std::string &s);
    static std::string toVarName(const ELF &elf, const std::string &gadgetAsm);

    static const std::string s_shebang;

private:
    // ELF files
    ELF m_elf;
    ELF m_libc;
    ELF m_ld;

    // By default, the exploit script starts the target binary with
    // the specified ld-linux-x86-64.so.2 and libc.so.6. We may also
    // choose to provde extra command-line arguments and environment variables.
    Process m_process;

    // Buffered lines.
    std::vector<std::string> m_ropPayloadLines;
};

}  // namespace s2e::plugins::crax

#endif  // S2E_PLUGINS_CRAX_EXPLOIT_H
