package main

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"fmt"
	"io"
	"os"

	"github.com/spacemonkeygo/openssl"
)

const (
	ChunkSize = 10 * 1024
	TagSize   = 16
)

func EncryptFile(inFile, outFile *os.File, key, iv []byte) error {
	ctx, err := openssl.NewGCMEncryptionCipherCtx(len(key)*8, nil, key, iv)
	if err != nil {
		return fmt.Errorf("Failed making GCM encryption ctx: %v", err)
	}

	reader := bufio.NewReader(inFile)
	chunk := make([]byte, ChunkSize)
	for {
		chunkSize, err := reader.Read(chunk)
		if err == io.EOF || chunkSize == 0 {
			break
		} else if err != nil {
			return fmt.Errorf("Failed to read a chunk: %v", err)
		}

		encData, err := ctx.EncryptUpdate(chunk[:chunkSize])
		if err != nil {
			return fmt.Errorf("Failed to perform an encryption: %v", err)
		}

		if _, err := outFile.Write(encData); err != nil {
			return fmt.Errorf("Failed to write an encrypted data: %v", err)
		}
	}

	encData, err := ctx.EncryptFinal()
	if err != nil {
		return fmt.Errorf("Failed to finalize encryption: %v", err)
	}
	if _, err := outFile.Write(encData); err != nil {
		return fmt.Errorf("Failed to write a final encrypted data: %v", err)
	}

	tag, err := ctx.GetTag()
	if err != nil {
		return fmt.Errorf("Failed to get GCM tag: %v", err)
	}
	if _, err := outFile.Write(tag); err != nil {
		return fmt.Errorf("Failed to write a gcm tag: %v", err)
	}

	return nil
}

func DecryptFile(inFile, outFile *os.File, key, iv []byte, fileSize int) error {
	ctx, err := openssl.NewGCMDecryptionCipherCtx(len(key)*8, nil, key, iv)
	if err != nil {
		return fmt.Errorf("Failed making GCM decryption ctx: %v", err)
	}

	reader := bufio.NewReader(inFile)
	chunk := make([]byte, ChunkSize)
	tag := &bytes.Buffer{}
	totalChunkSize := 0

	for {
		chunkSize, err := reader.Read(chunk)
		if err == io.EOF {
			break
		} else if err != nil {
			return fmt.Errorf("Failed to read an encrypted chunk: %v", err)
		}

		totalChunkSize += chunkSize

		if totalChunkSize > fileSize {
			d := totalChunkSize % fileSize

			d %= ChunkSize
			if d == 0 {
				d = chunkSize
			}

			tag.Write(chunk[chunkSize-d : chunkSize])
			chunkSize -= d
		}

		if chunkSize > 0 {
			data, err := ctx.DecryptUpdate(chunk[:chunkSize])
			if err != nil {
				return fmt.Errorf("Failed to perform a decryption: %v", err)
			}

			if _, err := outFile.Write(data); err != nil {
				return fmt.Errorf("Failed to write a decrypted data: %v", err)
			}
		}
	}

	if err := ctx.SetTag(tag.Bytes()); err != nil {
		return fmt.Errorf("Failed to set expected GCM tag: %v", err)
	}

	data, err := ctx.DecryptFinal()
	if err != nil {
		return fmt.Errorf("Failed to finalize decryption: %v", err)
	}
	if _, err := outFile.Write(data); err != nil {
		return fmt.Errorf("Failed to write a final decrypted data: %v", err)
	}

	return nil
}

func main() {
	inFilePath := "testdata"
	outEncFilePath := "testdata.enc"
	outDecFilePath := "newtestdata"
	key, _ := hex.DecodeString("81be5e09c111576103a8507658d47891")
	iv, _ := hex.DecodeString("81be5e09c111576103a85076")

	inFile, err := os.Open(inFilePath)
	if err != nil {
		panic(err)
	}
	defer inFile.Close()

	fileStat, err := os.Stat(inFilePath)
	if err != nil {
		panic(err)
	}
	fileSize := int(fileStat.Size())

	outEncFile, err := os.Create(outEncFilePath)
	if err != nil {
		panic(err)
	}
	defer outEncFile.Close()

	outDecFile, err := os.Create(outDecFilePath)
	if err != nil {
		panic(err)
	}
	defer outDecFile.Close()

	if err := EncryptFile(inFile, outEncFile, key, iv); err != nil {
		panic(err)
	}
	outEncFile.Seek(0, 0)
	if err := DecryptFile(outEncFile, outDecFile, key, iv, fileSize); err != nil {
		panic(err)
	}
}