from pwnypack.shellcode.linux import Linux
from pwnypack.shellcode.aarch64 import AArch64
from pwnypack.shellcode.mutable_data import gnu_as_mutable_data_finalizer
from pwnypack.shellcode.stack_data import stack_data_finalizer
from pwnypack.shellcode.types import SyscallDef, PTR


__all__ = ['LinuxAArch64Mutable', 'LinuxAArch64Stack']


class LinuxAArch64(Linux, AArch64):
    """
    An environment that targets a generic Linux AArch64 machine.
    """

    sys_rt_sigreturn = SyscallDef('sys_rt_sigreturn', PTR)  #:

    SYSCALL_ARG_MAP = [AArch64.X0, AArch64.X1, AArch64.X2, AArch64.X3, AArch64.X4, AArch64.X5]
    SYSCALL_REG = AArch64.X8
    SYSCALL_RET_REG = AArch64.X0
    SYSCALL_INSTR = 'svc #0'

    SYSCALL_MAP = {
        Linux.sys_io_setup: 0,
        Linux.sys_io_destroy: 1,
        Linux.sys_io_submit: 2,
        Linux.sys_io_cancel: 3,
        Linux.sys_io_getevents: 4,
        Linux.sys_setxattr: 5,
        Linux.sys_lsetxattr: 6,
        Linux.sys_fsetxattr: 7,
        Linux.sys_getxattr: 8,
        Linux.sys_lgetxattr: 9,
        Linux.sys_fgetxattr: 10,
        Linux.sys_listxattr: 11,
        Linux.sys_llistxattr: 12,
        Linux.sys_flistxattr: 13,
        Linux.sys_removexattr: 14,
        Linux.sys_lremovexattr: 15,
        Linux.sys_fremovexattr: 16,
        Linux.sys_getcwd: 17,
        Linux.sys_lookup_dcookie: 18,
        Linux.sys_eventfd2: 19,
        Linux.sys_epoll_create1: 20,
        Linux.sys_epoll_ctl: 21,
        Linux.sys_epoll_pwait: 22,
        Linux.sys_dup: 23,
        Linux.sys_dup3: 24,
        Linux.sys_fcntl: 25,
        Linux.sys_inotify_init1: 26,
        Linux.sys_inotify_add_watch: 27,
        Linux.sys_inotify_rm_watch: 28,
        Linux.sys_ioctl: 29,
        Linux.sys_ioprio_set: 30,
        Linux.sys_ioprio_get: 31,
        Linux.sys_flock: 32,
        Linux.sys_mknodat: 33,
        Linux.sys_mkdirat: 34,
        Linux.sys_unlinkat: 35,
        Linux.sys_symlinkat: 36,
        Linux.sys_linkat: 37,
        Linux.sys_renameat: 38,
        Linux.sys_umount2: 39,
        Linux.sys_mount: 40,
        Linux.sys_pivot_root: 41,
        Linux.sys_nfsservctl: 42,
        Linux.sys_statfs: 43,
        Linux.sys_fstatfs: 44,
        Linux.sys_truncate: 45,
        Linux.sys_ftruncate: 46,
        Linux.sys_fallocate: 47,
        Linux.sys_faccessat: 48,
        Linux.sys_chdir: 49,
        Linux.sys_fchdir: 50,
        Linux.sys_chroot: 51,
        Linux.sys_fchmod: 52,
        Linux.sys_fchmodat: 53,
        Linux.sys_fchownat: 54,
        Linux.sys_fchown: 55,
        Linux.sys_openat: 56,
        Linux.sys_close: 57,
        Linux.sys_vhangup: 58,
        Linux.sys_pipe2: 59,
        Linux.sys_quotactl: 60,
        Linux.sys_getdents64: 61,
        Linux.sys_lseek: 62,
        Linux.sys_read: 63,
        Linux.sys_write: 64,
        Linux.sys_readv: 65,
        Linux.sys_writev: 66,
        Linux.sys_pread64: 67,
        Linux.sys_pwrite64: 68,
        Linux.sys_preadv: 69,
        Linux.sys_pwritev: 70,
        Linux.sys_sendfile: 71,
        Linux.sys_pselect6: 72,
        Linux.sys_ppoll: 73,
        Linux.sys_signalfd4: 74,
        Linux.sys_vmsplice: 75,
        Linux.sys_splice: 76,
        Linux.sys_tee: 77,
        Linux.sys_readlinkat: 78,
        Linux.sys_fstatat64: 79,
        Linux.sys_fstat: 80,
        Linux.sys_sync: 81,
        Linux.sys_fsync: 82,
        Linux.sys_fdatasync: 83,
        Linux.sys_sync_file_range2: 84,
        Linux.sys_timerfd_create: 85,
        Linux.sys_timerfd_settime: 86,
        Linux.sys_timerfd_gettime: 87,
        Linux.sys_utimensat: 88,
        Linux.sys_acct: 89,
        Linux.sys_capget: 90,
        Linux.sys_capset: 91,
        Linux.sys_personality: 92,
        Linux.sys_exit: 93,
        Linux.sys_exit_group: 94,
        Linux.sys_waitid: 95,
        Linux.sys_set_tid_address: 96,
        Linux.sys_unshare: 97,
        Linux.sys_futex: 98,
        Linux.sys_set_robust_list: 99,
        Linux.sys_get_robust_list: 100,
        Linux.sys_nanosleep: 101,
        Linux.sys_getitimer: 102,
        Linux.sys_setitimer: 103,
        Linux.sys_kexec_load: 104,
        Linux.sys_init_module: 105,
        Linux.sys_delete_module: 106,
        Linux.sys_timer_create: 107,
        Linux.sys_timer_gettime: 108,
        Linux.sys_timer_getoverrun: 109,
        Linux.sys_timer_settime: 110,
        Linux.sys_timer_delete: 111,
        Linux.sys_clock_settime: 112,
        Linux.sys_clock_gettime: 113,
        Linux.sys_clock_getres: 114,
        Linux.sys_clock_nanosleep: 115,
        Linux.sys_syslog: 116,
        Linux.sys_ptrace: 117,
        Linux.sys_sched_setparam: 118,
        Linux.sys_sched_setscheduler: 119,
        Linux.sys_sched_getscheduler: 120,
        Linux.sys_sched_getparam: 121,
        Linux.sys_sched_setaffinity: 122,
        Linux.sys_sched_getaffinity: 123,
        Linux.sys_sched_yield: 124,
        Linux.sys_sched_get_priority_max: 125,
        Linux.sys_sched_get_priority_min: 126,
        Linux.sys_sched_rr_get_interval: 127,
        Linux.sys_restart_syscall: 128,
        Linux.sys_kill: 129,
        Linux.sys_tkill: 130,
        Linux.sys_tgkill: 131,
        Linux.sys_sigaltstack: 132,
        Linux.sys_rt_sigsuspend: 133,
        Linux.sys_rt_sigaction: 134,
        Linux.sys_rt_sigprocmask: 135,
        Linux.sys_rt_sigpending: 136,
        Linux.sys_rt_sigtimedwait: 137,
        Linux.sys_rt_sigqueueinfo: 138,
        sys_rt_sigreturn: 139,
        Linux.sys_setpriority: 140,
        Linux.sys_getpriority: 141,
        Linux.sys_reboot: 142,
        Linux.sys_setregid: 143,
        Linux.sys_setgid: 144,
        Linux.sys_setreuid: 145,
        Linux.sys_setuid: 146,
        Linux.sys_setresuid: 147,
        Linux.sys_getresuid: 148,
        Linux.sys_setresgid: 149,
        Linux.sys_getresgid: 150,
        Linux.sys_setfsuid: 151,
        Linux.sys_setfsgid: 152,
        Linux.sys_times: 153,
        Linux.sys_setpgid: 154,
        Linux.sys_getpgid: 155,
        Linux.sys_getsid: 156,
        Linux.sys_setsid: 157,
        Linux.sys_getgroups: 158,
        Linux.sys_setgroups: 159,
        Linux.sys_uname: 160,
        Linux.sys_sethostname: 161,
        Linux.sys_setdomainname: 162,
        Linux.sys_getrlimit: 163,
        Linux.sys_setrlimit: 164,
        Linux.sys_getrusage: 165,
        Linux.sys_umask: 166,
        Linux.sys_prctl: 167,
        Linux.sys_getcpu: 168,
        Linux.sys_gettimeofday: 169,
        Linux.sys_settimeofday: 170,
        Linux.sys_adjtimex: 171,
        Linux.sys_getpid: 172,
        Linux.sys_getppid: 173,
        Linux.sys_getuid: 174,
        Linux.sys_geteuid: 175,
        Linux.sys_getgid: 176,
        Linux.sys_getegid: 177,
        Linux.sys_gettid: 178,
        Linux.sys_sysinfo: 179,
        Linux.sys_mq_open: 180,
        Linux.sys_mq_unlink: 181,
        Linux.sys_mq_timedsend: 182,
        Linux.sys_mq_timedreceive: 183,
        Linux.sys_mq_notify: 184,
        Linux.sys_mq_getsetattr: 185,
        Linux.sys_msgget: 186,
        Linux.sys_msgctl: 187,
        Linux.sys_msgrcv: 188,
        Linux.sys_msgsnd: 189,
        Linux.sys_semget: 190,
        Linux.sys_semctl: 191,
        Linux.sys_semtimedop: 192,
        Linux.sys_semop: 193,
        Linux.sys_shmget: 194,
        Linux.sys_shmctl: 195,
        Linux.sys_shmat: 196,
        Linux.sys_shmdt: 197,
        Linux.sys_socket: 198,
        Linux.sys_socketpair: 199,
        Linux.sys_bind: 200,
        Linux.sys_listen: 201,
        Linux.sys_accept: 202,
        Linux.sys_connect: 203,
        Linux.sys_getsockname: 204,
        Linux.sys_getpeername: 205,
        Linux.sys_sendto: 206,
        Linux.sys_recvfrom: 207,
        Linux.sys_setsockopt: 208,
        Linux.sys_getsockopt: 209,
        Linux.sys_shutdown: 210,
        Linux.sys_sendmsg: 211,
        Linux.sys_recvmsg: 212,
        Linux.sys_readahead: 213,
        Linux.sys_brk: 214,
        Linux.sys_munmap: 215,
        Linux.sys_mremap: 216,
        Linux.sys_add_key: 217,
        Linux.sys_request_key: 218,
        Linux.sys_keyctl: 219,
        Linux.sys_clone: 220,
        Linux.sys_execve: 221,
        Linux.sys_mmap2: 222,
        Linux.sys_fadvise64: 223,
        Linux.sys_swapon: 224,
        Linux.sys_swapoff: 225,
        Linux.sys_mprotect: 226,
        Linux.sys_msync: 227,
        Linux.sys_mlock: 228,
        Linux.sys_munlock: 229,
        Linux.sys_mlockall: 230,
        Linux.sys_munlockall: 231,
        Linux.sys_mincore: 232,
        Linux.sys_madvise: 233,
        Linux.sys_remap_file_pages: 234,
        Linux.sys_mbind: 235,
        Linux.sys_get_mempolicy: 236,
        Linux.sys_set_mempolicy: 237,
        Linux.sys_migrate_pages: 238,
        Linux.sys_move_pages: 239,
        Linux.sys_rt_tgsigqueueinfo: 240,
        Linux.sys_perf_event_open: 241,
        Linux.sys_accept4: 242,
        Linux.sys_recvmmsg: 243,
        Linux.sys_wait4: 260,
        Linux.sys_prlimit64: 261,
        Linux.sys_fanotify_init: 262,
        Linux.sys_fanotify_mark: 263,
        Linux.sys_name_to_handle_at: 264,
        Linux.sys_open_by_handle_at: 265,
        Linux.sys_clock_adjtime: 266,
        Linux.sys_syncfs: 267,
        Linux.sys_setns: 268,
        Linux.sys_sendmmsg: 269,
        Linux.sys_process_vm_readv: 270,
        Linux.sys_process_vm_writev: 271,
        Linux.sys_kcmp: 272,
        Linux.sys_finit_module: 273,
        Linux.sys_sched_setattr: 274,
        Linux.sys_sched_getattr: 275,
        Linux.sys_renameat2: 276,
        Linux.sys_seccomp: 277,
        Linux.sys_getrandom: 278,
        Linux.sys_memfd_create: 279,
        Linux.sys_bpf: 280,
        Linux.sys_execveat: 281,
        Linux.sys_userfaultfd: 282,
        Linux.sys_membarrier: 283,
        Linux.sys_mlock2: 284,
        Linux.sys_copy_file_range: 285,
    }


class LinuxAArch64Mutable(LinuxAArch64):
    """
    An environment that targets a 64-bit Linux ARM machine in a writable segment.
    """

    data_finalizer = gnu_as_mutable_data_finalizer(lambda env, _: ['\tadr %s, __data' % env.OFFSET_REG], '//')


class LinuxAArch64Stack(LinuxAArch64):
    """
    An environment that targets a 64-bit Linux ARM machine that allocates the
    required data on the stack.
    """

    data_finalizer = stack_data_finalizer(16)
