package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/asn1"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

const KeyBitSize = 4096

func readFile(filename string) ([]byte, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}

	return ioutil.ReadAll(file)
}

func writeFile(filename string, data []byte) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}

	_, err = file.Write(data)
	return err
}

func signData(data []byte, key *rsa.PublicKey) (string, error) {
	based := base64.StdEncoding.EncodeToString(data)
	h := sha256.Sum256([]byte(based))
	dataHash := fmt.Sprintf("%x", h)


	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, key, []byte(dataHash))
	if err != nil {
		return "", err
	}

	signature := base64.StdEncoding.EncodeToString(ciphertext)
	return fmt.Sprintf("%s.%s", based, signature), nil
}

func checkSignature(signedData string, key *rsa.PrivateKey) (bool, error) {
	spl := strings.Split(signedData, ".")
	data := spl[0]
	signature := spl[1]

	dataHash := fmt.Sprintf("%x", sha256.Sum256([]byte(data)))

	decodedSign, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return false, err
	}

	signHash, err := decodeMessage(decodedSign, key)
	if err != nil {
		return false, err
	}

	return dataHash == string(signHash), nil
}

func encodeMessage(message []byte, key *rsa.PrivateKey) ([]byte, error) {
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, &key.PublicKey, message)
	if err != nil {
		return nil, err
	}
	return ciphertext, nil
}

func decodeMessage(ciphertext []byte, key *rsa.PrivateKey) ([]byte, error) {
	plaintext, err := rsa.DecryptPKCS1v15(rand.Reader, key, ciphertext)
	if err != nil { return nil, err }
	return plaintext, nil
}

func loadPublicPEM(filename string) (*rsa.PublicKey, error) {
	pemBytes, err := readFile(filename)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, errors.New("there is no key in file")
	}

	return x509.ParsePKCS1PublicKey(block.Bytes)
}

func loadPrivatePEM(filename string) (*rsa.PrivateKey, error) {
	pemBytes, err := readFile(filename)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, errors.New("there is no key in file")
	}

	return x509.ParsePKCS1PrivateKey(block.Bytes)
}

func generateKeys() (keys *rsa.PrivateKey, private []byte, public []byte, err error) {
	key, err := rsa.GenerateKey(rand.Reader, KeyBitSize)
	if err != nil {
		return nil, nil, nil, err
	}

	var privateKey = &pem.Block{
		Type: "PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}

	asn1Bytes, _ := asn1.Marshal(key.PublicKey)
	var publicKey = &pem.Block{
		Type: "PUBLIC KEY",
		Bytes: asn1Bytes,
	}

	bytePrivate := pem.EncodeToMemory(privateKey)
	bytePublic := pem.EncodeToMemory(publicKey)
	return key, bytePrivate, bytePublic, nil
}