Add option to decrypt data from an image

Add flag -decrypt that decrypts the input image and writes the data into -out

.gitignore

@@ -1,2 +1,4 @@
 *.png
+*.jpg
+*.txt
 .idea

cryptpng.go

@@ -14,60 +14,89 @@ import (
 	"os"
 )
 
-var inputFile string
-var filename string
-func main() {
-	flag.StringVar(&inputFile, "input", "input.txt","The file with the input data.")
-	flag.StringVar(&filename, "image", "image.png", "The path of the png file.")
-	flag.Parse()
-	fmt.Println(filename)
-	f, err := os.Open(filename)
+func check(err error) {
 	if err != nil {
 		log.Fatal(err)
 	}
-	defer f.Close()
-	valid, header := ValidatePng(f)
-	if valid {
-		fout, err := os.Create("encrypted-" + filename)
-		if err != nil {
-			log.Fatal(err)
 }
+
+var inputFile string
+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.")
+	flag.Parse()
+	if decryptImage {
+		f, err := os.Open(imageFile)
+		check(err)
+		defer f.Close()
+		info, _ := f.Stat()
+		fmt.Printf("size: %d\n",info.Size())
+		check(err)
+		fout, err := os.Create(outputFile)
+		check(err)
 		defer fout.Close()
-		_, _ = fout.Write(header)
-		chunk, err := ReadChunk(f)
-		if err != nil {
-			log.Fatal(err)
-		}
-		for chunk.name != "IDAT" {
-			fmt.Printf("l: %d, n: %s, c: %d\n", chunk.length,  chunk.name, chunk.crc)
-			_, _ = fout.Write(chunk.raw)
-			chunk, err = ReadChunk(f)
-			if err != nil {
-				log.Fatal(err)
+		DecryptDataPng(f, fout)
+	} else {
+		f, err := os.Open(imageFile)
+		check(err)
+		defer f.Close()
+		check(err)
+		fout, err := os.Create(outputFile)
+		check(err)
+		defer fout.Close()
+		fin, err := os.Open(inputFile)
+		check(err)
+		defer fin.Close()
+		EncryptDataPng(f, fin, fout)
 	}
 }
-		reader := bufio.NewReader(os.Stdin)
-		fmt.Print("Data to encrypt: ")
-		inputDataString, _ := reader.ReadString('\n')
-		inputData := []byte(inputDataString)
+
+func EncryptDataPng(f *os.File, fin *os.File, fout *os.File) {
+	png := PngData{}
+	err := png.Read(f)
+	check(err)
+	inputData := readFileFull(fin)
 	inputData, err = encryptData(inputData)
-		if err != nil {
-			panic(err)
-		}
+	check(err)
 	cryptChunk := CreateChunk(inputData, "crPt")
-		_, _ = fout.Write(cryptChunk.raw)
-		for {
-			_, _ = fout.Write(chunk.raw)
-			chunk, err = ReadChunk(f)
-			if err != nil {
-				break
-			}
+	png.AddMetaChunk(cryptChunk)
+	err = png.Write(fout)
+	check(err)
 }
+
+func DecryptDataPng(f *os.File, fout *os.File) {
+	png := PngData{}
+	err := png.Read(f)
+	check(err)
+	cryptChunk := png.GetChunk("crPt")
+	if cryptChunk != nil {
+		data, err := decryptData(cryptChunk.data)
+		check(err)
+		_, err = fout.Write(data)
+		check(err)
 	} else {
-		log.Fatal("Invalid png.")
+		log.Fatal("no encrypted data inside the input image")
 	}
 }
 
+// reads all bytes of a file
+func readFileFull(f *os.File) []byte {
+	tmp := make([]byte, 8)
+	data := make([]byte, 0)
+	_, err := f.Read(tmp)
+	for err != io.EOF {
+		data = append(data, tmp...)
+		_, err = f.Read(tmp)
+	}
+	data = append(data, tmp...)
+	return data
+}
+
 // creates an encrypted png chunk
 func encryptData(data []byte) ([]byte, error) {
 	reader := bufio.NewReader(os.Stdin)
@@ -78,6 +107,15 @@ func encryptData(data []byte) ([]byte, error) {
 	return encrypt(key, data)
 }
 
+func decryptData(data []byte) ([]byte, error) {
+	reader := bufio.NewReader(os.Stdin)
+	fmt.Print("Password: ")
+	pw, _ := reader.ReadString('\n')
+	key := make([]byte, 32 - len(pw))
+	key = append(key, []byte(pw)...)
+	return decrypt(key, data)
+}
+
 // encrypt and decrypt functions taken from
 // https://stackoverflow.com/questions/18817336/golang-encrypting-a-string-with-aes-and-base64
 

pngUtils.go

@@ -2,11 +2,13 @@ package main
 
 import (
 	"encoding/binary"
+	"errors"
 	"hash/crc32"
-	"log"
+	"io"
 	"os"
 )
 
+
 type ChunkData struct {
 	length uint32
 	name string
@@ -15,14 +17,91 @@ type ChunkData struct {
 	raw []byte
 }
 
+type PngData struct {
+	header []byte
+	chunks []ChunkData
+}
+
+// Reads the png data from a file into the struct
+func (p *PngData) Read(f *os.File) error {
+	valid, header := ValidatePng(f)
+	if valid {
+		p.header = header
+		err := p.readChunks(f)
+		if err != io.EOF {
+			return err
+		}
+	} else {
+		return errors.New("invalid png")
+	}
+	return nil
+}
+
+// writes all the data of the png into a new file
+func (p *PngData) Write(f *os.File) error {
+	_, err := f.Write(p.header)
+	if err != nil {
+		return err
+	}
+	err = p.writeChunks(f)
+	return err
+}
+
+// reads all chunks from a png file.
+// must be called after reading the header
+func (p *PngData) readChunks(f *os.File) error {
+	p.chunks = make([]ChunkData, 0)
+	chunk, err := ReadChunk(f)
+	for err == nil {
+		p.chunks = append(p.chunks, chunk)
+		chunk, err = ReadChunk(f)
+	}
+	p.chunks = append(p.chunks, chunk)
+	return err
+}
+
+// writes all chunks to the given file
+func (p *PngData) writeChunks(f *os.File) error {
+	for _, chunk := range p.chunks {
+		_, err := f.Write(chunk.raw)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// adds a meta chunk to the chunk data before the IDAT chunk.
+func (p *PngData) AddMetaChunk(metaChunk ChunkData) {
+	newChunks := make([]ChunkData, 0)
+	appended := false
+	for _, chunk := range p.chunks {
+		if chunk.name == "IDAT" && !appended {
+			newChunks = append(newChunks, metaChunk)
+			newChunks = append(newChunks, chunk)
+			appended = true
+		} else {
+			newChunks = append(newChunks, chunk)
+		}
+	}
+	p.chunks = newChunks
+}
+
+// Returns the reference of a chunk by name
+func (p *PngData) GetChunk(name string) *ChunkData {
+	for _, chunk := range p.chunks {
+		if chunk.name == name {
+			return &chunk
+		}
+	}
+	return nil
+}
 
 // validates the png by reading the header of the file
 func ValidatePng(f *os.File) (bool, []byte) {
 	headerBytes := make([]byte, 8)
 	_, err := f.Read(headerBytes)
-	if err != nil {
-		log.Fatal(err)
-	}
+	check(err)
 	firstByteMatch := headerBytes[0] == 0x89
 	pngAsciiMatch := string(headerBytes[1:4]) == "PNG"
 	dosCRLF := headerBytes[4] == 0x0d && headerBytes[5] == 0x0a