import os
import sys
from collections import namedtuple
import unittest

"""
This script applies the following assumptions:

1.  Each virtual call has the pattern:
    {callq *(reg + ofset)} ^ {reg is 64 bits width} ^ {ofset can be divided by 8}

2.  There is no other control-flow instructions execept a conditional jump instruction
    before a protected virtual call. The jump target has the following patter:
    "movabs xxx"
"""


# Store configurations to VLOOM-LLVM
gMapConf = {}


# Read the configuration options to VLOOM-LLVM
def read_llvm_config(szConf):
    # format:
    # VLOOM_SIZE:<type 'int'>:[8,64]:20
    arLines = []
    with open(szConf, "r") as f:
        arLines = [l.strip() for l in f.readlines()]

    for l in arLines:
        fs = l.split(":")
        if len(fs) != 4:
            continue
        K = fs[0]
        V = fs[3]
        gMapConf[K] = V


# Generate a Makefile to compile libc.cpp, exporting env-variables accoring to config
def generate_makefile(szMakefile="Makefile"):
    pream = """
VLOOM_LLVM_ROOT	:= $(shell pwd)/../..
CC	:= ${VLOOM_LLVM_ROOT}/build_debug/bin/clang
CXX	:= ${VLOOM_LLVM_ROOT}/build_debug/bin/clang++

CXXFLAGS	+= -std=c++98 -shared -fPIC
LDFLAGS		+= -fuse-ld=lld -Wl,--emit-relocs -Wl,-znotext"""

    post = """
all: lib.so

lib.so: lib.cpp
	$(CXX) ${CXXFLAGS} ${LDFLAGS} $^ -o $@

clean:
	rm -f lib.so"""

    arExport = ["export %s=%s" % (k, v) for k, v in gMapConf.items()]
    mid = "\n".join(arExport)

    with open(szMakefile, "w+") as f:
        f.writelines(pream)
        f.writelines("\n")
        f.writelines(mid)
        f.writelines("\n")
        f.writelines(post)
        f.writelines("\n")


# Read in a file generated by "objump -d -z binary"
def build_libcpp_asm_file(szMakefile="Makefile", szAsmFile="lib.so.asm"):
    szMake = "make -f %s" % szMakefile
    os.system(szMake)

    szObjdump = "objdump -d -z lib.so > %s" % szAsmFile
    os.system(szObjdump)


def read_asm_file(szFile="lib.so.asm"):
    arLines = []
    with open(szFile, "r") as f:
        arLines = [l.strip() for l in f.readlines()]
    return arLines


# Stores all the assembly codes
AsmInst = namedtuple("AsmInst", "addr t opc asm")
glstAsmCodes = []  # stores instruction sequencically
gmapAsmCodes = {}  # for quick searching: addr -> idx to glstAsmCodes
INST_TYPE = {"NORM": 1, "JCC": 2, "ICALL": 3, "CALL": 4, "JMP": 5, "IJMP": 6, "RET": 7}


# Parse assembly code
def parse_asm_code(assembly_lines):
    # objdump assembly code format ==> addr:\t 24-bytes \t asm-code #comment
    # 1724: 48 8b 05 bd 68 20 00  mov 0x2068bd(%rip),%rax  # 207fe8 <__gmon_start__>
    # 1b20: 64 48 8b 04 25 28 00  mov %fs:0x28,%rax
    # exception cases:
    # 1. 302f:       00
    for l in assembly_lines:
        nColmn = l.find(":")
        if nColmn == -1:
            continue
        try:
            addr = int(l[:nColmn], 16)
        except:
            continue

        if len(l) <= nColmn + 1 + 24:
            continue

        szInst = l[nColmn + 1 + 23:].split("#")[0].strip()
        arInst = szInst.split()
        szOPCode = arInst[0]
        if szOPCode[:3] == "jmp":
            szOPRand = arInst[1]
            if szOPRand[0] == "*":
                t = INST_TYPE["IJMP"]
            else:
                t = INST_TYPE["JMP"]

        elif szOPCode[:1] == "j":
            t = INST_TYPE["JCC"]

        elif szOPCode[:4] == "call":
            szOPRand = arInst[1]
            if szOPRand[0] == "*":
                t = INST_TYPE["ICALL"]
            else:
                t = INST_TYPE["CALL"]

        elif szOPCode[:3] == "ret":
            t = INST_TYPE["RET"]
        else:
            t = INST_TYPE["NORM"]

        r = AsmInst(addr, t, szOPCode, szInst)
        gmapAsmCodes[addr] = len(glstAsmCodes)
        glstAsmCodes.append(r)


# Collect all protected vcall-sites
def find_first_instrumented_vcallsite():
    # Example:
    # 2143: jmp    2167
    # 2145: mov    $0x24000102,%eax
    # 214a: nopl   0x0(%rax,%rax,1)
    # 2152: nopl   0x0(%rax,%rax,1)
    # 215a: nopl   0x0(%rax)
    # 215e: nop
    # 215f: add    %al,(%rax)
    # 2161: add    %al,(%rax)
    # 2163: add    %al,(%rax)
    # 2165: add    %al,(%rax)
    # 2167: callq  *(%rax)

    idx = 0
    while idx < len(glstAsmCodes):
        r = glstAsmCodes[idx]
        if r.t != INST_TYPE["JMP"]:
            idx += 1
        else:
            newidx, rvcall = check_vcfi_checker(idx, r)
            if rvcall is not None:
                rmovflag = glstAsmCodes[idx+1]
                return rmovflag
            idx = newidx + 1

    # end while
    return None


# Check if this instruction conditionally targets a cross-DSO checker
def check_vcfi_checker(idx, rjmp):
    if idx + 3 > len(glstAsmCodes):
        return idx, None
    # Followed by mov & nop & nop
    rmov = glstAsmCodes[idx+1]
    if (rmov.opc != "mov") or ("$0x" not in rmov.asm):
        return idx, None
    rnop1 = glstAsmCodes[idx+2]
    if rnop1.opc != "nop" and rnop1.opc != "nopl":
        return idx+1, None
    rnop2 = glstAsmCodes[idx+3]
    if rnop2.opc != "nop" and rnop2.opc != "nopl":
        return idx+2, None

    # The target block ended by ICALL or IJMP
    szOprand = rjmp.asm.split()[1]
    addr = int(szOprand, 16)
    idx_jmp = gmapAsmCodes.get(addr)

    # get the indirect call or indirect jmp instruction
    for i in range(idx_jmp, len(glstAsmCodes)):
        r = glstAsmCodes[i]
        if r.t == INST_TYPE["NORM"]:
            continue

        if (r.t == INST_TYPE["IJMP"]) or (r.t == INST_TYPE["ICALL"]):
            return i, r
        else:
            return i, None


gMapFlag = {}
# Parse the flags stored in mov instruction


def parse_movflag_instruction(mov_inst):
    if mov_inst is None:
        gMapFlag["VLOOM_DISABLED"] = "true"
        return None
    else:
        gMapFlag["VLOOM_DISABLED"] = "false"
    oprand = mov_inst.asm.split()[1]
    opr1 = oprand.split(",")[0]
    opr1_v = int(opr1[1:], 16)
    b1 = opr1_v & 0xff  # VLOOM_SIZE
    b2 = (opr1_v >> 8) & 0xff  # VLOOM_SCRATCH
    b3 = (opr1_v >> 16) & 0xff  # VLOOM_SIDE_EFFECT
    b4 = (opr1_v >> 24) & 0xff  # VLOOM_REGSTER_SET
    # print("%2x, %2x, %2x, %2x" % (b1, b2, b3, b4))
    gMapFlag["VLOOM_SIZE"] = str(b1)
    gMapFlag["VLOOM_SCRATCH"] = str(b2)
    if b3 == 0:
        szEffect = "false"
    else:
        szEffect = "true"
    gMapFlag["VLOOM_SIDE_EFFECT"] = szEffect
    gMapFlag["VLOOM_REGSTER_SET"] = str(b4)
    return (b1, b2, b3, b4)


class TestStringMethods(unittest.TestCase):
    def test_vloom_llvm(self):
        szFlagDis = gMapFlag.get("VLOOM_DISABLED")
        if (szFlagDis is not None) and (szFlagDis == "true"):
            szConfDis = gMapConf.get("VLOOM_DISABLED")
            self.assertTrue(szConfDis)

        szFlagSize = gMapFlag.get("VLOOM_SIZE")
        szConfSize = gMapConf.get("VLOOM_SIZE")
        self.assertEqual(szFlagSize, szConfSize)

        szFlagReg = gMapFlag.get("VLOOM_SCRATCH")
        szConfReg = gMapConf.get("VLOOM_SCRATCH")
        self.assertEqual(szFlagReg, szConfReg)

        szFlagEfft = gMapFlag.get("VLOOM_SIDE_EFFECT")
        szConfEfft = gMapConf.get("VLOOM_SIDE_EFFECT")
        self.assertEqual(szFlagEfft,  szConfEfft)

        szFlagRSET = gMapFlag.get("VLOOM_REGSTER_SET")
        szConfRSET = gMapConf.get("VLOOM_REGSTER_SET")
        self.assertEqual(szFlagRSET, szConfRSET)


def main(szConf, szOutFile=None):
    # 1. Read configurations
    read_llvm_config(szConf)

    # 2. Generate Makefile
    generate_makefile()

    # 3. Build the source code and objdump to get the assebmly code
    build_libcpp_asm_file()

    # 1. read asm file generated by objdump
    arLines = read_asm_file()

    # 2. filter code lines
    parse_asm_code(arLines)

    # 5. find vcall instructions
    rmov = find_first_instrumented_vcallsite()

    # simple report
    parse_movflag_instruction(rmov)

    # report_verify_result()
    sys.argv.pop()
    unittest.main()


if __name__ == "__main__":
    if len(sys.argv) < 2 or len(sys.argv) > 3:
        print("%s llvm-conf-file [output-file]" % sys.argv[0])
    elif len(sys.argv) == 2:
        main(sys.argv[1], None)
    else:
        main(sys.argv[1], sys.argv[2])
