# -*- coding: utf-8 -*-

# This code is part of Amoco
# Copyright (C) 2006-2011 Axel Tillequin (bdcht3@gmail.com)
# published under GPLv2 license

from amoco.system.core import *
from amoco.code import tag,xfunc

import amoco.arch.x64.cpu_x64 as cpu

PAGESIZE = 4096

class ELF(CoreExec):

    def __init__(self,p):
        CoreExec.__init__(self,p,cpu)
        if self.bin:
            self.symbols.update(self.bin.functions)
            self.symbols.update(self.bin.variables)

    # load the program into virtual memory (populate the mmap dict)
    def load_binary(self):
        p = self.bin
        if p!=None:
            # create text and data segments according to elf header:
            for s in p.Phdr:
                ms = p.loadsegment(s,PAGESIZE)
                if ms!=None:
                    vaddr,data = ms.popitem()
                    self.mmap.write(vaddr,data)
            # create the dynamic segments:
            self.load_shlib()
        # create the stack zone:
        self.mmap.newzone(cpu.rsp)

    # call dynamic linker to populate mmap with shared libs:
    # for now, the external libs are seen through the elf dynamic section:
    def load_shlib(self):
        for k,f in self.bin._Elf64__dynamic(None).items():
            self.mmap.write(k,cpu.ext(f,size=64))

    def initenv(self):
        from amoco.cas.mapper import mapper
        m = mapper()
        for k,v in ((cpu.rip, cpu.cst(self.bin.entrypoints[0],64)),
                    (cpu.rbp, cpu.cst(0,64)),
                    (cpu.rax, cpu.cst(0,64)),
                    (cpu.rbx, cpu.cst(0,64)),
                    (cpu.rcx, cpu.cst(0,64)),
                    (cpu.rdx, cpu.cst(0,64)),
                    (cpu.rsi, cpu.cst(0,64)),
                    (cpu.rdi, cpu.cst(0,64))):
            m[k] = v
        return m

    # seqhelper provides arch-dependent information to amoco.main classes
    def seqhelper(self,seq):
        for i in seq:
            # some basic hints:
            if i.mnemonic.startswith('RET'):
                i.misc[tag.FUNC_END]=1
                continue
            elif i.mnemonic in ('PUSH','ENTER'):
                i.misc[tag.FUNC_STACK]=1
                if i.operands and i.operands[0] is cpu.rbp:
                    i.misc[tag.FUNC_START]=1
                    continue
            elif i.mnemonic in ('POP','LEAVE'):
                i.misc[tag.FUNC_UNSTACK]=1
                if i.operands and i.operands[0] is cpu.rbp:
                    i.misc[tag.FUNC_END]=1
                    continue
            # provide hints of absolute location from relative offset:
            elif i.mnemonic in ('CALL','JMP','Jcc'):
                if i.mnemonic == 'CALL':
                    i.misc[tag.FUNC_CALL]=1
                    i.misc['retto'] = i.address+i.length
                else:
                    i.misc[tag.FUNC_GOTO]=1
                    if i.mnemonic == 'Jcc':
                        i.misc['cond'] = i.cond
                if (i.address is not None) and i.operands[0]._is_cst:
                    v = i.address+i.operands[0].signextend(64)+i.length
                    x = self.check_sym(v)
                    if x is not None: v=x
                    i.misc['to'] = v
                    if i.misc[tag.FUNC_CALL] and i.misc['retto']==v:
                        # this looks like a fake call
                        i.misc[tag.FUNC_CALL]=-1
                    continue
            # check operands (globals & .got calls):
            for op in i.operands:
                if op._is_mem:
                    if op.a.base is cpu.rbp:
                        if   op.a.disp<0: i.misc[tag.FUNC_VAR]=True
                        elif op.a.disp>=16: i.misc[tag.FUNC_ARG]=True
                    elif op.a.base._is_cst or (op.a.base is cpu.rip):
                        b = op.a.base
                        if b is cpu.rip: b=i.address+i.length
                        x = self.check_sym(b+op.a.disp)
                        if x is not None:
                            op.a.base=x
                            op.a.disp=0
                            if i.mnemonic == 'JMP': # PLT jumps:
                                i.misc[tag.FUNC_START]=1
                                i.misc[tag.FUNC_END]=1
                elif op._is_cst:
                    x = self.check_sym(op)
                    i.misc['imm_ref'] = x
        return seq

    def blockhelper(self,block):
        block._helper = block_helper_
        return CoreExec.blockhelper(self,block)

    def funchelper(self,f):
        # check single root node:
        roots = f.cfg.roots()
        if len(roots)==0:
            roots = filter(lambda n:n.data.misc[tag.FUNC_START],f.cfg.sV)
            if len(roots)==0:
                logger.warning("no entry to function %s found"%f)
        if len(roots)>1:
            logger.verbose('multiple entries into function %s ?!'%f)
        # check start symbol:
        elif roots[0].data.address == self.bin.entrypoints[0]:
            f.name = '_start'
        # get section symbol if any:
        f.misc['section'] = section = self.bin.getinfo(f.address.value)[0]
        rets = f.cfg.leaves()
        if len(rets)==0:
            logger.warning("no exit to function %s found"%f)
        if len(rets)>1:
            logger.verbose('multiple exits in function %s'%f)
        for r in rets:
           # export PLT external symbol name:
            if section and section.name=='.plt':
                if isinstance(r.data,xfunc): f.name = section.name+r.name
            if r.data.misc[tag.FUNC_CALL]:
                f.misc[tag.FUNC_CALL] += 1
        if f.map:
        # check vars & args: should reflect x64 register calling convention
            f.misc[tag.FUNC_VAR] = []
            f.misc[tag.FUNC_ARG] = []
            for x in set(f.map.inputs()):
                f.misc[tag.FUNC_IN] += 1
                if x._is_mem and x.a.base==cpu.rsp:
                    if x.a.disp>=8:
                        f.misc[tag.FUNC_ARG].append(x)
            for x in set(f.map.outputs()):
                if x in (cpu.rsp, cpu.rbp): continue
                f.misc[tag.FUNC_OUT] += 1
                if x._is_mem and x.a.base==cpu.rsp:
                    if x.a.disp<0:
                        f.misc[tag.FUNC_VAR].append(x)


def block_helper_(block,m):
    # annotations based on block semantics:
    sta,sto = block.support
    if m[cpu.mem(cpu.rbp-8,64)] == cpu.rbp:
        block.misc[tag.FUNC_START]=1
    if m[cpu.rip]==cpu.mem(cpu.rsp-8,64):
        block.misc[tag.FUNC_END]=1
    if m[cpu.mem(cpu.rsp,64)]==sto:
        block.misc[tag.FUNC_CALL]=1

# HOOKS DEFINED HERE :
#----------------------------------------------------------------------------

@stub_default
def pop_rip(m,**kargs):
    cpu.pop(m,cpu.rip)

@stub
def __libc_start_main(m,**kargs):
    "tags: func_call"
    m[cpu.rip] = m(cpu.rdi)
    cpu.push(m,cpu.ext('exit',size=64))

@stub
def exit(m,**kargs):
    m[cpu.rip] = top(64)
@stub
def abort(m,**kargs):
    m[cpu.rip] = top(64)
@stub
def __assert(m,**kargs):
    m[cpu.rip] = top(64)
@stub
def __assert_fail(m,**kargs):
    m[cpu.rip] = top(64)
@stub
def _assert_perror_fail(m,**kargs):
    m[cpu.rip] = top(64)

#----------------------------------------------------------------------------

# SYSCALLS:
#----------------------------------------------------------------------------
IDT={
  0: "read",
  1: "write",
  2: "open",
  3: "close",
  4: "stat",
  5: "fstat",
  6: "lstat",
  7: "poll",
  8: "lseek",
  9: "mmap",
 10: "mprotect",
 13: "rt_sigaction",
 14: "rt_sigprocmask",
 15: "rt_sigreturn",
 16: "ioctl",
 17: "pread64",
 18: "pwrite64",
 19: "readv",
 20: "writev",
 21: "access",
 22: "pipe",
 23: "select",
 24: "sched_yield",
 25: "mremap",
 26: "msync",
 27: "mincore",
 28: "madvise",
 29: "shmget",
 30: "shmat",
 31: "shmctl",
 32: "dup",
 33: "dup2",
 34: "pause",
 35: "nanosleep",
 36: "getitimer",
 37: "alarm",
 38: "setitimer",
 39: "getpid",
 40: "sendfile",
 41: "socket",
 42: "connect",
 43: "accept",
 44: "sendto",
 45: "recvfrom",
 46: "sendmsg",
 47: "recvmsg",
 48: "shutdown",
 49: "bind",
 50: "listen",
 51: "getsockname",
 52: "getpeername",
 53: "socketpair",
 54: "setsockopt",
 55: "getsockopt",
 56: "clone",
 57: "fork",
 58: "vfork",
 59: "execve",
 60: "exit",
 61: "wait4",
 62: "kill",
 63: "uname",
 64: "semget",
 65: "semop",
 66: "semctl",
 67: "shmdt",
 68: "msgget",
 69: "msgsnd",
 70: "msgrcv",
 71: "msgctl",
 72: "fcntl",
 73: "flock",
 74: "fsync",
 75: "fdatasync",
 76: "truncate",
 77: "ftruncate",
 78: "getdents",
 79: "getcwd",
 80: "chdir",
 81: "fchdir",
 82: "rename",
 83: "mkdir",
 84: "rmdir",
 85: "creat",
 86: "link",
 87: "unlink",
 88: "symlink",
 89: "readlink",
 90: "chmod",
 91: "fchmod",
 92: "chown",
 93: "fchown",
 94: "lchown",
 95: "umask",
 96: "gettimeofday",
 97: "getrlimit",
 98: "getrusage",
 99: "sysinfo",
100: "times",
101: "ptrace",
102: "getuid",
103: "syslog",
104: "getgid",
105: "setuid",
106: "setgid",
107: "geteuid",
108: "getegid",
109: "setpgid",
110: "getppid",
111: "getpgrp",
112: "setsid",
113: "setreuid",
114: "setregid",
115: "getgroups",
116: "setgroups",
117: "setresuid",
118: "getresuid",
119: "setresgid",
120: "getresgid",
121: "getpgid",
122: "setfsuid",
123: "setfsgid",
124: "getsid",
125: "capget",
126: "capset",
127: "rt_sigpending",
128: "rt_sigtimedwait",
129: "rt_sigqueueinfo",
130: "rt_sigsuspend",
131: "sigaltstack",
132: "utime",
133: "mknod",
134: "uselib",
135: "personality",
136: "ustat",
137: "statfs",
138: "fstatfs",
139: "sysfs",
140: "getpriority",
141: "setpriority",
142: "sched_setparam",
143: "sched_getparam",
144: "sched_setscheduler",
145: "sched_getscheduler",
146: "sched_get_priority_max",
147: "sched_get_priority_min",
148: "sched_rr_get_interval",
149: "mlock",
150: "munlock",
151: "mlockall",
152: "munlockall",
153: "vhangup",
154: "modify_ldt",
155: "pivot_root",
156: "_sysctl",
157: "prctl",
158: "arch_prctl",
159: "adjtimex",
160: "setrlimit",
161: "chroot",
162: "sync",
163: "acct",
164: "settimeofday",
165: "mount",
166: "umount2",
167: "swapon",
168: "swapoff",
169: "reboot",
170: "sethostname",
171: "setdomainname",
172: "iopl",
173: "ioperm",
174: "create_module",
175: "init_module",
176: "delete_module",
177: "get_kernel_syms",
178: "query_module",
179: "quotactl",
180: "nfsservctl",
181: "getpmsg",
182: "putpmsg",
183: "afs_syscall",
184: "tuxcall",
185: "security",
186: "gettid",
187: "readahead",
188: "setxattr",
189: "lsetxattr",
190: "fsetxattr",
191: "getxattr",
192: "lgetxattr",
193: "fgetxattr",
194: "listxattr",
195: "llistxattr",
196: "flistxattr",
197: "removexattr",
198: "lremovexattr",
199: "fremovexattr",
200: "tkill",
201: "time",
202: "futex",
203: "sched_setaffinity",
204: "sched_getaffinity",
205: "set_thread_area",
206: "io_setup",
207: "io_destroy",
208: "io_getevents",
209: "io_submit",
210: "io_cancel",
211: "get_thread_area",
212: "lookup_dcookie",
213: "epoll_create",
214: "epoll_ctl_old",
215: "epoll_wait_old",
216: "remap_file_pages",
217: "getdents64",
218: "set_tid_address",
219: "restart_syscall",
220: "semtimedop",
221: "fadvise64",
222: "timer_create",
223: "timer_settime",
224: "timer_gettime",
225: "timer_getoverrun",
226: "timer_delete",
227: "clock_settime",
228: "clock_gettime",
229: "clock_getres",
230: "clock_nanosleep",
231: "exit_group",
232: "epoll_wait",
233: "epoll_ctl",
234: "tgkill",
235: "utimes",
236: "vserver",
237: "mbind",
238: "set_mempolicy",
239: "get_mempolicy",
240: "mq_open",
241: "mq_unlink",
242: "mq_timedsend",
243: "mq_timedreceive",
244: "mq_notify",
245: "mq_getsetattr",
246: "kexec_load",
247: "waitid",
248: "add_key",
249: "request_key",
250: "keyctl",
251: "ioprio_set",
252: "ioprio_get",
253: "inotify_init",
254: "inotify_add_watch",
255: "inotify_rm_watch",
256: "migrate_pages",
257: "openat",
258: "mkdirat",
259: "mknodat",
260: "fchownat",
261: "futimesat",
262: "newfstatat",
263: "unlinkat",
264: "renameat",
265: "linkat",
266: "symlinkat",
267: "readlinkat",
268: "fchmodat",
269: "faccessat",
270: "pselect6",
271: "ppoll",
272: "unshare",
273: "set_robust_list",
274: "get_robust_list",
275: "splice",
276: "tee",
277: "sync_file_range",
278: "vmsplice",
279: "move_pages",
280: "utimensat",
281: "epoll_pwait",
282: "signalfd",
283: "timerfd_create",
284: "eventfd",
285: "fallocate",
286: "timerfd_settime",
287: "timerfd_gettime",
288: "accept4",
289: "signalfd4",
290: "eventfd2",
291: "epoll_create1",
292: "dup3",
293: "pipe2",
294: "inotify_init1",
295: "preadv",
296: "pwritev",
297: "rt_tgsigqueueinfo",
298: "perf_event_open",
299: "recvmmsg",
300: "fanotify_init",
301: "fanotify_mark",
302: "prlimit64",
303: "name_to_handle_at",
304: "open_by_handle_at",
305: "clock_adjtime",
306: "syncfs",
307: "sendmmsg",
308: "setns",
309: "getcpu",
310: "process_vm_readv",
311: "process_vm_writev",
312: "kcmp",
313: "finit_module",
314: "sched_setattr",
315: "sched_getattr",
316: "renameat2",
317: "seccomp",
318: "getrandom",
319: "memfd_create",
320: "kexec_file_load",
321: "bpf",
322: "execveat" }

