/*
gohpenc -- Go high-performance encryption utility
Copyright (C) 2017-2022 Sergey Matveev <stargrave@stargrave.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, version 3 of the License.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

// Go high-performance encryption utility
package main

import (
	"bufio"
	"crypto/rand"
	"crypto/sha512"
	"encoding/base32"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sync"

	"golang.org/x/crypto/chacha20"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/hkdf"
	"golang.org/x/crypto/poly1305"
)

const (
	Magic    = "GOHPENC\n"
	SaltSize = 16
)

var Base32Codec *base32.Encoding = base32.StdEncoding.WithPadding(base32.NoPadding)

type Worker struct {
	ctr      uint64
	buf      []byte
	readyIn  chan struct{}
	readyOut chan struct{}
	last     bool
}

func readBuf(dst []byte, src io.Reader) ([]byte, error) {
	var n, full int
	var err error
	for full < len(dst) {
		n, err = src.Read(dst[full:])
		full += n
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, err
		}
	}
	return dst[:full], err
}

type DummyReader struct{}

func (r *DummyReader) Read(b []byte) (int, error) {
	return len(b), nil
}

func main() {
	doRNG := flag.Bool("r", false, "Random number generator")
	doGen := flag.Bool("psk", false, "Generate key")
	doDec := flag.Bool("d", false, "Decrypt, instead of encrypt")
	bs := flag.Int("b", 1<<10, "Blocksize, KiB")
	jobs := flag.Int("c", runtime.NumCPU(), "Number of parallel threads")
	keyB32 := flag.String("k", "", "Encryption key")
	flag.Parse()
	log.SetFlags(log.Ldate | log.Lmicroseconds | log.Lshortfile)

	var err error
	key := make([]byte, chacha20poly1305.KeySize)
	if *doGen || *doRNG {
		if _, err := io.ReadFull(rand.Reader, key); err != nil {
			log.Fatalln(err)
		}
		if *doGen {
			fmt.Println(Base32Codec.EncodeToString(key))
			return
		}
	} else {
		key, err = Base32Codec.DecodeString(*keyB32)
		if err != nil {
			log.Fatalln(err)
		}
		if len(key) != chacha20poly1305.KeySize {
			log.Fatalln("invalid key size")
		}
	}

	salt := make([]byte, SaltSize)
	if *doDec {
		if _, err = io.ReadFull(os.Stdin, salt[:len(Magic)]); err != nil {
			log.Fatalln(err)
		}
		if string(salt[:len(Magic)]) != Magic {
			log.Fatalln("invalid magic")
		}
		if _, err = io.ReadFull(os.Stdin, salt[:4]); err != nil {
			log.Fatalln(err)
		}
		*bs = int(binary.BigEndian.Uint32(salt[:4]))
		if _, err = io.ReadFull(os.Stdin, salt); err != nil {
			log.Fatalln(err)
		}
	} else {
		if _, err = os.Stdout.WriteString(Magic); err != nil {
			log.Fatalln(err)
		}
		*bs = *bs * 1024
		binary.BigEndian.PutUint32(salt, uint32(*bs))
		if _, err = os.Stdout.Write(salt[:4]); err != nil {
			log.Fatalln(err)
		}
		if _, err = io.ReadFull(rand.Reader, salt); err != nil {
			log.Fatalln(err)
		}
		if _, err = os.Stdout.Write(salt); err != nil {
			log.Fatalln(err)
		}
	}

	kdf := hkdf.New(sha512.New, key, salt, []byte(Magic))
	if _, err = io.ReadFull(kdf, key); err != nil {
		log.Fatalln(err)
	}

	var wg sync.WaitGroup
	var lastMet bool
	workers := make([]*Worker, 0, *jobs)
	for i := 0; i < *jobs; i++ {
		w := Worker{
			buf:      make([]byte, *bs+chacha20poly1305.Overhead),
			readyIn:  make(chan struct{}),
			readyOut: make(chan struct{}),
		}
		go func() {
			ciph, err := chacha20poly1305.New(key)
			if err != nil {
				log.Fatalln(err)
			}
			nonce := make([]byte, chacha20poly1305.NonceSize)
			var ciphertext, tag []byte
			var s *chacha20.Cipher
			var p *poly1305.MAC
			for {
				w.readyIn <- struct{}{}
				<-w.readyIn
				binary.BigEndian.PutUint64(nonce, w.ctr)
				if *doDec {
					tag = w.buf[len(w.buf)-poly1305.TagSize:]
					ciphertext = w.buf[:len(w.buf)-poly1305.TagSize]
					s, err = chacha20.NewUnauthenticatedCipher(key, nonce)
					if err != nil {
						log.Fatalln(err)
					}
					var polyKey [32]byte
					s.XORKeyStream(polyKey[:], polyKey[:])
					s.SetCounter(1)
					p = poly1305.New(&polyKey)
					writeWithPadding(p, nil)
					writeWithPadding(p, ciphertext)
					writeUint64(p, 0)
					writeUint64(p, len(ciphertext))
					if p.Verify(tag) {
						w.buf = ciphertext
						s.XORKeyStream(ciphertext, ciphertext)
					} else {
						lastMet = true
						if _, err = io.ReadFull(kdf, key); err != nil {
							log.Fatalln(err)
						}
						ciph, err = chacha20poly1305.New(key)
						if err != nil {
							log.Fatalln(err)
						}
						w.buf, err = ciph.Open(w.buf[:0], nonce, w.buf, nil)
						if err != nil {
							log.Fatalln(err)
						}
						lastMet = true
					}
				} else {
					if w.last {
						if _, err = io.ReadFull(kdf, key); err != nil {
							log.Fatalln(err)
						}
						ciph, err = chacha20poly1305.New(key)
						if err != nil {
							log.Fatalln(err)
						}
					}
					w.buf = ciph.Seal(w.buf[:0], nonce, w.buf, nil)
				}
				w.readyOut <- struct{}{}
				<-w.readyOut
				wg.Done()
			}
		}()
		workers = append(workers, &w)
	}

	go func() {
		var ctr int64
		var w *Worker
		var err error
		for {
			w = workers[ctr%int64(len(workers))]
			<-w.readyOut
			if _, err = os.Stdout.Write(w.buf); err != nil {
				log.Fatalln(err)
			}
			w.readyOut <- struct{}{}
			ctr++
		}
	}()

	var stdin io.Reader
	if *doRNG {
		stdin = &DummyReader{}
	} else {
		stdin = bufio.NewReaderSize(os.Stdin, *bs+chacha20poly1305.Overhead)
	}
	var ctr uint64
	var w *Worker
	for {
		w = workers[ctr%uint64(len(workers))]
		<-w.readyIn
		if *doDec {
			w.buf, err = readBuf(w.buf[:*bs+chacha20poly1305.Overhead], stdin)
		} else {
			w.buf, err = readBuf(w.buf[:*bs], stdin)
		}
		if err != nil && err != io.EOF {
			log.Fatalln(err)
		}
		if *doDec && len(w.buf) < chacha20poly1305.Overhead {
			break
		}
		w.ctr = ctr
		wg.Add(1)
		if err == io.EOF {
			w.last = true
		}
		w.readyIn <- struct{}{}
		if err == io.EOF {
			break
		}
		ctr++
	}
	wg.Wait()
	if *doDec && !lastMet {
		log.Fatalln("did not meet explicit last block")
	}
}