Commit 8a4434a

@@ -40,19 +40,14 @@ module TFA
         return ""
       end
 
-      if yes? "Upgrade to #{yml_path}?"
+      if yes? "Upgrade to #{yaml_path}?"
         pstore_storage.each do |row|
           row.each do |name, secret|
             yaml_storage.save(name, secret) if yes?("Migrate `#{name}`?")
           end
         end
-        if yes? "Encrypt `#{yml_path}`?"
-          passphrase = ask "Enter passphrase:", echo: false
-          say passphrase
-        end
-        if yes? "Delete `#{pstore_path}`?"
-          File.delete(pstore_path)
-        end
+        yaml_storage.encrypt!(passphrase) if yes?("Encrypt?")
+        File.delete(pstore_path) if yes?("Delete `#{pstore_path}`?")
       end
       ""
     end
@@ -68,7 +63,7 @@ module TFA
     end
 
     def yaml_storage
-      @yaml_storage ||= Storage.new(yml_path)
+      @yaml_storage ||= Storage.new(yaml_path)
     end
 
     def filename
@@ -83,7 +78,7 @@ module TFA
       File.join(directory, ".#{filename}.pstore")
     end
 
-    def yml_path
+    def yaml_path
       File.join(directory, ".#{filename}.yml")
     end
 
@@ -94,5 +89,9 @@ module TFA
         secret
       end
     end
+
+    def passphrase
+      @passphrase ||= ask("Enter passphrase:", echo: false)
+    end
   end
 end

@@ -1,8 +1,10 @@
 module TFA
   class Storage
     include Enumerable
+    attr_reader :path
 
     def initialize(path)
+      @path = path
       @storage =
         if ".pstore" == File.extname(path)
           PStore.new(path)
@@ -41,6 +43,26 @@ module TFA
       end
     end
 
+    def encrypt!(passphrase)
+      cipher.encrypt
+      cipher.key = Digest::SHA256.digest(passphrase)
+      cipher.iv = cipher.random_iv
+
+      plain_text = read_all
+      cipher_text = cipher.update(plain_text) + cipher.final
+      flush(cipher_text)
+    end
+
+    def decrypt!(passphrase)
+      cipher_text = read_all
+      decipher = cipher
+      decipher.decrypt
+      decipher.iv = cipher_text[0..decipher.iv_len-1]
+      cipher.key = Digest::SHA256.digest(passphrase)
+      data = cipher_text[decipher.iv_len..-1]
+      flush(decipher.update(data) + decipher.final)
+    end
+
     private
 
     def open_readonly
@@ -48,5 +70,17 @@ module TFA
         yield @storage
       end
     end
+
+    def cipher
+      @cipher ||= OpenSSL::Cipher.new("AES-256-CBC")
+    end
+
+    def read_all
+      IO.read(path)
+    end
+
+    def flush(data)
+      IO.write(path, data)
+    end
   end
 end

@@ -1,8 +1,10 @@
+require "digest"
+require "openssl"
 require "pstore"
-require "yaml/store"
 require "rotp"
-require "openssl"
-require "tfa/version"
-require "tfa/totp_command"
-require "tfa/storage"
+require "yaml/store"
+
 require "tfa/cli"
+require "tfa/storage"
+require "tfa/totp_command"
+require "tfa/version"