From 5038198256cfa5a74921e18e147ca8e733709464 Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@gmail.com>
Date: Sun, 16 Feb 2020 09:25:01 +0100
Subject: [PATCH] Change cli syntax

---
 README.md   |  4 ++--
 cryptpng.go | 36 ++++++++++++++++++++++++------------
 2 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index b542f4e..cf49567 100644
--- a/README.md
+++ b/README.md
@@ -6,10 +6,10 @@ A way to store encrypted data inside a png without altering the image itself.
 
 ```shell script
 # encrypt
-cryptpng --image <name of the image> --in <input file> --out <output file>
+cryptpng encrypt --image <name of the image> --in <input file> --out <output file>
 
 # decrypt
-cryptpng --decrypt --image <crypt image> --out <decrypted output file>
+cryptpng decrypt --image <crypt image> --out <decrypted output file>
 ```
 
 ## Technical Information
diff --git a/cryptpng.go b/cryptpng.go
index dbea358..54ac289 100644
--- a/cryptpng.go
+++ b/cryptpng.go
@@ -37,12 +37,20 @@ var outputFile string
 var imageFile string
 var decryptImage bool
 func main() {
-	flag.StringVar(&imageFile, "image", "image.png", "The path of the png file.")
-	flag.BoolVar(&decryptImage, "decrypt", false, "If the input image should be decrypted.")
-	flag.StringVar(&outputFile, "out", "out.png", "The output file for the encrypted/decrypted data.")
-	flag.StringVar(&inputFile, "in", "input.txt","The file with the input data.")
+	encryptFlags := flag.NewFlagSet("encrypt", flag.ContinueOnError)
+	decryptFlags := flag.NewFlagSet("decrypt", flag.ContinueOnError)
+	encryptFlags.StringVar(&imageFile, "img", "image.png", "The path of the png file.")
+	encryptFlags.StringVar(&imageFile, "image", "image.png", "The path of the png file.")
+	decryptFlags.StringVar(&imageFile, "img", "image.png", "The path of the png file.")
+	decryptFlags.StringVar(&imageFile, "image", "image.png", "The path of the png file.")
+	encryptFlags.StringVar(&inputFile, "in", "input.txt","The file with the input data.")
+	encryptFlags.StringVar(&outputFile, "out", "out.png", "The output filename for the image.")
+	decryptFlags.StringVar(&outputFile, "out", "out.txt", "The output file for the decrypted data.")
 	flag.Parse()
-	if decryptImage {
+	switch os.Args[1] {
+	case "encrypt":
+		err := encryptFlags.Parse(os.Args[2:])
+		check(err)
 		f, err := os.Open(imageFile)
 		check(err)
 		defer f.Close()
@@ -50,8 +58,13 @@ func main() {
 		fout, err := os.Create(outputFile)
 		check(err)
 		defer fout.Close()
-		DecryptDataPng(f, fout)
-	} else {
+		fin, err := os.Open(inputFile)
+		check(err)
+		defer fin.Close()
+		EncryptDataPng(f, fin, fout)
+	case "decrypt":
+		err := decryptFlags.Parse(os.Args[2:])
+		check(err)
 		f, err := os.Open(imageFile)
 		check(err)
 		defer f.Close()
@@ -59,10 +72,7 @@ func main() {
 		fout, err := os.Create(outputFile)
 		check(err)
 		defer fout.Close()
-		fin, err := os.Open(inputFile)
-		check(err)
-		defer fin.Close()
-		EncryptDataPng(f, fin, fout)
+		DecryptDataPng(f, fout)
 	}
 }
 
@@ -149,6 +159,7 @@ func readPassword(passwordSalt *[]byte) ([]byte, []byte) {
 	}
 }
 
+// function to encrypt the data
 func encrypt(key, data []byte) ([]byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
@@ -164,13 +175,14 @@ func encrypt(key, data []byte) ([]byte, error) {
 	return cipherText, nil
 }
 
+// function to decrypt the data
 func decrypt(key, data []byte) ([]byte, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
 		return nil, err
 	}
 	if len(data) < aes.BlockSize {
-		return nil, errors.New("ciphertext too short")
+		return nil, errors.New("data too short")
 	}
 	iv := data[:aes.BlockSize]
 	data = data[aes.BlockSize:]