diff --git a/.gitignore b/.gitignore
index 88cafdb..b0aeead 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,5 +5,9 @@
 *.order
 *.o
 
+a.out
+*.proc
+
 .vscode
-app/target
\ No newline at end of file
+app/target
+**.idea
\ No newline at end of file
diff --git a/app/src/main.rs b/app/src/main.rs
index 917e14a..03b1316 100644
--- a/app/src/main.rs
+++ b/app/src/main.rs
@@ -8,6 +8,7 @@ use object::{Object, ObjectSymbol};
 use regex::Regex;
 use serde::{Deserialize, Serialize};
 use std::{
+    env,
     ffi::CStr,
     fs::{self, File, OpenOptions},
     io::{Read, Seek, Write},
@@ -35,10 +36,47 @@ struct ProbeRaw {
     rip: u64,
     rsp: u64,
     rbp: u64,
-    rdi: u64,
     ss: u64,
+    rsi: u64,
+    rdi: u64,
+    cs: u64,
+    ds: u64,
+    es: u64,
+    fs: u64,
+    gs: u64,
     exe: [u8; 4096],
 }
+
+fn new_empty_raw_probe() -> ProbeRaw {
+    ProbeRaw {
+        pid: 0,
+        rax: 0,
+        rbx: 0,
+        rcx: 0,
+        rdx: 0,
+        r8: 0,
+        r9: 0,
+        r10: 0,
+        r11: 0,
+        r12: 0,
+        r13: 0,
+        r14: 0,
+        r15: 0,
+        rip: 0,
+        rsp: 0,
+        rbp: 0,
+        ss: 0,
+        rsi: 0,
+        rdi: 0,
+        cs: 0,
+        ds: 0,
+        es: 0,
+        fs: 0,
+        gs: 0,
+        exe: [0; 4096],
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 struct Probe {
     pid: u64,
@@ -57,8 +95,14 @@ struct Probe {
     rip: u64,
     rsp: u64,
     rbp: u64,
-    rdi: u64,
     ss: u64,
+    rsi: u64,
+    rdi: u64,
+    cs: u64,
+    ds: u64,
+    es: u64,
+    fs: u64,
+    gs: u64,
     exe: String,
     stack_from: u64,
     stack: Vec<u8>,
@@ -92,72 +136,12 @@ enum Command {
 }
 
 fn main() {
+    env::set_var("RUST_BACKTRACE", "1");
     let cli = Cli::parse();
 
     match cli.command {
         Command::ReadRegs { pid } => {
-            let file = OpenOptions::new()
-                .read(true)
-                .write(true)
-                .open("/dev/get_task")
-                .expect("Failed to open device");
-
-            let mut data = ProbeRaw {
-                pid,
-                rax: 0,
-                rbx: 0,
-                rcx: 0,
-                rdx: 0,
-                r8: 0,
-                r9: 0,
-                r10: 0,
-                r11: 0,
-                r12: 0,
-                r13: 0,
-                r14: 0,
-                r15: 0,
-                rip: 0,
-                rsp: 0,
-                rbp: 0,
-                rdi: 0,
-                ss: 0,
-                exe: [0; 4096],
-            };
-
-            let exe;
-
-            unsafe {
-                read_regs(file.as_raw_fd(), &mut data).expect("ioctl failed");
-
-                exe = CStr::from_ptr(data.exe.as_ptr() as *const _)
-                    .to_str()
-                    .unwrap()
-                    .to_owned();
-            }
-
-            let data = Probe {
-                pid: data.pid,
-                rax: data.rax,
-                rbx: data.rbx,
-                rcx: data.rcx,
-                rdx: data.rdx,
-                r8: data.r8,
-                r9: data.r9,
-                r10: data.r10,
-                r11: data.r11,
-                r12: data.r12,
-                r13: data.r13,
-                r14: data.r14,
-                r15: data.r15,
-                rip: data.rip,
-                rsp: data.rsp,
-                rbp: data.rbp,
-                rdi: data.rdi,
-                ss: data.ss,
-                exe,
-                stack_from: 0,
-                stack: vec![],
-            };
+            let data = read_regs_with_get_tasks(pid);
 
             println!("ioctl succeeded");
             println!("{:#x?}", data);
@@ -174,170 +158,124 @@ fn main() {
             println!("{:x} {:x}", heap_from, heap_to);
         }
         Command::Dump { pid, output } => {
-            let file = OpenOptions::new()
-                .read(true)
-                .write(true)
-                .open("/dev/get_task")
-                .expect("Failed to open device");
-
-            let mut data = ProbeRaw {
-                pid,
-                rax: 0,
-                rbx: 0,
-                rcx: 0,
-                rdx: 0,
-                r8: 0,
-                r9: 0,
-                r10: 0,
-                r11: 0,
-                r12: 0,
-                r13: 0,
-                r14: 0,
-                r15: 0,
-                rip: 0,
-                rsp: 0,
-                rbp: 0,
-                rdi: 0,
-                ss: 0,
-                exe: [0; 4096],
-            };
-
-            let exe;
-
-            unsafe {
-                read_regs(file.as_raw_fd(), &mut data).expect("ioctl failed");
-
-                exe = CStr::from_ptr(data.exe.as_ptr() as *const _)
-                    .to_str()
-                    .unwrap()
-                    .to_owned();
-            }
-
-            let (stack_from, _) = get_mem_region_limits(Pid::from_raw(pid as i32), "stack");
-
-            let mut data = Probe {
-                pid: data.pid,
-                rax: data.rax,
-                rbx: data.rbx,
-                rcx: data.rcx,
-                rdx: data.rdx,
-                r8: data.r8,
-                r9: data.r9,
-                r10: data.r10,
-                r11: data.r11,
-                r12: data.r12,
-                r13: data.r13,
-                r14: data.r14,
-                r15: data.r15,
-                rip: data.rip,
-                rsp: data.rsp,
-                rbp: data.rbp,
-                rdi: data.rdi,
-                ss: data.ss,
-                exe,
-                stack_from: stack_from as u64,
-                stack: vec![],
-            };
-
-            read_mem_region(Pid::from_raw(pid as i32), "stack", &mut data.stack);
-
-            // read_mem_region(pid, "stack", &mut data.heap);
-
-            std::fs::File::create(output)
-                .expect("Failed to create output file")
-                .write_all(
-                    postcard::to_allocvec(&data)
-                        .expect("Failed to serialize data")
-                        .as_slice(),
-                )
-                .expect("Failed to write to output file");
+            dump_pid(pid, output);
         }
         Command::Restore { dump } => {
-            let dump = {
-                let mut bytes = vec![];
-                std::fs::File::open(dump)
-                    .expect("Failed to open dump file")
-                    .read_to_end(&mut bytes)
-                    .expect("Failed to read from dump file");
-
-                postcard::from_bytes::<Probe>(&bytes).expect("Invalid dump file")
-            };
-
-            let pid;
-            unsafe {
-                use libc::*;
-
-                pid = fork();
-                if pid == 0 {
-                    ptrace::traceme().expect("Failed to ptrace traceme");
-
-                    // Stop itself for first setup phase
-                    raise(SIGSTOP);
-
-                    // Execute binary
-                    execv(dump.exe.as_ptr() as *const i8, std::ptr::null_mut());
-                }
-            }
-            let pid = Pid::from_raw(pid);
-
-            // First setup phase
-
-            // Sync on stop of the child process
-            waitpid(pid, None).unwrap();
-
-            // Set option to stop execution at exec
-            ptrace::setoptions(pid, ptrace::Options::PTRACE_O_TRACEEXEC)
-                .expect("Error when setting ptrace option");
+            restore_from_dump(dump);
+        }
+    }
+}
 
-            // Continue the execution of the child
-            ptrace::cont(pid, None).unwrap();
+fn dump_pid(pid: u64, to: String) {
+    let mut data = read_regs_with_get_tasks(pid);
 
-            waitpid(pid, None).unwrap();
+    let stack_from = read_mem_region(Pid::from_raw(pid as i32), "stack", &mut data.stack);
+    data.stack_from = stack_from;
 
-            ptrace::write(pid, find_main_address(&dump.exe) as *mut libc::c_void, INT3).unwrap();
-            ptrace::cont(pid, None).unwrap();
+    File::create(to)
+        .expect("Failed to create output file")
+        .write_all(
+            postcard::to_allocvec(&data)
+                .expect("Failed to serialize data")
+                .as_slice(),
+        )
+        .expect("Failed to write to output file");
+}
 
-            let status = waitpid(pid, None).unwrap();
-            println!("{:#?}", status);
+fn restore_from_dump(dump: String) {
+    let dump = {
+        let mut bytes = vec![];
+        File::open(dump)
+            .expect("Failed to open dump file")
+            .read_to_end(&mut bytes)
+            .expect("Failed to read from dump file");
 
-            let mut regs =
-                ptrace::getregs(pid).expect("Error when retrieving child process registers");
+        postcard::from_bytes::<Probe>(&bytes).expect("Invalid dump file")
+    };
 
-            println!("{:x} -> {:x}", regs.rip, dump.rip);
+    let pid;
+    unsafe {
+        use libc::*;
 
-            regs.rax = dump.rax;
-            regs.rbx = dump.rbx;
-            regs.rcx = dump.rcx;
-            regs.rdx = dump.rdx;
-            regs.r8 = dump.r8;
-            regs.r9 = dump.r9;
-            regs.r10 = dump.r10;
-            regs.r11 = dump.r11;
-            regs.r12 = dump.r12;
-            regs.r13 = dump.r13;
-            regs.r14 = dump.r14;
-            regs.r15 = dump.r15;
-            regs.rip = dump.rip;
-            regs.rsp = dump.rsp;
-            regs.rbp = dump.rbp;
-            regs.ss = dump.ss;
-            regs.rdi = dump.rdi;
+        pid = fork();
+        if pid == 0 {
+            // ptrace::traceme().expect("Failed to ptrace traceme");
 
-            ptrace::setregs(pid, regs).expect("Error when setting registers of child process.");
+            println!("child process setup");
+            println!("exe: {:#?}", dump.exe);
 
-            println!("Write stack (from 0x{:x})", dump.stack_from);
-            write_mem_region(pid, dump.stack_from, &dump.stack);
+            // Stop itself for first setup phase
+            raise(SIGSTOP);
 
-            ptrace::cont(pid, None).unwrap();
+            // Execute binary
+            execv(dump.exe.as_ptr() as *const i8, std::ptr::null_mut());
 
-            let status = waitpid(pid, None).unwrap();
-            println!("{:#?}", status);
+            println!("child process error");
         }
     }
+    let pid = Pid::from_raw(pid);
+
+    // First setup phase
+
+    // Sync on stop of the child process
+    ptrace::attach(pid).expect("failed to seize process");
+    waitpid(pid, None).unwrap();
+
+    // Set option to stop execution at exec
+    ptrace::setoptions(pid, ptrace::Options::PTRACE_O_TRACEEXEC)
+        .expect("Error when setting ptrace option");
+
+    // Continue the execution of the child until it reaches main
+    ptrace::cont(pid, None).unwrap();
+
+    waitpid(pid, None).unwrap();
+
+    ptrace::write(pid, find_main_address(&dump.exe) as *mut libc::c_void, INT3).unwrap();
+    ptrace::cont(pid, None).unwrap();
+
+    let status = waitpid(pid, None).unwrap();
+    println!("status after child reached main: {:#?}", status);
+
+    // Restores registers of the child
+    let mut regs = ptrace::getregs(pid).expect("Error when retrieving child process registers");
+
+    regs.rax = dump.rax;
+    regs.rbx = dump.rbx;
+    regs.rcx = dump.rcx;
+    regs.rdx = dump.rdx;
+    regs.r8 = dump.r8;
+    regs.r9 = dump.r9;
+    regs.r10 = dump.r10;
+    regs.r11 = dump.r11;
+    regs.r12 = dump.r12;
+    regs.r13 = dump.r13;
+    regs.r14 = dump.r14;
+    regs.r15 = dump.r15;
+    regs.rip = dump.rip;
+    regs.rsp = dump.rsp;
+    regs.rbp = dump.rbp;
+    regs.ss = dump.ss;
+    regs.rsi = dump.rsi;
+    regs.rdi = dump.rdi;
+    regs.cs = dump.cs;
+    // TODO: find a way to get the real values in get-tasks.c or in another way
+    //regs.ds = dump.ds; regs.es = dump.es; regs.fs = dump.fs; regs.gs = dump.gs; regs.eflags = 0x202;
+
+    println!("registers restored: {:x} -> {:#?}", regs.rip, regs);
+
+    ptrace::setregs(pid, regs).expect("Error when setting registers of child process.");
+
+    // Restores stack of the child
+    println!("Write stack (from 0x{:x})", dump.stack_from);
+    write_mem_region(pid, dump.stack_from, &dump.stack);
+
+    ptrace::cont(pid, None).unwrap();
+
+    let status = waitpid(pid, None).unwrap();
+    println!("{:#?}", status);
 }
 
 fn find_main_address(exe: &str) -> i32 {
-
     // Read the binary file
     let binary = fs::read(exe).expect("Failed to read the binary file");
 
@@ -375,7 +313,7 @@ fn get_mem_region_limits(pid: Pid, region_name: &str) -> (usize, usize) {
     (region_from, region_to)
 }
 
-fn read_mem_region(pid: Pid, region_name: &str, data: &mut Vec<u8>) {
+fn read_mem_region(pid: Pid, region_name: &str, data: &mut Vec<u8>) -> u64 {
     let (region_from, region_to) = get_mem_region_limits(pid, region_name);
 
     println!("Read {} ({:x}-{:x})", region_name, region_from, region_to);
@@ -386,6 +324,8 @@ fn read_mem_region(pid: Pid, region_name: &str, data: &mut Vec<u8>) {
 
     data.resize(region_to - region_from, 0);
     mem.read_exact(data.as_mut_slice()).unwrap();
+
+    region_from as u64
 }
 
 fn write_mem_region(pid: Pid, mut from: u64, data: &Vec<u8>) {
@@ -412,4 +352,57 @@ fn write_mem_region(pid: Pid, mut from: u64, data: &Vec<u8>) {
         ptrace::write(pid, from as *mut libc::c_void, word).expect("Failed to write in memory");
         from += 8;
     }
+    println!("Writing to stack succeeded !");
+}
+
+fn read_regs_with_get_tasks(pid: u64) -> Probe {
+    let file = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .open("/dev/get_task")
+        .expect("Failed to open device");
+
+    let mut data = new_empty_raw_probe();
+    data.pid = pid;
+
+    let exe;
+
+    unsafe {
+        read_regs(file.as_raw_fd(), &mut data).expect("ioctl failed");
+
+        exe = CStr::from_ptr(data.exe.as_ptr() as *const _)
+            .to_str()
+            .unwrap()
+            .to_owned();
+    }
+
+    Probe {
+        pid: data.pid,
+        rax: data.rax,
+        rbx: data.rbx,
+        rcx: data.rcx,
+        rdx: data.rdx,
+        r8: data.r8,
+        r9: data.r9,
+        r10: data.r10,
+        r11: data.r11,
+        r12: data.r12,
+        r13: data.r13,
+        r14: data.r14,
+        r15: data.r15,
+        rip: data.rip,
+        rsp: data.rsp,
+        rbp: data.rbp,
+        ss: data.ss,
+        rsi: data.rsi,
+        rdi: data.rdi,
+        cs: data.cs,
+        ds: data.ds,
+        es: data.es,
+        fs: data.fs,
+        gs: data.gs,
+        exe,
+        stack_from: 0,
+        stack: vec![],
+    }
 }
diff --git a/get-tasks/get-tasks.c b/get-tasks/get-tasks.c
index 56bd276..3c3fab5 100644
--- a/get-tasks/get-tasks.c
+++ b/get-tasks/get-tasks.c
@@ -35,8 +35,14 @@ typedef struct {
   uint64_t rip;
   uint64_t rsp;
   uint64_t rbp;
-  uint64_t rdi;
   uint64_t ss;
+  uint64_t rsi;
+  uint64_t rdi;
+  uint64_t cs;
+  uint64_t ds;
+  uint64_t es;
+  uint64_t fs;
+  uint64_t gs;
   char exe[PATH_MAX];
 } regs_t;
 
@@ -140,8 +146,20 @@ static long etx_ioctl(struct file *file, unsigned int cmd, unsigned long arg) {
     probe.rip = regs->ip;
     probe.rsp = regs->sp;
     probe.rbp = regs->bp;
-    probe.rdi = regs->di;
     probe.ss = regs->ss;
+    probe.rsi = regs->si;
+    probe.rdi = regs->di;
+    probe.cs = regs->cs;
+
+    // TODO: find a way to get the correct values
+    //probe.ds = regs->ds;
+    probe.ds = 0;
+    //probe.es = regs->es;
+    probe.es = 0;
+    //probe.fs = regs->fs;
+    probe.fs = 0;
+    //probe.gs = regs->gs;
+    probe.gs = 0;
 
     if (copy_to_user((regs_t *)arg, &probe, sizeof(regs_t))) {
       pr_err("Failed to write results for task %d\n", 1);