From c1fe49608b3701b8b3b789366ef946dd7460df06 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Wed, 14 Apr 2021 08:02:20 +0800
Subject: [PATCH] update decrypt

---
 .../tools/tinyencrypt/TinyEncryptMain.java    | 13 ++++-
 .../encrypt/EncryptedFileUtil.java            | 56 ++++++++++++++++++-
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/src/main/java/me/hatter/tools/tinyencrypt/TinyEncryptMain.java b/src/main/java/me/hatter/tools/tinyencrypt/TinyEncryptMain.java
index 29025ce..cee8ce3 100644
--- a/src/main/java/me/hatter/tools/tinyencrypt/TinyEncryptMain.java
+++ b/src/main/java/me/hatter/tools/tinyencrypt/TinyEncryptMain.java
@@ -15,7 +15,7 @@ public class TinyEncryptMain {
 
     static {
         BCUtil.init();
-        LogConfig.initMuteInfoMode();
+//        LogConfig.initMuteInfoMode();
         log = LogTools.getLogTool(TinyEncryptMain.class);
     }
 
@@ -47,10 +47,17 @@ public class TinyEncryptMain {
             log.error("FILE is not assigned.");
             return;
         }
-
+        int total = tinyEncryptArgs.files.length;
         try {
+            int index = 1;
             for (File f : tinyEncryptArgs.files) {
-                EncryptedFileUtil.encryptFile(tinyEncryptArgs.key, f, tinyEncryptArgs.comment);
+                log.info("Start processing file: " + f + ", " + index + " of " + total);
+                if (tinyEncryptArgs.encrypt) {
+                    EncryptedFileUtil.encryptFile(tinyEncryptArgs.key, f, tinyEncryptArgs.comment);
+                } else {
+                    EncryptedFileUtil.decryptFile(f);
+                }
+                index++;
             }
         } catch (JumpOutException joe) {
             log.error(joe.getMessage());
diff --git a/src/main/java/me/hatter/tools/tinyencrypt/encrypt/EncryptedFileUtil.java b/src/main/java/me/hatter/tools/tinyencrypt/encrypt/EncryptedFileUtil.java
index e079e3b..cb4e7c2 100644
--- a/src/main/java/me/hatter/tools/tinyencrypt/encrypt/EncryptedFileUtil.java
+++ b/src/main/java/me/hatter/tools/tinyencrypt/encrypt/EncryptedFileUtil.java
@@ -1,10 +1,13 @@
 package me.hatter.tools.tinyencrypt.encrypt;
 
+import com.alibaba.fastjson.JSON;
+import me.hatter.tools.commons.assertion.AssertUtil;
 import me.hatter.tools.commons.io.CopyRollCounter;
 import me.hatter.tools.commons.io.IOUtil;
 import me.hatter.tools.commons.log.LogTool;
 import me.hatter.tools.commons.log.LogTools;
 import me.hatter.tools.commons.security.crypt.CryptOutputStream;
+import me.hatter.tools.commons.tlv.Tlv;
 import me.hatter.tools.commons.tlv.TlvUtil;
 import me.hatter.tools.tinyencrypt.util.BCUtil;
 
@@ -17,6 +20,38 @@ public class EncryptedFileUtil {
         BCUtil.init();
     }
 
+    public static void decryptFile(File file) {
+        File decFile = getDecryptFile(file);
+        if (decFile == null) {
+            log.warn("File is not tinyenc file, skip: " + decFile);
+            return;
+        }
+        if (decFile.exists()) {
+            log.warn("File exists, skip: " + decFile);
+            return;
+        }
+        try {
+            try (FileInputStream fis = new FileInputStream(file)) {
+                Tlv tlv = TlvUtil.readTlv(fis);
+                TinyEncryptMeta meta = tlv.getValueAsBytes().asJSONObject(TinyEncryptMeta.class);
+                try (FileOutputStream fos = new FileOutputStream(decFile)) {
+                    try (OutputStream newOs = getDecryptOutputStream(fos, meta)) {
+                        IOUtil.copy(fis, newOs, new CopyRollCounter() {
+                            @Override
+                            public void count(long count, long length) {
+                                // TODO ... ProcessBar
+                            }
+                        });
+                    }
+                }
+            }
+            log.info("Decrypt file success: " + file);
+        } catch (Exception e) {
+            log.error("Decrypt file filed: " + file + ", reason: " + e.getMessage());
+            log.debug("Decrypt file filed: " + file + ", reason: " + e.getMessage(), e);
+        }
+    }
+
     public static void encryptFile(String keyName, File file, String comment) {
         File encFile = getEncryptFile(file);
         if (encFile.exists()) {
@@ -29,11 +64,11 @@ public class EncryptedFileUtil {
                     TinyEncryptMeta meta = TinyEncryptMetaUtil.create(keyName, comment);
                     TlvUtil.writeTlv(fos, TlvUtil.create(1, TinyEncryptMetaUtil.toString(meta)));
                     fos.flush();
-                    try (OutputStream newOs = getCryptoOutputStream(fos, meta)) {
+                    try (OutputStream newOs = getEncryptOutputStream(fos, meta)) {
                         IOUtil.copy(fis, newOs, new CopyRollCounter() {
                             @Override
                             public void count(long count, long length) {
-                                // TODO ...
+                                // TODO ... ProcessBar
                             }
                         });
                     }
@@ -52,7 +87,22 @@ public class EncryptedFileUtil {
         return encFile;
     }
 
-    private static OutputStream getCryptoOutputStream(OutputStream out, TinyEncryptMeta tinyEncryptMeta) {
+    public static File getDecryptFile(File file) {
+        File absFile = file.getAbsoluteFile();
+        String fn = absFile.getName();
+        if (!fn.endsWith(".tinyenc")) {
+            return null;
+        }
+        return new File(absFile.getParent(), fn.substring(0, fn.length() - ".tinyenc".length()));
+    }
+
+    private static OutputStream getEncryptOutputStream(OutputStream out, TinyEncryptMeta tinyEncryptMeta) {
         return CryptOutputStream.gcmEncrypt(out, tinyEncryptMeta.getDataKey(), tinyEncryptMeta.getNonce());
     }
+
+    private static OutputStream getDecryptOutputStream(OutputStream out, TinyEncryptMeta tinyEncryptMeta) {
+        AssertUtil.notNull(tinyEncryptMeta.getDataKey(), "Data key cannot be null");
+        AssertUtil.notNull(tinyEncryptMeta.getNonce(), "Nonce cannot be null");
+        return CryptOutputStream.gcmDecrypt(out, tinyEncryptMeta.getDataKey(), tinyEncryptMeta.getNonce());
+    }
 }