From fa12355609aa0e84b30d24945c360e16b7e5df72 Mon Sep 17 00:00:00 2001 From: Andreas Baumann Date: Sat, 26 Sep 2015 19:24:22 +0200 Subject: some code cleanup and added tagging of output in SSH mode (not completly correct yet) --- src/cssh.c | 221 ++++++++++++++++++++++++++++++++--------------------- src/options.ggo.in | 6 ++ 2 files changed, 140 insertions(+), 87 deletions(-) diff --git a/src/cssh.c b/src/cssh.c index b9b1917..2281eea 100644 --- a/src/cssh.c +++ b/src/cssh.c @@ -16,6 +16,7 @@ #include #include #include +#include #include "port.h" #include "progressbar.h" @@ -45,6 +46,16 @@ typedef enum execution_mode_e { CSSH_EXECUTE_AS_SCP } execution_mode_e; +typedef struct ssh_data_t { + ssh_channel channel; + char *stdout_buf; + char *stdout_bptr; + char *stderr_buf; + char *stderr_bptr; + bool eof_sent; + bool is_eof; +} ssh_data_t; + typedef enum copy_direction_e { CSSH_COPY_DIRECTION_UPLOAD, CSSH_COPY_DIRECTION_DOWNLOAD @@ -422,18 +433,29 @@ static int read_hosts_file( const char *hosts_file, unsigned short default_port, } // TODO: cleanup also progressbarpool -static void cleanup_sessions( ssh_session **session, ssh_channel **channel, scp_data_t **scp_data, char **host, unsigned short *port, const int nof_sessions, bool verbose ) +static void cleanup_sessions( ssh_session **session, ssh_data_t **ssh_data, scp_data_t **scp_data, char **host, unsigned short *port, const int nof_sessions, bool verbose ) { for( unsigned int i = 0; i < nof_sessions; i++ ) { if( ssh_is_connected( (*session)[i] ) ) { if( verbose ) { fprintf( stderr, "Disconnecting from '%s', port %d..\n", host[i], port[i] ); } - if( channel != NULL ) { - if( ssh_channel_is_open( (*channel)[i] ) ) { - ssh_channel_send_eof( (*channel)[i] ); - ssh_channel_close( (*channel)[i] ); - ssh_channel_free( (*channel)[i] ); + if( ssh_data != NULL ) { + ssh_channel *channel = &(*ssh_data)[i].channel; + if( ssh_channel_is_open( *channel ) ) { + ssh_channel_send_eof( *channel ); + ssh_channel_close( *channel ); + ssh_channel_free( *channel ); + } + char *stdout_buf = (*ssh_data)[i].stdout_buf; + if( stdout_buf != NULL ) { + free( stdout_buf ); + stdout_buf = NULL; + } + char *stderr_buf = (*ssh_data)[i].stderr_buf; + if( stderr_buf != NULL ) { + free( stderr_buf ); + stderr_buf = NULL; } } if( scp_data != NULL ) { @@ -467,8 +489,8 @@ static void cleanup_sessions( ssh_session **session, ssh_channel **channel, scp_ } ssh_free( (*session)[i] ); } - if( channel != NULL ) { - free( *channel ); + if( ssh_data != NULL ) { + free( *ssh_data ); } if( scp_data != NULL ) { free( *scp_data ); @@ -478,6 +500,47 @@ static void cleanup_sessions( ssh_session **session, ssh_channel **channel, scp_ free( port ); } +static const char *buffer_contains_a_line( const char *buffer, size_t bufsize ) +{ + for( const char *p = buffer; ( p < buffer + bufsize ) && *p != '\0'; p++ ) { + if( *p == '\n' ) { + return p; + } + } + return NULL; +} + +static ssize_t output_buffer( FILE *f, const char *host, const char *buffer, const size_t bufsize, const bool tagging ) +{ + const char *b = buffer; + size_t bsize = bufsize; + ssize_t nwritten = 0; + const char *p; + while( ( p < buffer + bufsize ) && ( p = buffer_contains_a_line( b, bsize ) ) != NULL) { + if( tagging ) { + fprintf( f, "[%s]: ", host ); + } + + ssize_t n = p - b; + ssize_t wrc = fwrite( b, 1, n, f ); + if( wrc < 0 ) { + return wrc; + } + if( wrc != n ) { + fprintf( stderr, "ERROR: Buffer write mismatch on %s (%zu != %d)\n", + ( f == stdout ) ? "stdout" : "stderr", + wrc, n ); + return -1; + } + nwritten += n + 1; + b = p + 1; + + fprintf( f, "\n" ); + } + + return nwritten; +} + int main( int argc, char *argv[] ) { struct gengetopt_args_info args_info; @@ -744,6 +807,8 @@ int main( int argc, char *argv[] ) cssh_msleep( 10 ); } + + free( is_connected ); // authentication phase @@ -868,170 +933,154 @@ int main( int argc, char *argv[] ) case CSSH_EXECUTE_AS_SSH: { // explicit low-level handling of channels in SSH mode - ssh_channel *channel = (ssh_channel *)malloc( ( nof_sessions + 1 ) * sizeof( ssh_channel ) ); - memset( channel, 0, nof_sessions + 1 ); - if( channel == NULL ) { - fprintf( stderr, "ERROR: Memory allocation failed for ssh_channels" ); + + ssh_data_t *ssh_data = (ssh_data_t *)malloc( ( nof_sessions + 1 ) * sizeof( ssh_data_t ) ); + if( ssh_data == NULL ) { + fprintf( stderr, "ERROR: Memory allocation failed for ssh_data array" ); cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } + memset( ssh_data, 0, ( nof_sessions + 1 ) * sizeof( ssh_data_t ) ); + for( unsigned int i = 0; i < nof_sessions; i++ ) { - channel[i] = ssh_channel_new( session[i] ); - if( channel[i] == NULL ) { + ssh_data[i].channel = ssh_channel_new( session[i] ); + if( ssh_data[i].channel == NULL ) { fprintf( stderr, "ERROR: Unable to open SSH channel: %s\n", ssh_get_error( session[i] ) ); - cleanup_sessions( &session, &channel, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } } for( unsigned int i = 0; i < nof_sessions; i++ ) { - rc = ssh_channel_open_session( channel[i] ); + rc = ssh_channel_open_session( ssh_data[i].channel ); if( rc != SSH_OK ) { - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } } // execute command on all channels for( unsigned int i = 0; i < nof_sessions; i++ ) { - rc = ssh_channel_request_exec( channel[i], cmd ); + rc = ssh_channel_request_exec( ssh_data[i].channel, cmd ); if( rc != SSH_OK ) { fprintf( stderr, "ERROR: Executing SSH command '%s' failed: %s\n", cmd, ssh_get_error( session[i] ) ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } } bool all_eof = false; - bool *eof_sent = (bool *)malloc( ( nof_sessions + 1 ) * sizeof( bool ) ); - if( eof_sent == NULL ) { - fprintf( stderr, "ERROR: Memory allocation failed for 'eof_sent'" ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); - exit( EXIT_FAILURE ); - } - memset( eof_sent, false, ( nof_sessions + 1 ) * sizeof( bool ) ); - bool *is_eof = (bool *)malloc( ( nof_sessions + 1 ) * sizeof( bool ) ); - if( is_eof == NULL ) { - fprintf( stderr, "ERROR: Memory allocation failed for 'is_eof'" ); - free( eof_sent ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); - exit( EXIT_FAILURE ); + + for( unsigned int i = 0; i < nof_sessions; i++ ) { + ssh_data[i].eof_sent = false; + ssh_data[i].is_eof = false; + ssh_data[i].stdout_buf = (char *)malloc( BUFSIZE ); + if( ssh_data[i].stdout_buf == NULL ) { + fprintf( stderr, "ERROR: Memory allocation failed for receiving buffers of ssh_channels" ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + exit( EXIT_FAILURE ); + } + ssh_data[i].stderr_buf = (char *)malloc( BUFSIZE ); + if( ssh_data[i].stderr_buf == NULL ) { + fprintf( stderr, "ERROR: Memory allocation failed for receiving buffers of ssh_channels" ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + exit( EXIT_FAILURE ); + } + ssh_data[i].stdout_bptr = ssh_data[i].stdout_buf; + ssh_data[i].stderr_bptr = ssh_data[i].stderr_buf; } - memset( is_eof, false, ( nof_sessions + 1 ) * sizeof( bool ) ); while( !all_eof ) { all_eof = true; for( unsigned int i = 0; i < nof_sessions; i++ ) { - if( !ssh_channel_is_closed( channel[i] ) && !ssh_channel_is_eof( channel[i] ) ) { + if( !ssh_channel_is_closed( ssh_data[i].channel ) && !ssh_channel_is_eof( ssh_data[i].channel ) ) { all_eof = false; } } // no stdin sent to commands on remote machines for now for( unsigned int i = 0; i < nof_sessions; i++ ) { - if( is_eof[i] && !eof_sent[i] ) { - ssh_channel_send_eof( channel[i] ); - eof_sent[i] = true; + if( ssh_data[i].is_eof && !ssh_data[i].eof_sent ) { + ssh_channel_send_eof( ssh_data[i].channel ); + ssh_data[i].eof_sent = true; continue; } } for( unsigned int i = 0; i < nof_sessions; i++ ) { - char buffer[BUFSIZE]; - rc = ssh_channel_poll( channel[i], 0 ); + + // handle stdout + rc = ssh_channel_poll( ssh_data[i].channel, 0 ); if( rc == SSH_ERROR ) { fprintf( stderr, "ERROR: ssh_channel_poll on stdout failed: %s\n", ssh_get_error( session[i] ) ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } if( rc > 0 ) { - unsigned int nread = ssh_channel_read_nonblocking( channel[i], buffer, sizeof( buffer ), 0 ); + unsigned int nread = ssh_channel_read_nonblocking( ssh_data[i].channel, ssh_data[i].stdout_bptr, BUFSIZE - ( ssh_data[i].stdout_bptr - ssh_data[i].stdout_buf ), 0 ); if( nread == SSH_ERROR ) { fprintf( stderr, "ERROR: ssh_channel_read_nonblocking on stdout failed: %s\n", ssh_get_error( session[i] ) ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } if( nread == 0 ) { - is_eof[i] = true; + ssh_data[i].is_eof = true; } if( nread > 0 ) { - size_t wrc = fwrite( buffer, 1, nread, stdout ); + ssize_t wrc = output_buffer( stdout, host[i], ssh_data[i].stdout_buf, nread + ( ssh_data[i].stdout_bptr - ssh_data[i].stdout_buf ), args_info.tagging_given ); if( wrc < 0 ) { fprintf( stderr, "ERROR: while writing to stdout: %s\n", strerror( errno ) ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); - exit( EXIT_FAILURE ); - } - - if( wrc != nread ) { - fprintf( stderr, "ERROR: Write mismatch on stdout (%zu != %d)\n", - wrc, nread ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); + } else if( wrc < nread ) { + memmove( ssh_data[i].stdout_buf, ssh_data[i].stdout_bptr + wrc, nread - wrc ); + ssh_data[i].stdout_bptr = ssh_data[i].stdout_buf + nread - wrc + 1; } } } - - rc = ssh_channel_poll( channel[i], 1 ); + + // handle stderr + rc = ssh_channel_poll( ssh_data[i].channel, 1 ); if( rc == SSH_ERROR ) { fprintf( stderr, "ERROR: ssh_channel_poll on stderr failed: %s\n", ssh_get_error( session[i] ) ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } if( rc > 0 ) { - unsigned int nread = ssh_channel_read_nonblocking( channel[i], buffer, sizeof( buffer ), 1 ); + unsigned int nread = ssh_channel_read_nonblocking( ssh_data[i].channel, ssh_data[i].stderr_bptr, BUFSIZE - ( ssh_data[i].stdout_bptr - ssh_data[i].stderr_buf ), 1 ); if( nread == SSH_ERROR ) { fprintf( stderr, "ERROR: ssh_channel_read_nonblocking on stderr failed: %s\n", ssh_get_error( session[i] ) ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } if( nread > 0 ) { - size_t wrc = fwrite( buffer, 1, nread, stderr ); + ssize_t wrc = output_buffer( stderr, host[i], ssh_data[i].stderr_buf, nread + ( ssh_data[i].stdout_bptr - ssh_data[i].stderr_buf ) , args_info.tagging_given ); if( wrc < 0 ) { fprintf( stderr, "ERROR: while writting to stderr: %s\n", strerror( errno ) ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); - exit( EXIT_FAILURE ); - } - - if( wrc != nread ) { - fprintf( stderr, "ERROR: Write mismatch on stderr (%zu != %d)\n", - wrc, nread ); - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); + } else if( wrc < nread ) { + memmove( ssh_data[i].stderr_buf, ssh_data[i].stderr_bptr + wrc, nread - wrc ); + ssh_data[i].stderr_bptr += nread - wrc + 1; } } if( nread == 0 ) { - is_eof[i] = true; + ssh_data[i].is_eof = true; } } } @@ -1039,9 +1088,7 @@ int main( int argc, char *argv[] ) cssh_msleep( 1 ); } - free( eof_sent ); - free( is_eof ); - cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); + cleanup_sessions( &session, &ssh_data, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); } break; @@ -1049,7 +1096,7 @@ int main( int argc, char *argv[] ) scp_data_t *scp_data = (scp_data_t *)malloc( ( nof_sessions + 1 ) * sizeof( scp_data_t ) ); if( scp_data == NULL ) { - fprintf( stderr, "ERROR: Memory allocation failed for ssh_data array" ); + fprintf( stderr, "ERROR: Memory allocation failed for scp_data array" ); cleanup_sessions( &session, NULL, NULL, host, port, nof_sessions, args_info.verbose_given > 0 ); exit( EXIT_FAILURE ); } diff --git a/src/options.ggo.in b/src/options.ggo.in index 8a90a72..0f003ea 100644 --- a/src/options.ggo.in +++ b/src/options.ggo.in @@ -29,6 +29,12 @@ you can specify a port immediatelly after a colon \ after the hostname. Example 'user@host:port'" string typestr="hosts-file" optional + +section "SSH options" + + option "tagging" - + "Enable tagging of each line with the name of the host" + optional section "SCP options" -- cgit v1.2.3-54-g00ecf