lib/pwkeep/main.rb in pwkeep-0.0.3 vs lib/pwkeep/main.rb in pwkeep-0.0.4

@@ -6,10 +6,28 @@
      attr :opts
   
      def initialize
        @opts = { :home => ENV['PWKEEP_HOME'] || '~/.pwkeep' } # required value
      end
+
+     def keypair_load
+       counter = 0
+       while true 
+         begin
+           pw = ask("Enter your password:") { |q| q.echo = false }
+           @storage.keypair_load pw
+         rescue OpenSSL::PKey::RSAError => e
+           say "<%= color('Invalid password', RED) %>"
+           counter = counter + 1
+           if (counter>2) 
+             raise e
+           end
+           next
+         end
+         break
+       end
+     end
   
      def setup
        @opts = Trollop::options do
          version "0.0.1 (c) 2014 Aki Tuomi"
          banner <<-EOS
@@ -29,10 +47,11 @@
          opt :delete, "Delete entry", :short => '-d'
          opt :search, "Search for system or username", :type => :string, :short => '-s'
          opt :list, "List all known systems", :short => '-l'
          opt :help, "Show usage", :short => '-h'
          opt :home, "Home directory", :short => '-H', :type => :string, :default => ( ENV['PWKEEP_HOME'] || '~/.pwkeep' )
+         opt :migrate, "Migrate from old storage format(s)"
          opt :version, "Show version", :short => '-V'
        end
   
        # validate options
        Trollop::die :system, "must be given for create/show/edit/delete" if opts[:system].nil? and (opts[:edit] or opts[:view] or opts[:delete] or opts[:create])
@@ -42,11 +61,11 @@
        Trollop::die :delete, "can only have one mode of operation" if opts[:delete] and (opts[:edit] or opts[:view] or opts[:create] or opts[:search] or opts[:initialize] or opts[:list])
        Trollop::die :search, "can only have one mode of operation" if opts[:search] and (opts[:edit] or opts[:view] or opts[:delete] or opts[:create] or opts[:initialize] or opts[:list])
        Trollop::die :initialize, "can only have one mode of operation" if opts[:initialize] and (opts[:edit] or opts[:view] or opts[:delete] or opts[:create] or opts[:search] or opts[:list])
        Trollop::die :list, "can only have one mode of operation" if opts[:list] and (opts[:edit] or opts[:view] or opts[:delete] or opts[:create] or opts[:search] or opts[:initialize])
   
-       Trollop::die "You must choose one mode of operation" unless opts[:create] or opts[:edit] or opts[:view] or opts[:delete] or opts[:search] or opts[:initialize] or opts[:list]
+       Trollop::die "You must choose one mode of operation" unless opts[:create] or opts[:edit] or opts[:view] or opts[:delete] or opts[:search] or opts[:initialize] or opts[:list] or opts[:migrate]
      end
   
      def self.run
        Main.new.run
      end
@@ -90,12 +109,11 @@
          end
    
          raise "Storage not initialized (run with --initialize)" unless @storage.valid? 
   
          if opts[:view]
-           pw = ask("Enter your password:") { |q| q.echo = false }
-           @storage.keypair_load pw
+           keypair_load
    
            data = @storage.load_system opts[:system]
   
            say("Last edited: #{data[:stored_at]}\nSystem: #{data[:system]}\n\n")
            say(data[:data])
@@ -109,20 +127,18 @@
   
            unless result[0]
              raise PWKeep::Exception, "Not modified"
            end
   
-           pw = ask("Enter your password:") { |q| q.echo = false }
-           @storage.keypair_load pw
+           keypair_load
            @storage.save_system opts[:system], result[1]
            say("<%= color('Changes stored', GREEN)%>")
            return
          end
   
          if opts[:edit]
-           pw = ask("Enter your password:") { |q| q.echo = false }
-           @storage.keypair_load pw
+           keypair_load
            data = @storage.load_system opts[:system]
            result = PWKeep.run_editor(data[:data], {})
            unless result[0]
              raise PWKeep::Exception, "Not modified"
            end
@@ -130,45 +146,52 @@
            say("<%= color('Changes stored', GREEN)%>")
            return
          end
   
          if opts[:delete]
-           pw = ask("Enter your password:") { |q| q.echo = false }
-           @storage.keypair_load pw
+           keypair_load
            data = @storage.load_system opts[:system]
            # just to be sure
            unless agree("Are you <%=color('SURE',BOLD)%> you want to delete #{data[:system]}?")
              @storage.delete data[:system]
            end
            say("<%= color('System deleted', YELLOW)%>")
            return
          end
   
          if opts[:search]
-           pw = ask("Enter your password:") { |q| q.echo = false }
-           @storage.keypair_load pw
+           keypair_load
            say("All matching systems\n")
            @storage.list_all_systems.sort.each do |system|
              if system.match opts[:search]
                say("  - #{system}")
              end
            end
            return
          end
   
          if opts[:list]
-           pw = ask("Enter your password:") { |q| q.echo = false }
-           @storage.keypair_load pw
+           keypair_load
            say("All known systems\n")
            @storage.list_all_systems.sort.each do |system| 
              say("  - #{system}")
            end
            return
          end
+
+         if opts[:migrate]
+           keypair_load
+           say("Migrating systems...\n")
+           count = @storage.migrate
+           say("Migrated #{count} systems\n")
+           return
+         end
        rescue PWKeep::Exception => e1
         PWKeep::logger.error e1.message.colorize(:red)
        rescue OpenSSL::PKey::RSAError => e2
         PWKeep::logger.error "Cannot load private key".colorize(:red)
+       rescue SystemExit,Interrupt
+        # ignore
        end
      end
   end
 end