package main

import (
	"bufio"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"github.com/urfave/cli"
	"io/ioutil"
	"log"
	"os"
)

const BLOCKSIZE_BYTE = 16
const BLOCKSIZE_BIT = 16 * 8

func main() {
	app := cli.NewApp()
	app.Name = "pscrypt"
	app.Usage = "pscrypt "
	app.Description = "Encrypt video files in AES CTR provided a key and optionally and IV"
	app.Version = "0.1.0"
	app.Action = pscrypt

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "in, i",
			Usage: "Input file",
		},
		cli.StringFlag{
			Name:  "out, o",
			Usage: "Output file",
		},
		cli.StringFlag{
			Name:  "key, k",
			Usage: fmt.Sprintf("Key as %d-byte hex string (%d characters)", BLOCKSIZE_BYTE, hex.EncodedLen(BLOCKSIZE_BYTE)),
		},
		cli.BoolFlag{
			Name:  "implicit-encrypt, E",
			Usage: fmt.Sprintf("Generate a random %d-byte IV and prepend it to the output file", BLOCKSIZE_BYTE),
		},
		cli.BoolFlag{
			Name:  "implicit-decrypt, D",
			Usage: fmt.Sprintf("Use first %d bytes of file as IV for decryption", BLOCKSIZE_BYTE),
		},
		cli.StringFlag{
			Name:  "iv",
			Usage: fmt.Sprintf("IV as %d-byte hex string (%d characters)", BLOCKSIZE_BYTE, hex.EncodedLen(BLOCKSIZE_BYTE)),
			Value: "",
		},
	}

	app.Run(os.Args)
}

func pscrypt(ctx *cli.Context) error {
	if ctx.String("in") == "" ||
		ctx.String("out") == "" ||
		ctx.String("key") == "" ||
		ctx.Bool("implicit-decrypt") && ctx.Bool("implicit-encrypt") ||
		(ctx.String("iv") == "" && !(ctx.Bool("implicit-decrypt") || ctx.Bool("implicit-encrypt"))) {

		cli.ShowAppHelp(ctx)
		return nil
	}

	in, err := os.Open(ctx.String("in"))
	check(err)

	ir := bufio.NewReader(in)

	out, err := os.Create(ctx.String("out"))
	check(err)

	ow := bufio.NewWriter(out)

	key := make([]byte, BLOCKSIZE_BYTE)
	n, err := hex.Decode(key, []byte(ctx.String("key")))
	check(err)
	if n != BLOCKSIZE_BYTE {
		log.Fatalf("Provided key is not %d b long", BLOCKSIZE_BIT)
	}

	iv := make([]byte, BLOCKSIZE_BYTE)

	source, err := ioutil.ReadAll(ir)

	if ctx.String("iv") != "" {
		n, err = hex.Decode(key, []byte(ctx.String("iv")))
		check(err)
		if n != BLOCKSIZE_BYTE {
			log.Fatalf("Provided IV is not %d b long", BLOCKSIZE_BIT)
		}
	} else {
		if ctx.Bool("implicit-decrypt") {
			copy(iv, source[:BLOCKSIZE_BYTE])
			source = source[BLOCKSIZE_BYTE:]
		} else if ctx.Bool("implicit-encrypt") {
			n, err = rand.Read(iv)
			check(err)
			if n != BLOCKSIZE_BYTE {
				log.Fatalf("Error gathering %d random bytes for IV", BLOCKSIZE_BIT)
			}
			ow.Write(iv)
		} else {
			log.Fatal("You need to provide an IV or use -I/-E option.")
		}
	}

	aesCipher, err := aes.NewCipher(key)
	check(err)

	ctr := cipher.NewCTR(aesCipher, iv)

	destination := make([]byte, len(source))
	check(err)

	ctr.XORKeyStream(destination, source)

	ow.Write(destination)

	return nil
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}