Made a working ssh proxy server

12 files changed, 569 insertions, 406 deletions

diff --git a/.gitignore b/.gitignore
index fa0ede0..a8ff552 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 bastion
 libssh/
+id_rsa
 ssh_host_*_key
 test_client.c
 *.[do]
diff --git a/GNUmakefile b/GNUmakefile
index f6e0f7c..3bbefd6 100644
--- a/GNUmakefile
+++ b/GNUmakefile
@@ -1,6 +1,6 @@
 CC=clang
 DEBUG=-g
-CFLAGS= ${DEBUG} -Wall -Werror -Wextra
+CFLAGS= ${DEBUG} -Wall -Werror -Wextra -Weverything -Wno-missing-prototypes -Wno-disabled-macro-expansion
 
 sources=$(wildcard src/*.c)
 OBJ=$(sources:.c=.o)
@@ -8,7 +8,8 @@ OBJ=$(sources:.c=.o)
 all: bastion
 
 bastion: $(OBJ)
-	$(CC) ${DEBUG} -o bastion $(OBJ) -lssh -lpthread -lssh_threads -lutil
+	#$(CC) ${DEBUG} -o bastion $(OBJ) -lssh -lutil -lpthread -lssh_threads
+	$(CC) ${DEBUG} -o bastion $(OBJ) -lssh -lutil
 
 clean:
 	$(RM) bastion *.[do] src/*.[do]
@@ -22,5 +23,5 @@ clean:
 	@rm -f $*.d.tmp
 
 valgrind:
-	valgrind --leak-check=full --show-leak-kinds=all --suppressions=$HOME/.valgrind_suppressions ./bastion
-	#valgrind -v --leak-check=full --show-leak-kinds=all --suppressions=$HOME/.valgrind_suppressions --gen-suppressions=yes ./bastion
+	valgrind --leak-check=full --show-leak-kinds=all --trace-children=yes --suppressions=${HOME}/.valgrind_suppressions ./bastion
+	#valgrind -v --leak-check=full --show-leak-kinds=all --trace-children=yes --suppressions=${HOME}/.valgrind_suppressions --gen-suppressions=yes ./bastion
diff --git a/src/client.c b/src/client.c
new file mode 100644
index 0000000..64df1eb
--- /dev/null
+++ b/src/client.c
@@ -0,0 +1,187 @@
+#include <libssh/callbacks.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "../config.h"
+#include "client.h"
+
+// callback function for channel data and exceptions
+static int client_data_function(ssh_session session, ssh_channel channel, void *data,
+                         uint32_t len, int is_stderr, void *userdata) {
+    struct client_channel_data_struct *cdata = (struct client_channel_data_struct *) userdata;
+    (void) session;
+    (void) channel;
+    (void) is_stderr;
+
+    if (ssh_channel_is_open(cdata->proxy_channel))
+        return ssh_channel_write(cdata->proxy_channel, (char*) data, len);
+    else
+        return SSH_ERROR;
+}
+
+static void client_channel_eof_callback (ssh_session session, ssh_channel channel, void *userdata)
+{
+    struct client_channel_data_struct *cdata = (struct client_channel_data_struct *) userdata;
+    (void) session;
+    (void) channel;
+
+    if (ssh_channel_is_open(cdata->proxy_channel))
+        ssh_channel_send_eof(cdata->proxy_channel);
+}
+
+static void client_channel_close_callback (ssh_session session, ssh_channel channel, void *userdata)
+{
+    struct client_channel_data_struct *cdata = (struct client_channel_data_struct *) userdata;
+    (void) session;
+    (void) channel;
+
+    if (ssh_channel_is_open(cdata->proxy_channel))
+        ssh_channel_close(cdata->proxy_channel);
+}
+
+static void client_channel_exit_status_callback (ssh_session session, ssh_channel channel, int exit_status, void *userdata)
+{
+    (void) session;
+    (void) channel;
+    (void) userdata;
+    printf("client exit status callback %d\n", exit_status);
+}
+
+static void client_channel_signal_callback (ssh_session session, ssh_channel channel,
+                                    const char *signal, void *userdata) {
+    (void) session;
+    (void) channel;
+    (void) signal;
+    (void) userdata;
+    printf("client signal callback\n");
+}
+
+static void client_channel_exit_signal_callback (ssh_session session, ssh_channel channel,
+                                    const char *signal, int core, const char *errmsg,
+                                    const char *lang, void *userdata) {
+    (void) session;
+    (void) channel;
+    (void) signal;
+    (void) core;
+    (void) errmsg;
+    (void) lang;
+    (void) userdata;
+    printf("client exit signal callback\n");
+}
+
+struct client_channel_data_struct* client_dial(ssh_event event, struct proxy_channel_data_struct *pdata, const char * hostname)
+{
+    struct client_channel_data_struct *cdata = malloc(sizeof(*cdata));
+    cdata->event = event;
+    cdata->my_session = NULL;
+    cdata->my_channel = NULL;
+    cdata->proxy_channel = pdata->my_channel;
+    cdata->client_channel_cb = NULL;
+
+    /* First we try to add the private key that the server will accept */
+    ssh_key privkey;
+    if (ssh_pki_import_privkey_file(PRIVKEY_PATH, NULL, NULL, NULL, &privkey) != SSH_OK) {
+        printf("Error importing private key");
+        goto privkey_clean;
+    }
+
+    /* We try to connect to the remote server */
+    printf("Connecting to %s\n", hostname);
+    cdata->my_session = ssh_new();
+
+    ssh_options_set(cdata->my_session, SSH_OPTIONS_HOST, hostname);
+    ssh_options_set(cdata->my_session, SSH_OPTIONS_USER, USER_TO_LOGIN_AS);
+    int verbosity = SSH_LOG_PROTOCOL;
+    ssh_options_set(cdata->my_session, SSH_OPTIONS_LOG_VERBOSITY, &verbosity);
+
+    if (ssh_connect(cdata->my_session) != SSH_OK) {
+        printf("Error connecting to %s: %s\n", hostname, ssh_get_error(cdata->my_session));
+        goto session_clean;
+    }
+
+    /* We now validate the remote server's public key */
+    ssh_key server_pub_key;
+    unsigned char * hash = NULL;
+    size_t hlen;
+    char * hexa = NULL;
+    if (ssh_get_server_publickey(cdata->my_session, &server_pub_key) != SSH_OK) {
+        fprintf(stderr, "Error getting server publickey: %s\n", ssh_get_error(cdata->my_session));
+        goto pubkey_clean;
+    }
+    if (ssh_get_publickey_hash(server_pub_key, SSH_PUBLICKEY_HASH_SHA1, &hash, &hlen) != SSH_OK) {
+        fprintf(stderr, "Error getting publickey hash: %s\n", ssh_get_error(cdata->my_session));
+        goto pubkey_hash_clean;
+    }
+    hexa = ssh_get_hexa(hash, hlen);
+    printf("Server public key hash : %s\n", hexa); // TODO validate the key
+    // if error goto pubkey_nomatch_clean
+
+    /* With the server checked, we can authenticate */
+    if(ssh_userauth_publickey(cdata->my_session, NULL, privkey) == SSH_AUTH_SUCCESS){
+        printf("Authentication success\n");
+    } else {
+        printf("Error private key was rejected\n");
+        goto auth_clean;
+    }
+
+    /* we open the client channel */
+    cdata->my_channel = ssh_channel_new(cdata->my_session);
+    if (cdata->my_channel == NULL) {
+        printf("Couldn't open client channel to %s\n", hostname);
+        goto channel_clean;
+    }
+
+    /* we open a session channel for the future shell, not suitable for tcp
+     * forwarding */
+    if (ssh_channel_open_session(cdata->my_channel) != SSH_OK) {
+        printf("Couldn't open the session channel\n");
+        goto channel_clean;
+    }
+
+    cdata->client_channel_cb = malloc(sizeof(*cdata->client_channel_cb));
+    cdata->client_channel_cb->userdata = cdata;
+    cdata->client_channel_cb->channel_data_function = client_data_function;
+    cdata->client_channel_cb->channel_eof_function = client_channel_eof_callback;
+    cdata->client_channel_cb->channel_close_function = client_channel_close_callback;
+    cdata->client_channel_cb->channel_exit_status_function = client_channel_exit_status_callback;
+    cdata->client_channel_cb->channel_signal_function = client_channel_signal_callback;
+    cdata->client_channel_cb->channel_exit_signal_function = client_channel_exit_signal_callback;
+
+    ssh_callbacks_init(cdata->client_channel_cb);
+    ssh_set_channel_callbacks(cdata->my_channel, cdata->client_channel_cb);
+    ssh_event_add_session(event, cdata->my_session);
+
+    ssh_string_free_char(hexa);
+    ssh_clean_pubkey_hash(&hash);
+    ssh_key_free(server_pub_key);
+    ssh_key_free(privkey);
+    return cdata;
+
+channel_clean:
+    ssh_channel_free(cdata->my_channel);
+auth_clean:
+    // TODO
+//pubkey_nomatch_clean:
+    ssh_string_free_char(hexa);
+pubkey_hash_clean:
+    ssh_clean_pubkey_hash(&hash);
+pubkey_clean:
+    ssh_key_free(server_pub_key);
+session_clean:
+    ssh_disconnect(cdata->my_session);
+    ssh_free(cdata->my_session);
+privkey_clean:
+    ssh_key_free(privkey);
+    free(cdata);
+    return NULL;
+}
+
+void client_cleanup(struct client_channel_data_struct *cdata)
+{
+    ssh_event_remove_session(cdata->event, cdata->my_session);
+    ssh_channel_free(cdata->my_channel);
+    ssh_disconnect(cdata->my_session);
+    ssh_free(cdata->my_session);
+    free(cdata->client_channel_cb);
+    free(cdata);
+}
diff --git a/src/client.h b/src/client.h
new file mode 100644
index 0000000..3826b5f
--- /dev/null
+++ b/src/client.h
@@ -0,0 +1,22 @@
+#ifndef CLIENT_H_
+#define CLIENT_H_
+
+#include <libssh/libssh.h>
+
+#include "proxy.h"
+#include "session.h"
+
+/* A userdata struct for channel. */
+struct client_channel_data_struct {
+    /* Event which is used to poll */
+    ssh_event event;
+    ssh_session my_session;
+    ssh_channel my_channel;
+    ssh_channel proxy_channel;
+    struct ssh_channel_callbacks_struct * client_channel_cb;
+};
+
+struct client_channel_data_struct* client_dial(ssh_event event, struct proxy_channel_data_struct *pdata, const char * hostname);
+void client_cleanup(struct client_channel_data_struct *cdata);
+
+#endif
diff --git a/src/main.c b/src/main.c
index cd8e74f..0c6b6cb 100644
--- a/src/main.c
+++ b/src/main.c
@@ -1,15 +1,10 @@
 #include <libssh/callbacks.h>
 #include <libssh/server.h>
-//#include <fcntl.h>
-//#include <libutil.h>
-//#include <poll.h>
-//#include <signal.h>
 #include <stdio.h>
 #include <stdlib.h>
-//#include <sys/ioctl.h>
 #include <sys/wait.h>
-//#include <util.h>
 
+#include "../config.h"
 #include "session.h"
 
 /* SIGCHLD handler for cleaning up dead children. */
@@ -18,6 +13,19 @@ static void sigchld_handler(int signo) {
     while (waitpid(-1, NULL, WNOHANG) > 0);
 }
 
+/* SIGINT handler for cleaning up on forced exit. */
+static ssh_bind sshbind;
+static ssh_session session;
+
+__attribute__((noreturn)) void sigint_handler(int signo)
+{
+    (void) signo;
+    ssh_free(session);
+    ssh_bind_free(sshbind);
+    ssh_finalize();
+    exit(0);
+}
+
 int main()
 {
     // Set up SIGCHLD handler
@@ -29,30 +37,40 @@ int main()
         fprintf(stderr, "Failed to register SIGCHLD handler\n");
         return 1;
     }
+    // Set up SIGINT handler
+    struct sigaction sa2;
+    sa2.sa_handler = sigint_handler;
+    sigemptyset(&sa2.sa_mask);
+    sa2.sa_flags = 0;
+    if (sigaction(SIGINT, &sa2, NULL) != 0) {
+        fprintf(stderr, "Failed to register SIGINT handler\n");
+        return 1;
+    }
 
     // Initializing ssh context
-    ssh_threads_set_callbacks(ssh_threads_get_pthread());
     ssh_init();
 
     // Initializing ssh_bind
-    ssh_bind sshbind = ssh_bind_new();
+    sshbind = ssh_bind_new();
     if (sshbind == NULL) {
         fprintf(stderr, "Error initializing ssh_bind\n");
         exit(-1);
     }
     int port = 2222;
     ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_BINDPORT, &port);
-    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_DSAKEY, "ssh_host_dsa_key");
-    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_RSAKEY, "ssh_host_rsa_key");
-    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_ECDSAKEY, "ssh_host_ecdsa_key");
+    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_DSAKEY, DSAKEY_PATH);
+    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_RSAKEY, RSAKEY_PATH);
+    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_ECDSAKEY, ECDSAKEY_PATH);
 
     if (ssh_bind_listen(sshbind) < 0) {
         printf("Error listening to socket: %s\n", ssh_get_error(sshbind));
+        ssh_bind_free(sshbind);
+        ssh_finalize();
         return 1;
     }
 
     while (1) {
-        ssh_session session = ssh_new();
+        session = ssh_new();
         if (session == NULL) {
             fprintf(stderr, "Error initializing ssh_session\n");
             break;
@@ -67,14 +85,12 @@ int main()
                 /* Remove the SIGCHLD handler inherited from parent. */
                 sa.sa_handler = SIG_DFL;
                 sigaction(SIGCHLD, &sa, NULL);
-                /* Remove socket binding, which allows us to restart the
-                 * parent process, without terminating existing sessions. */
+                /* Remove socket binding, which allows us to restart the parent process, without terminating existing sessions. */
                 ssh_bind_free(sshbind);
 
                 ssh_event event = ssh_event_new();
                 if (event != NULL) {
-                    /* Blocks until the SSH session ends by either
-                     * child process exiting, or client disconnecting. */
+                    /* Blocks until the SSH session ends */
                     handle_session(event, session);
                     ssh_event_free(event);
                 } else {
@@ -82,21 +98,24 @@ int main()
                 }
                 ssh_disconnect(session);
                 ssh_free(session);
+                ssh_finalize();
 
-                exit(0);
+                return 0;
               case -1:
                 fprintf(stderr, "Failed to fork\n");
             }
         } else {
             fprintf(stderr, "Error accepting a connection : %s\n", ssh_get_error(sshbind));
-            exit(1);
+            ssh_disconnect(session);
+            ssh_free(session);
+            ssh_bind_free(sshbind);
+            ssh_finalize();
+            return 1;
         }
-        /* Since the session has been passed to a child fork, do some cleaning
-         * up at the parent process. */
+        /* Since the session has been passed to a child fork, do some cleaning up at the parent process. */
         ssh_disconnect(session);
         ssh_free(session);
     }
-
     ssh_bind_free(sshbind);
     ssh_finalize();
     return 0;
diff --git a/src/proxy.c b/src/proxy.c
new file mode 100644
index 0000000..7d3290e
--- /dev/null
+++ b/src/proxy.c
@@ -0,0 +1,230 @@
+#include <libssh/callbacks.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "client.h"
+#include "proxy.h"
+
+// callback function for channel data and exceptions
+static int proxy_data_function(ssh_session session, ssh_channel channel, void *data,
+                         uint32_t len, int is_stderr, void *userdata) {
+    struct proxy_channel_data_struct *pdata = (struct proxy_channel_data_struct *) userdata;
+    (void) session;
+    (void) channel;
+    (void) is_stderr;
+
+    if (ssh_channel_is_open(pdata->client_channel))
+        return ssh_channel_write(pdata->client_channel, (char*) data, len);
+    else
+        return SSH_ERROR;
+}
+
+// callback function for SSH channel PTY request from a client
+static int proxy_pty_request(ssh_session session, ssh_channel channel,
+                const char *term, int cols, int rows, int py, int px,
+                void *userdata) {
+    struct proxy_channel_data_struct *pdata = (struct proxy_channel_data_struct *)userdata;
+
+    (void) session;
+    (void) channel;
+    (void) py;
+    (void) px;
+
+    printf("pty request\n");
+    int rc = ssh_channel_request_pty_size(pdata->client_channel, term, cols, rows);
+    if (rc == SSH_OK) {
+        printf("pty request successfull\n");
+    } else {
+        printf("pty request failed\n");
+    }
+    return rc;
+}
+
+// callback function for SSH channel PTY resize from a client
+static int proxy_pty_resize(ssh_session session, ssh_channel channel, int cols,
+               int rows, int py, int px, void *userdata) {
+    struct proxy_channel_data_struct *pdata = (struct proxy_channel_data_struct *)userdata;
+
+    (void) session;
+    (void) channel;
+    (void) py;
+    (void) px;
+
+    if (pdata->client_channel == NULL || ssh_channel_is_open(pdata->client_channel) == 0) {
+        fprintf(stderr, "proxy pty oups!!!!!\n");
+        return SSH_ERROR;
+    }
+    printf("pty resize\n");
+    int rc = ssh_channel_change_pty_size(pdata->client_channel, cols, rows);
+    if (rc == SSH_OK) {
+        printf("pty resize successfull\n");
+    } else {
+        printf("pty resize failed\n");
+    }
+    return rc;
+}
+
+static int proxy_exec_request(ssh_session session, ssh_channel channel,
+                        const char *command, void *userdata) {
+    struct proxy_channel_data_struct *pdata = (struct proxy_channel_data_struct *) userdata;
+
+    (void) session;
+    (void) channel;
+
+    printf("exec request : %s\n", command); // TODO
+    int rc = ssh_channel_request_exec(pdata->client_channel, command);
+    if (rc == SSH_OK) {
+        printf("exec request successfull\n");
+    } else {
+        printf("exec request failed\n");
+    }
+    return rc;
+}
+
+static int proxy_shell_request(ssh_session session, ssh_channel channel,
+                         void *userdata) {
+    struct proxy_channel_data_struct *pdata = (struct proxy_channel_data_struct *) userdata;
+
+    (void) session;
+    (void) channel;
+
+    printf("shell request\n");
+    int rc = ssh_channel_request_shell(pdata->client_channel);
+    if (rc == SSH_OK) {
+        printf("shell request successfull\n");
+    } else {
+        printf("shell request failed\n");
+    }
+    return rc;
+}
+
+static int proxy_subsystem_request(ssh_session session, ssh_channel channel,
+                             const char *subsystem, void *userdata) {
+    ///* subsystem requests behave simillarly to exec requests. */
+    //if (strcmp(subsystem, "sftp") == 0) {
+    //    printf("sftp request\n"); // TODO
+    //    return exec_request(session, channel, SFTP_SERVER_PATH, userdata);
+    //}
+    (void) session;
+    (void) channel;
+    (void) subsystem;
+    (void) userdata;
+    return SSH_ERROR; // TODO
+}
+
+static void proxy_channel_eof_callback (ssh_session session, ssh_channel channel, void *userdata)
+{
+    (void) session;
+    (void) channel;
+    (void) userdata;
+    printf("proxy eof callback\n");
+}
+
+static void proxy_channel_close_callback (ssh_session session, ssh_channel channel, void *userdata)
+{
+    (void) session;
+    (void) channel;
+    (void) userdata;
+    printf("proxy close callback\n");
+}
+
+static void proxy_channel_exit_status_callback (ssh_session session, ssh_channel channel, int exit_status, void *userdata)
+{
+    (void) session;
+    (void) channel;
+    (void) exit_status;
+    (void) userdata;
+    printf("proxy exit status callback\n");
+}
+
+static void proxy_channel_signal_callback (ssh_session session, ssh_channel channel,
+                                    const char *signal, void *userdata) {
+    (void) session;
+    (void) channel;
+    (void) signal;
+    (void) userdata;
+    printf("proxy signal callback\n");
+}
+
+static void proxy_channel_exit_signal_callback (ssh_session session, ssh_channel channel,
+                                    const char *signal, int core, const char *errmsg,
+                                    const char *lang, void *userdata) {
+    (void) session;
+    (void) channel;
+    (void) signal;
+    (void) core;
+    (void) errmsg;
+    (void) lang;
+    (void) userdata;
+    printf("proxy exit signal callback\n");
+}
+
+void handle_proxy_session(ssh_event event, ssh_session session, ssh_channel my_channel, const char * hostname)
+{
+    struct client_channel_data_struct * cdata;
+
+    struct proxy_channel_data_struct pdata = {
+        .event = event,
+        .my_session = session,
+        .my_channel = my_channel,
+        .client_channel = NULL,
+    };
+
+    //ssh_event_remove_session(event, session);
+    cdata = client_dial(event, &pdata, hostname);
+    //for (int n = 0; n < 10; n++) {
+    //    ssh_event_dopoll(event, 100);
+    //}
+    //ssh_event_add_session(event, session);
+
+    if (cdata == NULL) {
+        return;
+    }
+    pdata.client_channel = cdata->my_channel;
+
+    /* We tie everything together */
+    struct ssh_channel_callbacks_struct channel_cb = {
+        .userdata = &pdata,
+        .channel_data_function = proxy_data_function,
+        .channel_eof_function = proxy_channel_eof_callback,
+        .channel_close_function = proxy_channel_close_callback,
+        .channel_signal_function = proxy_channel_signal_callback,
+        .channel_exit_status_function = proxy_channel_exit_status_callback,
+        .channel_exit_signal_function = proxy_channel_exit_signal_callback,
+        .channel_pty_request_function = proxy_pty_request,
+        .channel_shell_request_function = proxy_shell_request,
+        .channel_pty_window_change_function = proxy_pty_resize,
+        .channel_exec_request_function = proxy_exec_request,
+        .channel_subsystem_request_function = proxy_subsystem_request,
+  /** This function will be called when a client requests agent
+   * authentication forwarding.
+   */
+  //ssh_channel_auth_agent_req_callback channel_auth_agent_req_function;
+  /** This function will be called when a client requests X11
+   * forwarding.
+   */
+  //ssh_channel_x11_req_callback channel_x11_req_function;
+  /** This function will be called when a client requests an environment
+   * variable to be set.
+   */
+  /** This function will be called when the channel write is guaranteed
+   * not to block.
+   */
+  //    .channel_write_wontblock_function = proxy_channel_write_wontblock,
+    };
+    ssh_callbacks_init(&channel_cb);
+    ssh_set_channel_callbacks(my_channel, &channel_cb);
+
+    do {
+        /* Poll the main event which takes care of the sessions and channels */
+        if (ssh_event_dopoll(event, -1) == SSH_ERROR) {
+            break;
+        }
+    } while(ssh_channel_is_open(my_channel) && ssh_channel_is_open(pdata.client_channel));
+    if (ssh_channel_is_open(my_channel))
+        ssh_channel_close(my_channel);
+    if (ssh_channel_is_open(cdata->my_channel))
+        ssh_channel_close(cdata->my_channel);
+
+    client_cleanup(cdata);
+}
diff --git a/src/proxy.h b/src/proxy.h
new file mode 100644
index 0000000..3756544
--- /dev/null
+++ b/src/proxy.h
@@ -0,0 +1,18 @@
+#ifndef PROXY_H_
+#define PROXY_H_
+
+#include <libssh/libssh.h>
+
+#include "session.h"
+
+/* A userdata struct for channel. */
+struct proxy_channel_data_struct {
+    /* Event which is used to poll */
+    ssh_event event;
+    ssh_session my_session;
+    ssh_channel my_channel;
+    ssh_channel client_channel;
+};
+void handle_proxy_session(ssh_event event, ssh_session session, ssh_channel my_channel, const char * hostname);
+
+#endif
diff --git a/src/pty.c b/src/pty.c
deleted file mode 100644
index c888df3..0000000
--- a/src/pty.c
+++ /dev/null
@@ -1,128 +0,0 @@
-#include <pty.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <utmp.h>
-
-#include "pty.h"
-
-// callback function for SSH channel PTY request from a client
-int pty_request(ssh_session session, ssh_channel channel,
-                const char *term, int cols, int rows, int py, int px,
-                void *userdata) {
-    struct channel_data_struct *cdata = (struct channel_data_struct *)userdata;
-
-    (void) session;
-    (void) channel;
-    (void) term;
-
-    cdata->winsize->ws_row = rows;
-    cdata->winsize->ws_col = cols;
-    cdata->winsize->ws_xpixel = px;
-    cdata->winsize->ws_ypixel = py;
-
-    if (openpty(&cdata->pty_master, &cdata->pty_slave, NULL, NULL,
-                cdata->winsize) != 0) {
-        fprintf(stderr, "Failed to open pty\n");
-        return SSH_ERROR;
-    }
-    return SSH_OK;
-}
-
-// callback function for SSH channel PTY resize from a client
-int pty_resize(ssh_session session, ssh_channel channel, int cols,
-               int rows, int py, int px, void *userdata) {
-    struct channel_data_struct *cdata = (struct channel_data_struct *)userdata;
-
-    (void) session;
-    (void) channel;
-
-    cdata->winsize->ws_row = rows;
-    cdata->winsize->ws_col = cols;
-    cdata->winsize->ws_xpixel = px;
-    cdata->winsize->ws_ypixel = py;
-
-    if (cdata->pty_master != -1) {
-        return ioctl(cdata->pty_master, TIOCSWINSZ, cdata->winsize);
-    }
-
-    return SSH_ERROR;
-}
-
-int exec_pty(const char *mode, const char *command,
-             struct channel_data_struct *cdata) {
-    switch(cdata->pid = fork()) {
-        case -1:
-            close(cdata->pty_master);
-            close(cdata->pty_slave);
-            fprintf(stderr, "Failed to fork\n");
-            return SSH_ERROR;
-        case 0:
-            close(cdata->pty_master);
-            if (login_tty(cdata->pty_slave) != 0) {
-                exit(1);
-            }
-            execl("/bin/sh", "sh", mode, command, NULL);
-            exit(0);
-        default:
-            close(cdata->pty_slave);
-            /* pty fd is bi-directional */
-            cdata->child_stdout = cdata->child_stdin = cdata->pty_master;
-    }
-    return SSH_OK;
-}
-
-int exec_nopty(const char *command, struct channel_data_struct *cdata) {
-    int in[2], out[2], err[2];
-
-    /* Do the plumbing to be able to talk with the child process. */
-    if (pipe(in) != 0) {
-        goto stdin_failed;
-    }
-    if (pipe(out) != 0) {
-        goto stdout_failed;
-    }
-    if (pipe(err) != 0) {
-        goto stderr_failed;
-    }
-
-    switch(cdata->pid = fork()) {
-        case -1:
-            goto fork_failed;
-        case 0:
-            /* Finish the plumbing in the child process. */
-            close(in[1]);
-            close(out[0]);
-            close(err[0]);
-            dup2(in[0], STDIN_FILENO);
-            dup2(out[1], STDOUT_FILENO);
-            dup2(err[1], STDERR_FILENO);
-            close(in[0]);
-            close(out[1]);
-            close(err[1]);
-            /* exec the requested command. */
-            execl("/bin/sh", "sh", "-c", command, NULL);
-            exit(0);
-    }
-
-    close(in[0]);
-    close(out[1]);
-    close(err[1]);
-
-    cdata->child_stdin = in[1];
-    cdata->child_stdout = out[0];
-    cdata->child_stderr = err[0];
-
-    return SSH_OK;
-
-fork_failed:
-    close(err[0]);
-    close(err[1]);
-stderr_failed:
-    close(out[0]);
-    close(out[1]);
-stdout_failed:
-    close(in[0]);
-    close(in[1]);
-stdin_failed:
-    return SSH_ERROR;
-}
diff --git a/src/pty.h b/src/pty.h
deleted file mode 100644
index c130f4a..0000000
--- a/src/pty.h
+++ /dev/null
@@ -1,15 +0,0 @@
-#ifndef _PTY_H_
-#define _PTY_H_
-
-#include "util.h"
-
-int pty_request(ssh_session session, ssh_channel channel,
-                const char *term, int cols, int rows, int py, int px,
-                void *userdata);
-int pty_resize(ssh_session session, ssh_channel channel, int cols,
-               int rows, int py, int px, void *userdata);
-int exec_pty(const char *mode, const char *command,
-             struct channel_data_struct *cdata);
-int exec_nopty(const char *command, struct channel_data_struct *cdata);
-
-#endif
diff --git a/src/session.c b/src/session.c
index 6186e6c..90e2855 100644
--- a/src/session.c
+++ b/src/session.c
@@ -8,178 +8,76 @@
 #include <string.h>
 #include <sys/wait.h>
 
-#include "pty.h"
+#include "../config.h"
+#include "proxy.h"
 #include "session.h"
 
-#define USER "julien"
-#define PASS "graou"
-#define BUF_SIZE 1048576
-#define SESSION_END (SSH_CLOSED | SSH_CLOSED_ERROR)
-#define SFTP_SERVER_PATH "/usr/lib/sftp-server"
-
-// callback function for channel data and exceptions
-static int data_function(ssh_session session, ssh_channel channel, void *data,
-                         uint32_t len, int is_stderr, void *userdata) {
-    struct channel_data_struct *cdata = (struct channel_data_struct *) userdata;
-
-    (void) session;
-    (void) channel;
-    (void) is_stderr;
-
-    if (len == 0 || cdata->pid < 1 || kill(cdata->pid, 0) < 0) {
-        return 0;
-    }
-
-    return write(cdata->child_stdin, (char *) data, len);
-}
-
-static int exec_request(ssh_session session, ssh_channel channel,
-                        const char *command, void *userdata) {
-    struct channel_data_struct *cdata = (struct channel_data_struct *) userdata;
-
-
-    (void) session;
-    (void) channel;
-
-    if(cdata->pid > 0) {
-        return SSH_ERROR;
-    }
-
-    if (cdata->pty_master != -1 && cdata->pty_slave != -1) {
-        return exec_pty("-c", command, cdata);
-    }
-    return exec_nopty(command, cdata);
-}
-
-static int shell_request(ssh_session session, ssh_channel channel,
-                         void *userdata) {
-    struct channel_data_struct *cdata = (struct channel_data_struct *) userdata;
-
+int auth_pubkey(ssh_session session, const char *user,
+                       struct ssh_key_struct *pubkey,
+                       char signature_state, void *userdata) {
+    struct session_data_struct *sdata = (struct session_data_struct *) userdata;
     (void) session;
-    (void) channel;
 
-    if(cdata->pid > 0) {
+    // For some reason, libssh can call this twice for the same key
+    if (sdata->authenticated == 1)
         return SSH_ERROR;
-    }
 
-    if (cdata->pty_master != -1 && cdata->pty_slave != -1) {
-        return exec_pty("-l", NULL, cdata);
+    if (signature_state != SSH_PUBLICKEY_STATE_NONE && signature_state != SSH_PUBLICKEY_STATE_VALID) {
+        fprintf(stderr, "Invalid signature state\n");
+        sdata->auth_attempts++;
+        return SSH_AUTH_DENIED;
     }
-    /* Client requested a shell without a pty, let's pretend we allow that */
-    return SSH_OK;
-}
 
-static int subsystem_request(ssh_session session, ssh_channel channel,
-                             const char *subsystem, void *userdata) {
-    /* subsystem requests behave simillarly to exec requests. */
-    if (strcmp(subsystem, "sftp") == 0) {
-        return exec_request(session, channel, SFTP_SERVER_PATH, userdata);
-    }
-    return SSH_ERROR;
-}
+    // TODO check for an invite
 
-static int auth_password(ssh_session session, const char *user,
-                         const char *pass, void *userdata) {
-    struct session_data_struct *sdata = (struct session_data_struct *) userdata;
-
-    (void) session;
-
-    if (strcmp(user, USER) == 0 && strcmp(pass, PASS) == 0) {
+    ssh_key reference_key = ssh_key_new();
+    ssh_pki_import_pubkey_base64(USER_RSA_PUBKEY, SSH_KEYTYPE_RSA, &reference_key); // TODO fetch all pubkeys from db
+    if (!ssh_key_cmp(pubkey, reference_key, SSH_KEY_CMP_PUBLIC)) {
         sdata->authenticated = 1;
+        ssh_key_free(reference_key);
+        size_t len = strnlen(user, MAX_HOSTNAME_LENGTH + 1);
+        if (len == MAX_HOSTNAME_LENGTH + 1) {
+            fprintf(stderr, "Hostname too long, max length is %d.\n", MAX_HOSTNAME_LENGTH);
+            return SSH_ERROR;
+        }
+        sdata->login_username = malloc(len+1);
+        memset(sdata->login_username, 0, len+1);
+        strncpy(sdata->login_username, user, len);
         return SSH_AUTH_SUCCESS;
+    } else {
+        ssh_key_free(reference_key);
+        sdata->auth_attempts++;
+        return SSH_AUTH_DENIED;
     }
-
-    sdata->auth_attempts++;
-    return SSH_AUTH_DENIED;
 }
 
-static ssh_channel channel_open(ssh_session session, void *userdata) {
+ssh_channel channel_open(ssh_session session, void *userdata) {
     struct session_data_struct *sdata = (struct session_data_struct *) userdata;
 
-    sdata->channel = ssh_channel_new(session);
-    return sdata->channel;
-}
-
-static int process_stdout(socket_t fd, int revents, void *userdata) {
-    char buf[BUF_SIZE];
-    int n = -1;
-    ssh_channel channel = (ssh_channel) userdata;
-
-    if (channel != NULL && (revents & POLLIN) != 0) {
-        n = read(fd, buf, BUF_SIZE);
-        if (n > 0) {
-            ssh_channel_write(channel, buf, n);
-            // TODO implement logging here
-        }
-    }
-
-    return n;
-}
-
-static int process_stderr(socket_t fd, int revents, void *userdata) {
-    char buf[BUF_SIZE];
-    int n = -1;
-    ssh_channel channel = (ssh_channel) userdata;
-
-    if (channel != NULL && (revents & POLLIN) != 0) {
-        n = read(fd, buf, BUF_SIZE);
-        if (n > 0) {
-            ssh_channel_write_stderr(channel, buf, n);
-        }
+    if (sdata->channel == NULL) {
+        sdata->channel = ssh_channel_new(session);
+        return sdata->channel;
+    } else {
+        // Only one channel allowed
+        return NULL;
     }
-
-    return n;
 }
 
 void handle_session(ssh_event event, ssh_session session) {
-    int n, rc;
-
-    /* Structure for storing the pty size. */
-    struct winsize wsize = {
-        .ws_row = 0,
-        .ws_col = 0,
-        .ws_xpixel = 0,
-        .ws_ypixel = 0
-    };
-
-    /* Our struct holding information about the channel. */
-    struct channel_data_struct cdata = {
-        .pid = 0,
-        .pty_master = -1,
-        .pty_slave = -1,
-        .child_stdin = -1,
-        .child_stdout = -1,
-        .child_stderr = -1,
-        .event = NULL,
-        .winsize = &wsize
-    };
-
     /* Our struct holding information about the session. */
     struct session_data_struct sdata = {
         .channel = NULL,
         .auth_attempts = 0,
-        .authenticated = 0
-    };
-
-    struct ssh_channel_callbacks_struct channel_cb = {
-        .userdata = &cdata,
-        .channel_pty_request_function = pty_request,
-        .channel_pty_window_change_function = pty_resize,
-        .channel_shell_request_function = shell_request,
-        .channel_exec_request_function = exec_request,
-        .channel_data_function = data_function,
-        .channel_subsystem_request_function = subsystem_request
+        .authenticated = 0,
+        .login_username = NULL
     };
 
     struct ssh_server_callbacks_struct server_cb = {
         .userdata = &sdata,
-        .auth_password_function = auth_password,
+        .auth_pubkey_function = auth_pubkey,
         .channel_open_request_session_function = channel_open,
     };
-
     ssh_callbacks_init(&server_cb);
-    ssh_callbacks_init(&channel_cb);
-
     ssh_set_server_callbacks(session, &server_cb);
 
     if (ssh_handle_key_exchange(session) != SSH_OK) {
@@ -187,14 +85,18 @@ void handle_session(ssh_event event, ssh_session session) {
         return;
     }
 
-    ssh_set_auth_methods(session, SSH_AUTH_METHOD_PASSWORD);
+    ssh_set_auth_methods(session, SSH_AUTH_METHOD_PUBLICKEY);
     ssh_event_add_session(event, session);
 
-    n = 0;
-    while (sdata.authenticated == 0 || sdata.channel == NULL) {
+    for (int n=0; sdata.authenticated == 0 || sdata.channel == NULL; n++) {
         /* If the user has used up all attempts, or if he hasn't been able to
          * authenticate in 10 seconds (n * 100ms), disconnect. */
-        if (sdata.auth_attempts >= 3 || n >= 100) {
+        if (sdata.auth_attempts >= 3) {
+            fprintf(stderr, "Closing connection after 3 failed auth attempts\n");
+            return;
+        }
+        if (n >= 100) {
+            fprintf(stderr, "Closing connection after 10 seconds without successfull authentication\n");
             return;
         }
 
@@ -202,70 +104,18 @@ void handle_session(ssh_event event, ssh_session session) {
             fprintf(stderr, "%s\n", ssh_get_error(session));
             return;
         }
-        n++;
     }
 
-    ssh_set_channel_callbacks(sdata.channel, &channel_cb);
-
-    do {
-        /* Poll the main event which takes care of the session, the channel and
-         * even our child process's stdout/stderr (once it's started). */
-        if (ssh_event_dopoll(event, -1) == SSH_ERROR) {
-          ssh_channel_close(sdata.channel);
-        }
+    handle_proxy_session(event, session, sdata.channel, sdata.login_username);
 
-        /* If child process's stdout/stderr has been registered with the event,
-         * or the child process hasn't started yet, continue. */
-        if (cdata.event != NULL || cdata.pid == 0) {
-            continue;
-        }
-        /* Executed only once, once the child process starts. */
-        cdata.event = event;
-        /* If stdout valid, add stdout to be monitored by the poll event. */
-        if (cdata.child_stdout != -1) {
-            if (ssh_event_add_fd(event, cdata.child_stdout, POLLIN, process_stdout,
-                                 sdata.channel) != SSH_OK) {
-                fprintf(stderr, "Failed to register stdout to poll context\n");
-                ssh_channel_close(sdata.channel);
-            }
-        }
-
-        /* If stderr valid, add stderr to be monitored by the poll event. */
-        if (cdata.child_stderr != -1){
-            if (ssh_event_add_fd(event, cdata.child_stderr, POLLIN, process_stderr,
-                                 sdata.channel) != SSH_OK) {
-                fprintf(stderr, "Failed to register stderr to poll context\n");
-                ssh_channel_close(sdata.channel);
-            }
-        }
-    } while(ssh_channel_is_open(sdata.channel) &&
-            (cdata.pid == 0 || waitpid(cdata.pid, &rc, WNOHANG) == 0));
-
-    close(cdata.pty_master);
-    close(cdata.child_stdin);
-    close(cdata.child_stdout);
-    close(cdata.child_stderr);
-
-    /* Remove the descriptors from the polling context, since they are now
-     * closed, they will always trigger during the poll calls. */
-    ssh_event_remove_fd(event, cdata.child_stdout);
-    ssh_event_remove_fd(event, cdata.child_stderr);
-
-    /* If the child process exited. */
-    if (kill(cdata.pid, 0) < 0 && WIFEXITED(rc)) {
-        rc = WEXITSTATUS(rc);
-        ssh_channel_request_send_exit_status(sdata.channel, rc);
-    /* If client terminated the channel or the process did not exit nicely,
-     * but only if something has been forked. */
-    } else if (cdata.pid > 0) {
-        kill(cdata.pid, SIGKILL);
+    if (ssh_channel_is_open(sdata.channel)) {
+        ssh_channel_close(sdata.channel);
     }
 
-    ssh_channel_send_eof(sdata.channel);
-    ssh_channel_close(sdata.channel);
-
     /* Wait up to 5 seconds for the client to terminate the session. */
-    for (n = 0; n < 50 && (ssh_get_status(session) & SESSION_END) == 0; n++) {
+    for (int n = 0; n < 50 && (ssh_get_status(session) & SESSION_END) == 0; n++) {
         ssh_event_dopoll(event, 100);
     }
+    free(sdata.login_username);
+    ssh_event_remove_session(event, session);
 }
diff --git a/src/session.h b/src/session.h
index 2a4ea8d..50b5fe8 100644
--- a/src/session.h
+++ b/src/session.h
@@ -1,9 +1,19 @@
-#ifndef _SESSION_H_
-#define _SESSION_H_
+#ifndef SESSION_H_
+#define SESSION_H_
 
 #include <libssh/libssh.h>
 
-#include "util.h"
+#define SESSION_END (SSH_CLOSED | SSH_CLOSED_ERROR)
+
+/* A userdata struct for session. */
+struct session_data_struct {
+    /* Pointer to the channel the session will allocate. */
+    ssh_channel channel;
+    int auth_attempts;
+    int authenticated;
+    // ssh user name when login
+    char * login_username;
+};
 
 void handle_session(ssh_event event, ssh_session session);
 
diff --git a/src/util.h b/src/util.h
deleted file mode 100644
index 48a83ff..0000000
--- a/src/util.h
+++ /dev/null
@@ -1,32 +0,0 @@
-#ifndef _UTIL_H_
-#define _UTIL_H_
-
-#include <libssh/libssh.h>
-
-/* A userdata struct for channel. */
-struct channel_data_struct {
-    /* pid of the child process the channel will spawn. */
-    pid_t pid;
-    /* For PTY allocation */
-    socket_t pty_master;
-    socket_t pty_slave;
-    /* For communication with the child process. */
-    socket_t child_stdin;
-    socket_t child_stdout;
-    /* Only used for subsystem and exec requests. */
-    socket_t child_stderr;
-    /* Event which is used to poll the above descriptors. */
-    ssh_event event;
-    /* Terminal size struct. */
-    struct winsize *winsize;
-};
-
-/* A userdata struct for session. */
-struct session_data_struct {
-    /* Pointer to the channel the session will allocate. */
-    ssh_channel channel;
-    int auth_attempts;
-    int authenticated;
-};
-
-#endif