package shttp

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"fmt"

	"golang.org/x/crypto/ssh"
)

type RSA struct {
	bits    int
	key     *rsa.PrivateKey
	Comment string
}

func (r *RSA) SetBits(b int) {
	r.bits = b
}

func (r *RSA) New() (err error) {
	if r.bits < 1024 || r.bits > 4096 {
		r.bits = 2048
	}
	r.key, err = rsa.GenerateKey(rand.Reader, r.bits)
	return
}

func (r *RSA) Key() *rsa.PrivateKey {
	return r.key
}

func (r *RSA) PublicKey() string {
	if r.key == nil {
		return ""
	}
	return string(pem.EncodeToMemory(&pem.Block{Type: "RSA PUBLIC KEY", Bytes: x509.MarshalPKCS1PublicKey(&r.key.PublicKey)}))
}

func (r *RSA) PrivateKey() string {
	if r.key == nil {
		return ""
	}
	return string(pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(r.key)}))
}

func (r *RSA) EncryptByte(msg []byte) ([]byte, error) {
	if r.key == nil {
		return nil, fmt.Errorf("no key found")
	}
	b, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, &r.key.PublicKey, msg, []byte(""))
	if err != nil {
		return nil, err
	}
	return b, nil
}

func (r *RSA) Encrypt(msg string) (string, error) {
	b, err := r.EncryptByte([]byte(msg))
	return hex.EncodeToString(b), err
}

func (r *RSA) DecryptByte(msg []byte) ([]byte, error) {
	if r.key == nil {
		return nil, fmt.Errorf("no key found")
	}
	b, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, r.key, msg, []byte(""))
	return b, err
}

func (r *RSA) Decrypt(cypher string) (string, error) {
	block, err := hex.DecodeString(cypher)
	if err != nil {
		return "", err
	}
	msg, err := r.DecryptByte(block)
	if err != nil {
		return "", err
	}
	return string(msg), nil
}

func (r *RSA) LoadPublic(pub string) error {
	if r.key == nil {
		r.key = &rsa.PrivateKey{}
	}
	parsed, comment, _, _, err := ssh.ParseAuthorizedKey([]byte(pub))
	if err != nil {
		block, rest := pem.Decode([]byte(pub))
		if block == nil || block.Type != "RSA PUBLIC KEY" {
			return fmt.Errorf("ssh: no key found")
		}
		if len(rest) != 0 {
			r.Comment = string(rest)
		}
		p, err := x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return err
		}
		r.key.PublicKey = *p
		return nil
	}
	parsedCryptoKey := parsed.(ssh.CryptoPublicKey)
	pubCrypto := parsedCryptoKey.CryptoPublicKey()
	p := pubCrypto.(*rsa.PublicKey)

	r.key.PublicKey = *p
	r.Comment = comment
	return nil
}

func (r *RSA) LoadPrivate(pri string) error {
	var p *rsa.PrivateKey
	var err error
	opssh, err := ssh.ParseRawPrivateKey([]byte(pri))
	if err != nil {
		return err
	}
	p = opssh.(*rsa.PrivateKey)
	if r.key != nil {
		p.PublicKey = r.key.PublicKey
	}
	r.key = p
	return nil
}

func (r *RSA) LoadPrivateWithPassphrase(pri, passphrase string) error {
	var p *rsa.PrivateKey
	var err error
	opssh, err := ssh.ParseRawPrivateKeyWithPassphrase([]byte(pri), []byte(passphrase))
	if err != nil {
		return err
	}
	p = opssh.(*rsa.PrivateKey)
	if r.key != nil {
		p.PublicKey = r.key.PublicKey
	}
	r.key = p
	return nil
}