/*

 * This file is part of the SSH Library

 *

 * Copyright (c) 2010 by Aris Adamantiadis

 *

 * The SSH Library is free software; you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as published by

 * the Free Software Foundation; either version 2.1 of the License, or (at your

 * option) any later version.

 *

 * The SSH Library is distributed in the hope that it will be useful, but

 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY

 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public

 * License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with the SSH Library; see the file COPYING.  If not, write to

 * the Free Software Foundation, Inc., 59 Temple Place - Suite 330, Boston,

 * MA 02111-1307, USA.

 */

#include "benchmarks.h"

#include <string.h>

#include <stdlib.h>

#include <stdio.h>

#define PYTHON_PATH "/usr/bin/python"

const char python_eater[]=

"#!/usr/bin/python\n"

"import sys\n"

"print 'go'\n"

"sys.stdout.flush()\n"

"toread=XXXXXXXXXX\n"

"read=0\n"

"while(read < toread):\n"

"    buffersize=toread-read\n"

"    if(buffersize > 4096):\n"

"        buffersize=4096\n"

"    r=len(sys.stdin.read(buffersize))\n"

"    read+=r\n"

"    if(r<=0):\n"

"        print 'error'\n"

"        exit()\n"

"print 'done'\n";

static char *get_python_eater(unsigned long bytes){

  char *eater=malloc(sizeof(python_eater));

  char *ptr;

  char buf[12];

  memcpy(eater,python_eater,sizeof(python_eater));

  ptr=strstr(eater,"XXXXXXXXXX");

  if(!ptr){

    free(eater);

    return NULL;

  }

  sprintf(buf,"0x%.8lx",bytes);

  memcpy(ptr,buf,10);

  return eater;

}

/** @internal

 * @brief uploads a script (python or other) at a specific path on the

 * remote host

 * @param[in] session an active SSH session

 * @param[in] path to copy the file

 * @param[in] content of the file to copy

 * @return 0 on success, -1 on error

 */

static int upload_script(ssh_session session, const char *path,

    const char *script){

  ssh_channel channel;

  char cmd[128];

  int err;

  channel=ssh_channel_new(session);

  if(!channel)

    goto error;

  if(ssh_channel_open_session(channel) == SSH_ERROR)

    goto error;

  snprintf(cmd,sizeof(cmd),"cat > %s",path);

  if(ssh_channel_request_exec(channel,cmd) == SSH_ERROR)

    goto error;

  err=ssh_channel_write(channel,script,strlen(script));

  if(err == SSH_ERROR)

    goto error;

  if(ssh_channel_send_eof(channel) == SSH_ERROR)

    goto error;

  if(ssh_channel_close(channel) == SSH_ERROR)

    goto error;

  ssh_channel_free(channel);

  return 0;

error:

  fprintf(stderr,"Error while copying script : %s\n",ssh_get_error(session));

  return -1;

}

/** @internal

 * @brief benchmarks a raw upload (simple upload in a SSH channel) using an

 * existing SSH session.

 * @param[in] session Open SSH session

 * @param[in] args Parsed command line arguments

 * @param[out] bps The calculated bytes per second obtained via benchmark.

 * @return 0 on success, -1 on error.

 */

int benchmarks_raw_up (ssh_session session, struct argument_s *args,

    float *bps){

  unsigned long bytes;

  char *script;

  char cmd[128];

  int err;

  ssh_channel channel;

  struct timestamp_struct ts;

  float ms=0.0;

  unsigned long total=0;

  bytes = args->datasize * 1024 * 1024;

  script =get_python_eater(bytes);

  err=upload_script(session,"/tmp/eater.py",script);

  free(script);

  if(err<0)

    return err;

  channel=ssh_channel_new(session);

  if(channel == NULL)

    goto error;

  if(ssh_channel_open_session(channel)==SSH_ERROR)

    goto error;

  snprintf(cmd,sizeof(cmd),"%s /tmp/eater.py", PYTHON_PATH);

  if(ssh_channel_request_exec(channel,cmd)==SSH_ERROR)

    goto error;

  if((err=ssh_channel_read(channel,buffer,sizeof(buffer)-1,0))==SSH_ERROR)

    goto error;

  buffer[err]=0;

  if(!strstr(buffer,"go")){

    fprintf(stderr,"parse error : %s\n",buffer);

    ssh_channel_close(channel);

    ssh_channel_free(channel);

    return -1;

  }

  if(args->verbose>0)

    fprintf(stdout,"Starting upload of %lu bytes now\n",bytes);

  timestamp_init(&ts);

  while(total < bytes){

    unsigned long towrite = bytes - total;

    int w;

    if(towrite > args->chunksize)

      towrite = args->chunksize;

    w=ssh_channel_write(channel,buffer,towrite);

    if(w == SSH_ERROR)

      goto error;

    total += w;

  }

  if(args->verbose>0)

    fprintf(stdout,"Finished upload, now waiting the ack\n");

  if((err=ssh_channel_read(channel,buffer,5,0))==SSH_ERROR)

      goto error;

  buffer[err]=0;

  if(!strstr(buffer,"done")){

    fprintf(stderr,"parse error : %s\n",buffer);

    ssh_channel_close(channel);

    ssh_channel_free(channel);

    return -1;

  }

  ms=elapsed_time(&ts);

  *bps=8000 * (float)bytes / ms;

  if(args->verbose > 0)

    fprintf(stdout,"Upload took %f ms for %lu bytes, at %f bps\n",ms,

        bytes,*bps);

  ssh_channel_close(channel);

  ssh_channel_free(channel);

  return 0;

error:

  fprintf(stderr,"Error during raw upload : %s\n",ssh_get_error(session));

  if(channel){

    ssh_channel_close(channel);

    ssh_channel_free(channel);

  }

  return -1;

}

const char python_giver[] =

"#!/usr/bin/python\n"

"import sys\n"

"r=sys.stdin.read(2)\n"

"towrite=XXXXXXXXXX\n"

"wrote=0\n"

"mtu = 32786\n"

"buf = 'A'*mtu\n"

"while(wrote < towrite):\n"

"    buffersize=towrite-wrote\n"

"    if(buffersize > mtu):\n"

"        buffersize=mtu\n"

"    if(buffersize == mtu):\n"

"        sys.stdout.write(buf)\n"

"    else:\n"

"        sys.stdout.write('A'*buffersize)\n"

"    wrote+=buffersize\n"

"sys.stdout.flush()\n";

static char *get_python_giver(unsigned long bytes){

  char *giver=malloc(sizeof(python_giver));

  char *ptr;

  char buf[12];

  memcpy(giver,python_giver,sizeof(python_giver));

  ptr=strstr(giver,"XXXXXXXXXX");

  if(!ptr){

    free(giver);

    return NULL;

  }

  sprintf(buf,"0x%.8lx",bytes);

  memcpy(ptr,buf,10);

  return giver;

}

/** @internal

 * @brief benchmarks a raw download (simple upload in a SSH channel) using an

 * existing SSH session.

 * @param[in] session Open SSH session

 * @param[in] args Parsed command line arguments

 * @param[out] bps The calculated bytes per second obtained via benchmark.

 * @return 0 on success, -1 on error.

 */

int benchmarks_raw_down (ssh_session session, struct argument_s *args,

    float *bps){

  unsigned long bytes;

  char *script;

  char cmd[128];

  int err;

  ssh_channel channel;

  struct timestamp_struct ts;

  float ms=0.0;

  unsigned long total=0;

  bytes = args->datasize * 1024 * 1024;

  script =get_python_giver(bytes);

  err=upload_script(session,"/tmp/giver.py",script);

  free(script);

  if(err<0)

    return err;

  channel=ssh_channel_new(session);

  if(channel == NULL)

    goto error;

  if(ssh_channel_open_session(channel)==SSH_ERROR)

    goto error;

  snprintf(cmd,sizeof(cmd),"%s /tmp/giver.py", PYTHON_PATH);

  if(ssh_channel_request_exec(channel,cmd)==SSH_ERROR)

    goto error;

  if((err=ssh_channel_write(channel,"go",2))==SSH_ERROR)

    goto error;

  if(args->verbose>0)

    fprintf(stdout,"Starting download of %lu bytes now\n",bytes);

  timestamp_init(&ts);

  while(total < bytes){

    unsigned long toread = bytes - total;

    int r;

    if(toread > args->chunksize)

      toread = args->chunksize;

    r=ssh_channel_read(channel,buffer,toread,0);

    if(r == SSH_ERROR)

      goto error;

    total += r;

  }

  if(args->verbose>0)

    fprintf(stdout,"Finished download\n");

  ms=elapsed_time(&ts);

  *bps=8000 * (float)bytes / ms;

  if(args->verbose > 0)

    fprintf(stdout,"Download took %f ms for %lu bytes, at %f bps\n",ms,

        bytes,*bps);

  ssh_channel_close(channel);

  ssh_channel_free(channel);

  return 0;

error:

  fprintf(stderr,"Error during raw upload : %s\n",ssh_get_error(session));

  if(channel){

    ssh_channel_close(channel);

    ssh_channel_free(channel);

  }

  return -1;

}