Merge pull request #2 from kongjun18/main

Some Improvement

5 files changed, 62 insertions, 16 deletions

diff --git a/lua/encrypt/decrypt.lua b/lua/encrypt/decrypt.lua
index 2ae4e2a..a1d6b66 100644
--- a/lua/encrypt/decrypt.lua
+++ b/lua/encrypt/decrypt.lua
@@ -5,13 +5,17 @@ local helpers = require("encrypt.helpers")
 ---@param password string
 local function decrypt_lines(lines, password)
   return vim.fn.systemlist(
-    "base64 --decode | openssl enc -d -aes-256-cbc -pbkdf2 -salt -in - -out - -k " .. password,
+    "base64 --decode | openssl enc -d -aes-256-cbc -pbkdf2 -salt -in - -out - -k " .. vim.fn.shellescape(password),
     lines
   )
 end
 
 local function decrypt()
   local password = helpers.getPassword()
+  if not password then
+    vim.notify("Password can not be empty", vim.log.levels.ERROR)
+    return
+  end
   local encrypted_lines = vim.api.nvim_buf_get_lines(0, 1, -1, false)
   local decrypted_lines = decrypt_lines(encrypted_lines, password)
   if vim.v.shell_error ~= 0 then
@@ -20,6 +24,7 @@ local function decrypt()
   end
 
   vim.fn.timer_start(0, function()
+    vim.bo.modifiable = true
     vim.api.nvim_buf_set_lines(0, 0, -1, false, decrypted_lines)
     vim.bo.modified = false
   end)
diff --git a/lua/encrypt/encrypt.lua b/lua/encrypt/encrypt.lua
index 5cdf8a9..c7efb1f 100644
--- a/lua/encrypt/encrypt.lua
+++ b/lua/encrypt/encrypt.lua
@@ -5,13 +5,17 @@ local helpers = require("encrypt.helpers")
 ---@param password string
 local function encrypt_lines(lines, password)
   return vim.fn.systemlist(
-    "openssl enc -aes-256-cbc -pbkdf2 -salt -in - -out - -k " .. password .. " | base64",
+    "openssl enc -aes-256-cbc -pbkdf2 -salt -in - -out - -k " .. vim.fn.shellescape(password) .. " | base64",
     lines
   )
 end
 
 local function encrypt()
   local password = helpers.getPassword()
+  if not password then
+    vim.notify("Password can not be empty", vim.log.levels.ERROR)
+    return
+  end
   local buf_lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
   local encrypted_lines = encrypt_lines(buf_lines, password)
 
diff --git a/lua/encrypt/helpers.lua b/lua/encrypt/helpers.lua
index 37a96b3..c88b448 100644
--- a/lua/encrypt/helpers.lua
+++ b/lua/encrypt/helpers.lua
@@ -1,13 +1,39 @@
----@return string
-local function getPassword()
-  local password = vim.b["password"]
-  if password == nil then
-    password = vim.fn.inputsecret("Enter password: ")
-    vim.b["password"] = password
+---@return function string
+local getPasswordFactory = function()
+  local bufferPasswordMap = {}
+
+  return function()
+    local key = string.format("%s", vim.fn.bufnr())
+    if bufferPasswordMap[key] == nil then
+      password = vim.fn.inputsecret("Enter password: ")
+      if password == "" then
+        vim.notify("Password is cancelled", vim.log.levels.WARN)
+      end
+      bufferPasswordMap[key] = password
+    end
+    return bufferPasswordMap[key]
   end
-  return password
 end
 
+local getPassword = getPasswordFactory()
 local ENCRYPTED_PREFIX = "# <<<encrypted>>>"
 
-return { getPassword = getPassword, ENCRYPTED_PREFIX = ENCRYPTED_PREFIX }
+---@enum buftype
+local BUFTYPE = {
+  encrypted = 1,
+  plaintext = 2,
+}
+
+local function bufferEncrypted()
+  local first_line = vim.api.nvim_buf_get_lines(0, 0, 1, false)[1]
+  if first_line == ENCRYPTED_PREFIX then
+    return true
+  end
+end
+
+return {
+  getPassword = getPassword,
+  ENCRYPTED_PREFIX = ENCRYPTED_PREFIX,
+  BUFTYPE = BUFTYPE,
+  bufferEncrypted = bufferEncrypted,
+}
diff --git a/lua/encrypt/init.lua b/lua/encrypt/init.lua
index 53fc259..ea5368f 100644
--- a/lua/encrypt/init.lua
+++ b/lua/encrypt/init.lua
@@ -5,15 +5,19 @@ local helpers = require("encrypt.helpers")
 
 local function setup()
   vim.api.nvim_create_user_command("X", function()
-    setupBuffer()
-    encrypt()
+    if helpers.bufferEncrypted() then
+      setupBuffer(helpers.BUFTYPE.encrypted)
+      decrypt()
+    else
+      setupBuffer(helpers.BUFTYPE.plaintext)
+      encrypt()
+    end
   end, {})
 
   vim.api.nvim_create_autocmd({ "BufReadPost" }, {
     callback = function()
-      local first_line = vim.api.nvim_buf_get_lines(0, 0, 1, false)[1]
-      if first_line == helpers.ENCRYPTED_PREFIX then
-        setupBuffer()
+      if helpers.bufferEncrypted() then
+        setupBuffer(helpers.BUFTYPE.encrypted)
         decrypt()
       end
     end,
diff --git a/lua/encrypt/setup-buffer.lua b/lua/encrypt/setup-buffer.lua
index 8b950ec..e03f71c 100644
--- a/lua/encrypt/setup-buffer.lua
+++ b/lua/encrypt/setup-buffer.lua
@@ -1,4 +1,5 @@
 local encrypt = require("encrypt.encrypt")
+local helpers = require("encrypt.helpers")
 
 local function createWriteAutoCmd()
   if vim.b["encryptionAutoCmd"] ~= nil then
@@ -18,12 +19,18 @@ local function createWriteAutoCmd()
   })
 end
 
-local function setupBuffer()
+---@param buftype buftype
+local function setupBuffer(buftype)
   -- buftype="acwrite" means save using the BufWriteCmd command
   vim.bo.buftype = "acwrite"
   vim.bo.swapfile = false
   vim.bo.undofile = false
   vim.b["encrypted"] = true
+  -- Lock the encrypted file to read-only mode to avoid corruption
+  -- caused by unintended modifications.
+  if buftype == helpers.BUFTYPE.encrypted then
+    vim.bo.modifiable = false
+  end
   createWriteAutoCmd()
 end