package main

import (
	"bytes"
	"log"

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

type proxy struct {
	secretiveAgent agent.ExtendedAgent
	keyringAgent   agent.ExtendedAgent
}

func NewProxy(secretiveAgent agent.ExtendedAgent, keyringAgent agent.ExtendedAgent) agent.ExtendedAgent {
	return &proxy{
		secretiveAgent: secretiveAgent,
		keyringAgent:   keyringAgent,
	}
}

func (p *proxy) List() ([]*agent.Key, error) {
	secretiveKeys, err := p.secretiveAgent.List()
	if err != nil {
		return nil, err
	}

	keyringKeys, err := p.keyringAgent.List()
	if err != nil {
		return nil, err
	}

	var allKeys []*agent.Key
	allKeys = append(allKeys, secretiveKeys...)
	allKeys = append(allKeys, keyringKeys...)

	return allKeys, nil
}

func (p *proxy) Add(key agent.AddedKey) error {
	log.Printf("added key %s\n", key.Comment)
	return p.keyringAgent.Add(key)
}

func (p *proxy) Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error) {
	return p.SignWithFlags(key, data, 0)
}

func (p *proxy) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) {
	keyringKeys, err := p.keyringAgent.List()
	if err != nil {
		return nil, err
	}

	for _, k := range keyringKeys {
		if bytes.Equal(k.Blob, key.Marshal()) {
			log.Printf("signing with keyring\n")
			return p.keyringAgent.SignWithFlags(key, data, flags)
		}
	}

	log.Printf("signing with secretive\n")
	return p.secretiveAgent.SignWithFlags(key, data, flags)
}

func (p *proxy) Remove(key ssh.PublicKey) error {
	log.Printf("removed key\n")
	return p.keyringAgent.Remove(key)
}

func (p *proxy) RemoveAll() error {
	log.Printf("cleared keys\n")
	return p.keyringAgent.RemoveAll()
}

func (p *proxy) Lock(passphrase []byte) error {
	return p.keyringAgent.Lock(passphrase)
}

func (p *proxy) Unlock(passphrase []byte) error {
	return p.keyringAgent.Unlock(passphrase)
}

func (p *proxy) Signers() ([]ssh.Signer, error) {
	secretiveSigners, err := p.secretiveAgent.Signers()
	if err != nil {
		return nil, err
	}

	keyringSigners, err := p.keyringAgent.Signers()
	if err != nil {
		return nil, err
	}

	var allSigners []ssh.Signer
	allSigners = append(allSigners, secretiveSigners...)
	allSigners = append(allSigners, keyringSigners...)

	return allSigners, nil
}

// The keyring does not support any extensions
func (p *proxy) Extension(extensionType string, contents []byte) ([]byte, error) {
	return nil, agent.ErrExtensionUnsupported
}