(long) need advice to improve code

Full code follows at the bottom.


I've recently started hacking around on this again, and I'm stuck on a few 
things.  Firstly, this little work is a real memory hog, but I'm not skilled 
enough to be able to figure out why (anyone?).  Secondly, I have some minor 
issues that I'd like to solve, but I'm not sure how.  First, a little 
background of what this does.  This program basically logs you into any 
number of systems, using telnet and the same username and password for all 
of them.  Then you can run the same command on all of them by only typing it 
once.  So, I can say ./stelnet.rb host1 .. host20, log in just one time and 
then look at the root filesystems on all of them by running "ls /" one time 
and scrolling through the output.  So, on to my questions.  I'm having a 
hard time figuring out how to run some interactive commands which don't 
return any type of prompt (i.e. ex).  If I try to start an ex session (to 
edit a file) my program times out waiting for a prompt to return.  How can I 
get around this?  Also, how could I pass a CTRL-c through to ex instead of 
having my program die over it?  I guess the basic question is, how do I get 
this to act more like a real telnet session (i.e. more solid)?

Thanks

Code follows

#!/usr/gnu/bin/ruby

#  %W%        %G%
#
#  Program Name:        stelnet.rb
#  Date Created:        02/06/01
#  Creator:             Mike Wilson
#
#  SYNOPSIS
#       stelnet.rb host1..hostn
#
#  DESCRIPTION
#       stelnet will log you in to multiple systems, using telnet, and run
#       the specified commands.
#
require 'net/telnet'

module Net
    class Telnet
        # Allow us to get the hostname by printing the connection
        def to_s
            @options['Host']
        end
    end
end

module SuperTelnet
    class User
        # I'll add some stuff to this class at some point.
        attr_accessor :username, :passwd

        def initialize(user, pass)
            @username = user
            @passwd = pass
        end
    end

    class Connection
        # Keep track of all the open connections
        @@uservars = {}

        # This will be used when we only want to execute something once
        # i.e. help should only be displayed 1 time -- not 1 time per
        # connection.
        @@flag = 0

        def initialize(aHash, userObj)
            raise ArgumentError if not userObj.instance_of? 
SuperTelnet::User

            @conOptions = aHash
            @thisUser = userObj

            # Each user gets an array which will hold his connections
            # Hopefully, we'll be able to handle multiple users at some 
point.
            @@uservars[@thisUser] ||= []

            # Set some defaults so that you really only need to provide
            # the hostname to connect to within the conOptions hash.
            @conOptions['Prompt'] ||= Regexp.new('[#$%>] *\z', false, 'n')
            @conOptions['Timeout'] ||= 10
            @conOptions['Port'] ||= 23

            begin
	        @thisConnection = Net::Telnet.new(@conOptions)
            rescue SocketError
                yield "SocketError on #{@conOptions['Host']}, bad host?" if
                    block_given?
            rescue TimeoutError
                yield "TimeoutError on #{@conOptions['Host']}, host down?" 
if
                    block_given?
            rescue Errno::ECONNREFUSED
                yield "Connection refused on #{@conOptions['Host']}" if
                    block_given?
            else
                @@uservars[@thisUser] << self
            end

            @conOptions['Timeout'] = 30
        end

        # If the user enters something that looks like a meta command, but
        # isn't, method_missing will catch it and send it to the executeCmd
        # method to be issued on the command line.
        def method_missing(methodId)
            executeCmd(methodId.id2name) do |c|
                yield c if block_given?
            end
        end

        def to_s
            @thisConnection.to_s
        end

        def Connection.uservars
            @@uservars
        end

        # Always call this after looping through all of the hosts within
        # your implementation of this module.
        def Connection.flagReset
            @@flag = 0
        end

        # Feel free to redefine this as you add commands
        def Connection.help
            "\nAvaiable commands:\n" +
            "  !login(host [hostn])             -> autologin to hosts\n" +
            "  !logoff(host [hostn])            -> logoff hosts\n" +
            "  !su(host [hostn])                -> change user on hosts\n" +
            "  !su                              -> to su on all hosts\n" +
            "  !on(host [hostn]) somecommand    -> run only on these 
hosts\n" +
            "  !off(host [hostn]) somecommand   -> don't run on these\n" +
            "  !set key=value                   -> set user variable\n" +
            "                                      (reference via @key)\n" +
            "  !unset key                       -> unset user variable\n" +
            "  !quit(host [hostn])              -> logoff hosts\n" +
            "  !help                            -> this screen\n\n"
        end

        # Call this for each connection that's been created for thisUser.
        def loginOnConnection
            return if @thisConnection.nil?
            begin
                @thisConnection.login(@thisUser.username, @thisUser.passwd)
            rescue TimeoutError
                yield "TimeoutError on #{@thisConnection}, bad login?" if
                    block_given?
                logoffConnection
            end
        end

        # Call this when you're all done with the connection.  I use this
        # within the loginOnConnection method to remove the connection if
        # the login fails (i.e. bad login).
        def logoffConnection
            @@uservars[@thisUser].delete(self)
            @thisConnection.close
        end

        ####
        # Commands:  So far, you have 2 types of commands -- those that look
        # like !this and those that look like ?this.  !this being a regular
        # command, and ?this being more like a query. (i.e. ?threads #=> on)
        #
        # To add a new command, simply add the corresponding method.  Either
        # Bang_commandName(args, extras) or Qmark_commandName(args, extras)
        # depending on whether it's query or not.  Use the flag if
        # you only want it to execute once.  Other than that, you're done.
        ####

        # !help will print the help message.  Another reason to call
        # Connection.flagReset when you're done
        def Bang_help(args, extras)
            if @@flag == 0
                @@flag = 1
                Connection.help
            end
        end

        # This command equates to !su and !su(host1 host2) and !su host1 
host2
        # Here, we're allowing a method to perform su's so that you can
        # implement a way to hide the password.
        def Bang_su(args, extras)
            user, pass, *hosts = args
            if hosts.size.zero? or hosts.detect do |c|
                    c =~ @thisConnection.to_s
                end
                tmpHash = @conOptions

                # I'm assuming that you want the environment (i.e. the '-')
                tmpHash['String'] = "su - #{user}"
                tmpHash['Match'] = /[pP]assword:/

                @thisConnection.cmd(tmpHash) do |c| yield c if block_given? 
end

                tmpHash['String'] = pass
                tmpHash['Match'] = @conOptions['Prompt']
                @thisConnection.cmd(tmpHash) do |c| yield c if block_given? 
end
            end
        end

        # This command equates to !login(host1 host2) and !login host1 host2
        # Here, you'll be logged into the new hosts without having to 
provide
        # login info -- it'll just use your initial info.
        #
        # This also shows a case where you'd need to call 
Connection.flagReset
        def Bang_login(args, extras)
            if @@flag == 0
                @@flag = 1
                args.each do |arg|
                    Connection.new({'Host' => arg}, @thisUser) do |c|
                        puts c
                    end.loginOnConnection do |c|
                        puts c
                    end
                end
            end
            return
        end

        # This command is !logoff and !logoff(host1 host2) and !logoff host1
        # self-explanatory.
        def Bang_logoff(args, extras)
            raise EOFError if args.size.zero?
            if args.detect do |c|
                    c == @thisConnection.to_s
                end
                logoffConnection
            end
        end

        # Pretty much like the Unix env command at a basic level.  !env
        def Bang_env(args, extras)
            output = []
            if @@flag == 0
                @@flag = 1
                @@uservars.each do |k, v|
                    v = v.join(' ') if v.type == Array
                    output << "#{k}=#{v}\n"
                end
                output
            end
        end

        # This command equates to "!on(host1 host2) do some command"
        def Bang_on(args, extras)
            output = []
            args.each do |arg|
                arg == @thisConnection.to_s and
                    output << runCmdOnConnection(extras)
            end
            output
        end

        # Opposite of above, don't run on these hosts. "!off(host1 host2) 
..."
        def Bang_off(args, extras)
            output = []
            args.each do |arg|
                arg != @thisConnection.to_s and
                    output << runCmdOnConnection(extras)
            end
            output
        end

        # Set user variable.  Use the variable like @this.
        # So, !set foo=bar  can be used as @foo
        def Bang_set(args, extras)
            if @@flag == 0
                @@flag = 1
                k, v = args.join(' ').to_s.split('=')
                @@uservars[k] = v
            end
        end

        # Unset user variable.  !unset foo
        def Bang_unset(args, extras)
            if @@flag == 0
                @@flag = 1
                args.each do |arg|
                    @@uservars.delete(arg)
                end
            end
        end

        # Everything typed by the user (all commands, whether they're to
        # be processed internally or sent to the shell on each box) should
        # pass through this method.  So, really this should probably be 
called
        # proxyCmd or something.
        def executeCmd(cmd)
            # Substitute user variables out before anything else
            @@uservars.each do |k, v|
                cmd.sub!('@' + k.to_s, v.to_s)
            end

            output = case cmd
                when /^([!?=].*)/
                    # It might be an internal command.  Try to process it as
                    # such and if it fails, pass it to the shell
                    begin
                        executeMetaCmd($1)
                    rescue ArgumentError
                        runCmdOnConnection($1)
                    end
                else
                    # Didn't look like anything we're supposed to deal with.
                    # Send it to the shell.
                    runCmdOnConnection(cmd)
            end
            if output and not output.empty? and block_given?
                # Use yields rather than prints so that we can throw a GUI
                # frontend on this someday.
                yield "#{output}\n\n"
            end
            output
        end

        # This is kind of like a proxy for our internal commands.
        def executeMetaCmd(cmd)
            # extras is necessary for commands like:
            # ?on(host1 host2) ls -ltr /var
            # extras would equal "ls -ltr /var" here.
            extras = cmd.sub(/[!\?]\w+\s*\(.*\)\s*(.*)$/, '\1')

            # This is the command (i.e. login) and the hosts (generally)
            metaCmd, *args = cmd.sub(/([!\?]\w+)\s*\((.*)\).*$/, '\1 
\2').split

            case metaCmd
                when /!(\w+)/
                    # If the command was: !login foo bar
                    # call Bang_login(['foo', 'bar'], "")
                    meth = 'Bang_' + $1
                when /\?(\w+)/
                    # If the command was: ?on(oscar) ls -ltr /var
                    # call Qmark_on(['oscar'], "ls -ltr /var")
                    meth = 'Qmark_' + $1
            end
            send(meth, args, extras)
        end

        # This will send the command to the shell for this connection.
        def runCmdOnConnection(cmd)
            output = ["\n#{@thisConnection}:\n"]

            @thisConnection.cmd({
                'String'  => cmd,
                'Match'   => @conOptions['Prompt'],
                'Timeout' => false
            }) do |c|
                if c and c.chomp != cmd
                    output << c
                end
            end
            output
        end

        alias_method :Bang_quit, :Bang_logoff
        alias_method :Bang_exit, :Bang_logoff
    end
end

if __FILE__ == $0
    # We need to add a command "!threads(on|off)" so that we can have
    # sequential or threaded access.  We add all of the support methods, 
along
    # with redefining Connection.help to account for the new command, and
    # finally we add the actual command method Bang_threads.
    module SuperTelnet
        class Connection
            @@threading = 'on'

            def Connection.threading=(toggle)
                toggle = @@threading if toggle != 'on' and toggle != 'off'
                @@threading = toggle
                return
            end

            def Connection.threading
                @@threading
            end

            def Connection.help
                "\nAvaiable commands:\n" +
                "  !login(host [hostn])             -> autologin to hosts\n" 
+
                "  !logoff(host [hostn])            -> logoff hosts\n" +
                "  !su(host [hostn])                -> change user on 
hosts\n" +
                "  !su                              -> to su on all hosts\n" 
+
                "  !on(host [hostn]) somecommand    -> run only on these 
hosts\n" +
                "  !off(host [hostn]) somecommand   -> don't run on these\n" 
+
                "  !set key=value                   -> set user variable\n" 
+
                "                                      (reference via 
@key)\n" +
                "  !unset key                       -> unset user 
variable\n" +
                "  !threads <on|off>                -> set threaded or 
sequential\n" +
                "  ?threads                         -> return thread 
status\n" +
                "  !quit(host [hostn])              -> logoff hosts\n" +
                "  !help                            -> this screen\n\n"
            end

            def Bang_threads(args, extras)
                Connection.threading = args.to_s
            end

            def Qmark_threads(args, extras)
                Connection.threading
            end
        end
    end

    include SuperTelnet

    progname = File.basename $0

    def getUser
        print 'Username: '
        $stdin.gets.chomp
    end

    def getPwd
        print 'Password: '
        begin
            system 'stty -echo'
            $stdin.gets.chomp
        ensure
            puts
            system 'stty echo'
        end
    end

    def usage(progname)
        print "Usage: " + progname
        puts " host1..hostn"
        print Connection.help
        exit
    end

    usage(progname) if ARGV.size.zero? or ARGV[0] =~ /^-[?h]/
    user = User.new(getUser, getPwd)

    ARGV.each do |aHost|
        Connection.new({'Host'=>aHost}, user) do |c| puts c end
    end

    loginThreads = []
    Connection.uservars[user].each do |aCon|
        loginThreads << Thread.new(aCon) do |con|
            con.loginOnConnection do |c| puts c end
        end
    end
    loginThreads.each do |aThread| aThread.join end

    cmdThreads = []
    # Set threading on
    Connection.threading = 'on'
    begin
        while true
            exit if Connection.uservars[user].size.zero?

            print "#{progname}> "
            cmd = $stdin.readline.chomp

            # Within our module, the Bang_su method is called as
            # Bang_su(['username', 'password', 'host1', 'host2'], "")
            # Obviously, we don't want you to type it in on the commandline
            # like that (you'll expose your password), so we proxy instead.
            # You enter !su(host1 host2) and we call getUser and getPwd to
            # prompt you for you info.
            if cmd =~ /!su\s*\(?(.*)+\)?$/
                cmd = "!su(#{getUser} #{getPwd} #{$1})"
            end

            if Connection.threading == 'on'
                Connection.uservars[user].each do |aCon|
                    cmdThreads << Thread.new(aCon) do |con|
                        begin
                            con.executeCmd(cmd) do |c| puts c end
                        rescue Errno::ECONNRESET, Errno::EPIPE
                            puts "#{con}: Connection reset by peer - 
autologoff"
                            con.logoffConnection
                        end
                    end
                end
                cmdThreads.each do |aThread| aThread.join end
            else
                Connection.uservars[user].each do |con|
                    begin
                        con.executeCmd(cmd) do |c| puts c end
                    rescue Errno::ECONNRESET, Errno::EPIPE
                        puts "#{con}: Connection reset by peer - autologoff"
                        con.logoffConnection
                    end
                end
            end
            # Always do this when you've completed processing each command.
            Connection.flagReset
        end
    rescue Interrupt, EOFError
        Connection.uservars[user].each do |con|
            con.logoffConnection
        end
    rescue SocketError
        retry
    end
end




_________________________________________________________________
Chat with friends online, try MSN Messenger: http://messenger.msn.com