package main

import (
	"bytes"
	"crypto/rsa"
	"errors"
	"fmt"
	"math/big"
	"os"
	"strconv"
	"strings"
)

type keytype uint

const (
	Plan9	keytype = iota
	SSH
)

var pmap = map[keytype]struct{
	s	string
	pub	func(string) (*rsa.PublicKey, error)
	priv	func(string) (*rsa.PrivateKey, error)
}{
	Plan9:	{"plan9",	plan9pub,	plan9priv},
	SSH:	{"ssh",	sshpub,	sshpriv},
}

func plan9pairs(key string) (map[string]string, error) {
	fields := strings.Fields(key)
	if len(fields) < 1 {
		return nil, errors.New("missing fields")
	}
	pairs := make(map[string]string, 16)
	for _, pair := range fields[1:] {	// skip over "key"
		kv := strings.Split(pair, "=")
		if len(kv) < 2 {
			continue
		}
		k := kv[0]
		v := kv[1]
		pairs[k] = v
	}

	return pairs, nil
}

func plan9pubpairs(pairs map[string]string) (*rsa.PublicKey, error) {
	if pairs["proto"] != "rsa" {
		return nil, errors.New("proto != rsa")
	}

	ek, err := strconv.ParseInt(pairs["ek"], 16, 0)
	if err != nil {
		return nil, fmt.Errorf("ek: %s", err.Error())
	}

	n := new(big.Int)
	if _, ok := n.SetString(pairs["n"], 16); !ok {
		return nil, errors.New("can't parse n")
	}

	pub := &rsa.PublicKey{
		N:	n,
		E:	int(ek),
	}
	return pub, nil
}

func plan9pub(key string) (*rsa.PublicKey, error) {
	pairs, err := plan9pairs(key)
	if err != nil {
		return nil, err
	}

	pub, err := plan9pubpairs(pairs)
	if err != nil {
		return nil, err
	}
	return pub, nil
}

func plan9priv(key string) (*rsa.PrivateKey, error) {
	pairs, err := plan9pairs(key)
	if err != nil {
		return nil, err
	}

	if pairs["proto"] != "rsa" {
		return nil, errors.New("proto != rsa")
	}

	dk := new(big.Int)
	if _, ok := dk.SetString(pairs["!dk"], 16); !ok {
		return nil, errors.New("can't parse dk")
	}

	p := new(big.Int)
	if _, ok := p.SetString(pairs["!p"], 16); !ok {
		return nil, errors.New("can't parse p")
	}

	q := new(big.Int)
	if _, ok := q.SetString(pairs["!q"], 16); !ok {
		return nil, errors.New("can't parse q")
	}

	pub, err := plan9pubpairs(pairs)
	if err != nil {
		return nil, fmt.Errorf("can't parse pub: %s", err.Error())
	}

	priv := &rsa.PrivateKey{
		PublicKey:	*pub,
		D:	dk,
		Primes:	[]*big.Int{p, q},
	}
	return priv, nil
}

func sshpub(key string) (*rsa.PublicKey, error) {
	fields := strings.Fields(key)
	if len(fields) < 3 {
		return nil, errors.New("missing fields")
	}

	ek, err := strconv.ParseInt(fields[1], 10, 0)
	if err != nil {
		return nil, fmt.Errorf("ek: %s", err.Error())
	}

	n := new(big.Int)
	if _, ok := n.SetString(fields[2], 10); !ok {
		return nil, errors.New("can't parse n")
	}

	pub := &rsa.PublicKey{
		N:	n,
		E:	int(ek),
	}
	return pub, nil
}

func sshpriv(string) (*rsa.PrivateKey, error) {
	return nil, errors.New("can't parse private ssh keys")
}

func identKey(key string) keytype {
	if strings.HasPrefix(key, "key") {
		return Plan9
	} else {
		return SSH
	}
	panic("bad programming")
}

func ParsePublicKey(key string) (*rsa.PublicKey, error) {
	parser, ok := pmap[identKey(key)]
	if !ok {
		return nil, errors.New("unknown key type")
	}

	pub, err := parser.pub(key)
	if err != nil {
		return nil, fmt.Errorf("%s: %s", parser.s, err.Error())
	}
	return pub, nil
}

func ParsePublicKeyFile(filename string) (*rsa.PublicKey, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	buf := new(bytes.Buffer)
	if _, err := buf.ReadFrom(f); err != nil {
		return nil, err
	}
	return ParsePublicKey(buf.String())
}

func ParsePrivateKey(key string) (*rsa.PrivateKey, error) {
	parser, ok := pmap[identKey(key)]
	if !ok {
		return nil, errors.New("unknown key type")
	}

	priv, err := parser.priv(key)
	if err != nil {
		return nil, fmt.Errorf("%s: %s", parser.s, err.Error())
	}
	return priv, nil
}

func ParsePrivateKeyFile(filename string) (*rsa.PrivateKey, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	buf := new(bytes.Buffer)
	if _, err := buf.ReadFrom(f); err != nil {
		return nil, err
	}
	return ParsePrivateKey(buf.String())
}