// remoteCmdOutput runs the given command on a remote server at the given hostname as the given user.
func remoteCmdOutput(username, hostname, cmd string, privateKey []byte) (b []byte, err error) {
	p, err := ssh.ParseRawPrivateKey(privateKey)
	if err != nil {
		return b, err
	}
	s, err := ssh.NewSignerFromKey(p)
	if err != nil {
		return b, err
	}
	pub := ssh.PublicKeys(s)
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{pub},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		return b, errors.New("ERROR: Failed to dial: " + err.Error())
	}
	defer client.Close()
	session, err := client.NewSession()
	if err != nil {
		return b, errors.New("ERROR: Failed to create session: " + err.Error())
	}
	defer session.Close()
	b, err = session.Output(cmd)
	if err != nil {
		return b, fmt.Errorf("ERROR: Failed to run cmd on host %s: %s", hostname, err.Error())
	}
	return b, nil
}

// NewLinuxKVMCoreOSFactory returns a new HostedProgramFactory that can
// create docker containers to wrap programs.
// TODO(kwalsh) fix comment.
func NewLinuxKVMCoreOSHostFactory(sockPath string, cfg *CoreOSLinuxhostConfig) (HostedProgramFactory, error) {

	// Create a key to use to connect to the instance and set up LinuxHost
	// there.
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	sshpk, err := ssh.NewPublicKey(&priv.PublicKey)
	if err != nil {
		return nil, err
	}
	pkstr := "ssh-rsa " + base64.StdEncoding.EncodeToString(sshpk.Marshal()) + " linux_host"

	sshpriv, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		return nil, err
	}

	return &LinuxKVMCoreOSHostFactory{
		Cfg:        cfg,
		SocketPath: sockPath,
		PublicKey:  pkstr,
		PrivateKey: sshpriv,
	}, nil
}

// NewRandomSigner generates a random key of a desired bit length.
func NewRandomSigner(bits int) (ssh.Signer, error) {
	key, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, err
	}
	return ssh.NewSignerFromKey(key)
}

func newSigner(privKeyFile string) (*signer, error) {
	signer := new(signer)

	if len(privKeyFile) > 0 {
		privateBytes, err := ioutil.ReadFile(privKeyFile)
		if err != nil {
			return nil, errors.New("Failed to load private host key")
		}

		signer.Signer, err = ssh.ParsePrivateKey(privateBytes)
		if err != nil {
			return nil, errors.New("Failed to parse private host key")
		}

		return signer, nil
	}

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, errors.New("Failed to generate server key pair")
	}

	signer.Signer, err = ssh.NewSignerFromKey(key)
	if err != nil {
		return nil, errors.New("Failed to extract private key from generated key pair")
	}
	signer.generated = true

	return signer, nil
}

func sshOnConn(conn net.Conn, h conf.Host) (*ssh.Client, error) {
	var auths []ssh.AuthMethod

	if h.Pass != "" {
		auths = append(auths, ssh.Password(h.Pass))
		auths = append(auths, ssh.KeyboardInteractive(kbdInteractive(h.Pass)))
	}

	if h.Key != "" {
		k := &keyring{}
		err := k.loadPEM([]byte(h.Key))
		if err != nil {
			return nil, err
		}
		for _, k := range k.keys {
			s, _ := ssh.NewSignerFromKey(k)
			auths = append(auths, ssh.PublicKeys(s))
		}
	}

	config := &ssh.ClientConfig{
		User: h.User,
		Auth: auths,
	}

	debugln("handshake & authenticate")
	cc, nc, reqs, err := ssh.NewClientConn(conn, conn.RemoteAddr().String(), config)
	if err != nil {
		return nil, err
	}
	client := ssh.NewClient(cc, nc, reqs)
	return client, nil
}

// Insert adds a private key to the keyring. If a certificate
// is given, that certificate is added as public key. Note that
// any constraints given are ignored.
func (r *keyring) Add(key AddedKey) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.locked {
		return errLocked
	}
	signer, err := ssh.NewSignerFromKey(key.PrivateKey)

	if err != nil {
		return err
	}

	if cert := key.Certificate; cert != nil {
		signer, err = ssh.NewCertSigner(cert, signer)
		if err != nil {
			return err
		}
	}

	p := privKey{
		signer:  signer,
		comment: key.Comment,
	}

	if key.LifetimeSecs > 0 {
		t := time.Now().Add(time.Duration(key.LifetimeSecs) * time.Second)
		p.expire = &t
	}

	r.keys = append(r.keys, p)

	return nil
}

func (c *client) insertCert(s interface{}, cert *ssh.Certificate, comment string, constraints []byte) error {
	var req []byte
	switch k := s.(type) {
	case *rsa.PrivateKey:
		if len(k.Primes) != 2 {
			return fmt.Errorf("agent: unsupported RSA key with %d primes", len(k.Primes))
		}
		k.Precompute()
		req = ssh.Marshal(rsaCertMsg{
			Type:        cert.Type(),
			CertBytes:   cert.Marshal(),
			D:           k.D,
			Iqmp:        k.Precomputed.Qinv,
			P:           k.Primes[0],
			Q:           k.Primes[1],
			Comments:    comment,
			Constraints: constraints,
		})
	case *dsa.PrivateKey:
		req = ssh.Marshal(dsaCertMsg{
			Type:      cert.Type(),
			CertBytes: cert.Marshal(),
			X:         k.X,
			Comments:  comment,
		})
	case *ecdsa.PrivateKey:
		req = ssh.Marshal(ecdsaCertMsg{
			Type:      cert.Type(),
			CertBytes: cert.Marshal(),
			D:         k.D,
			Comments:  comment,
		})
	default:
		return fmt.Errorf("agent: unsupported key type %T", s)
	}

	// if constraints are present then the message type needs to be changed.
	if len(constraints) != 0 {
		req[0] = agentAddIdConstrained
	}

	signer, err := ssh.NewSignerFromKey(s)
	if err != nil {
		return err
	}
	if bytes.Compare(cert.Key.Marshal(), signer.PublicKey().Marshal()) != 0 {
		return errors.New("agent: signer and cert have different public key")
	}

	resp, err := c.call(req)
	if err != nil {
		return err
	}
	if _, ok := resp.(*successAgentMsg); ok {
		return nil
	}
	return errors.New("agent: failure")
}

func loadPrivateKey(path string) (ssh.AuthMethod, error) {
	// Read file
	keyData, err := ioutil.ReadFile(path)
	if err != nil {
		log.Printf("error: could not read key file '%s': %s", path, err)
		return nil, err
	}

	// Get first PEM block
	block, _ := pem.Decode(keyData)
	if err != nil {
		log.Printf("error: no key found in file '%s': %s", path, err)
		return nil, err
	}

	// If it's encrypted...
	var (
		signer    ssh.Signer
		signerErr error
	)

	if x509.IsEncryptedPEMBlock(block) {
		// Get the passphrase
		prompt := fmt.Sprintf("Enter passphrase for key '%s': ", path)
		pass, err := speakeasy.Ask(prompt)
		if err != nil {
			log.Printf("error: error getting passphrase: %s", err)
			return nil, err
		}

		block.Bytes, err = x509.DecryptPEMBlock(block, []byte(pass))
		if err != nil {
			log.Printf("error: error decrypting key: %s", err)
			return nil, err
		}

		key, err := ParsePEMBlock(block)
		if err != nil {
			log.Printf("error: could not parse PEM block: %s", err)
			return nil, err
		}

		signer, signerErr = ssh.NewSignerFromKey(key)
	} else {
		signer, signerErr = ssh.ParsePrivateKey(keyData)
	}

	if signerErr != nil {
		log.Printf("error: error parsing private key '%s': %s", path, signerErr)
		return nil, signerErr
	}

	return ssh.PublicKeys(signer), nil
}

//export c_NewSignerFromKey
func c_NewSignerFromKey(key uint64) (uint64, int, *C.char) {
	obj, ok := GetObject(Handle(key))
	if !ok {
		return IH, ErrorCodeNotFound, C.CString(MessageNotFound)
	}
	signer, err := ssh.NewSignerFromKey(obj)
	if err != nil {
		return IH, ErrorCodeInternal, C.CString(err.Error())
	}
	return uint64(RegisterObject(&signer)), ErrorCodeSuccess, nil
}

func testLockAgent(agent Agent, t *testing.T) {
	if err := agent.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"], Comment: "comment 1"}); err != nil {
		t.Errorf("Add: %v", err)
	}
	if err := agent.Add(AddedKey{PrivateKey: testPrivateKeys["dsa"], Comment: "comment dsa"}); err != nil {
		t.Errorf("Add: %v", err)
	}
	if keys, err := agent.List(); err != nil {
		t.Errorf("List: %v", err)
	} else if len(keys) != 2 {
		t.Errorf("Want 2 keys, got %v", keys)
	}

	passphrase := []byte("secret")
	if err := agent.Lock(passphrase); err != nil {
		t.Errorf("Lock: %v", err)
	}

	if keys, err := agent.List(); err != nil {
		t.Errorf("List: %v", err)
	} else if len(keys) != 0 {
		t.Errorf("Want 0 keys, got %v", keys)
	}

	signer, _ := ssh.NewSignerFromKey(testPrivateKeys["rsa"])
	if _, err := agent.Sign(signer.PublicKey(), []byte("hello")); err == nil {
		t.Fatalf("Sign did not fail")
	}

	if err := agent.Remove(signer.PublicKey()); err == nil {
		t.Fatalf("Remove did not fail")
	}

	if err := agent.RemoveAll(); err == nil {
		t.Fatalf("RemoveAll did not fail")
	}

	if err := agent.Unlock(nil); err == nil {
		t.Errorf("Unlock with wrong passphrase succeeded")
	}
	if err := agent.Unlock(passphrase); err != nil {
		t.Errorf("Unlock: %v", err)
	}

	if err := agent.Remove(signer.PublicKey()); err != nil {
		t.Fatalf("Remove: %v", err)
	}

	if keys, err := agent.List(); err != nil {
		t.Errorf("List: %v", err)
	} else if len(keys) != 1 {
		t.Errorf("Want 1 keys, got %v", keys)
	}
}

func clientAuth(pemData []byte) (ssh.AuthMethod, error) {
	privKey, err := ssh.ParseRawPrivateKey(pemData)
	if err != nil {
		return nil, err
	}
	signer, err := ssh.NewSignerFromKey(privKey)
	if err != nil {
		return nil, err
	}
	return ssh.PublicKeys(signer), nil
}

func testAgentInterface(t *testing.T, agent Agent, key interface{}, cert *ssh.Certificate, lifetimeSecs uint32) {
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		t.Fatalf("NewSignerFromKey(%T): %v", key, err)
	}
	// The agent should start up empty.
	if keys, err := agent.List(); err != nil {
		t.Fatalf("RequestIdentities: %v", err)
	} else if len(keys) > 0 {
		t.Fatalf("got %d keys, want 0: %v", len(keys), keys)
	}

	// Attempt to insert the key, with certificate if specified.
	var pubKey ssh.PublicKey
	if cert != nil {
		err = agent.Add(AddedKey{
			PrivateKey:   key,
			Certificate:  cert,
			Comment:      "comment",
			LifetimeSecs: lifetimeSecs,
		})
		pubKey = cert
	} else {
		err = agent.Add(AddedKey{PrivateKey: key, Comment: "comment", LifetimeSecs: lifetimeSecs})
		pubKey = signer.PublicKey()
	}
	if err != nil {
		t.Fatalf("insert(%T): %v", key, err)
	}

	// Did the key get inserted successfully?
	if keys, err := agent.List(); err != nil {
		t.Fatalf("List: %v", err)
	} else if len(keys) != 1 {
		t.Fatalf("got %v, want 1 key", keys)
	} else if keys[0].Comment != "comment" {
		t.Fatalf("key comment: got %v, want %v", keys[0].Comment, "comment")
	} else if !bytes.Equal(keys[0].Blob, pubKey.Marshal()) {
		t.Fatalf("key mismatch")
	}

	// Can the agent make a valid signature?
	data := []byte("hello")
	sig, err := agent.Sign(pubKey, data)
	if err != nil {
		t.Fatalf("Sign(%s): %v", pubKey.Type(), err)
	}

	if err := pubKey.Verify(data, sig); err != nil {
		t.Fatalf("Verify(%s): %v", pubKey.Type(), err)
	}
}

func addKeyAuth(auths []ssh.AuthMethod, keypath string) []ssh.AuthMethod {
	if len(keypath) == 0 {
		return auths
	}

	// read the file
	pemBytes, err := ioutil.ReadFile(keypath)
	if err != nil {
		log.Print(err)
		os.Exit(1)
	}

	// get first pem block
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		log.Printf("no key found in %s", keypath)
		return auths
	}

	// handle plain and encrypted keyfiles
	if x509.IsEncryptedPEMBlock(block) {
		prompt := fmt.Sprintf("Enter passphrase for key '%s': ", keypath)
		pass, err := getpass(prompt)
		if err != nil {
			return auths
		}
		block.Bytes, err = x509.DecryptPEMBlock(block, []byte(pass))
		if err != nil {
			log.Print(err)
			return auths
		}
		key, err := ParsePemBlock(block)
		if err != nil {
			log.Print(err)
			return auths
		}
		signer, err := ssh.NewSignerFromKey(key)
		if err != nil {
			log.Print(err)
			return auths
		}
		return append(auths, ssh.PublicKeys(signer))
	} else {
		signer, err := ssh.ParsePrivateKey(pemBytes)
		if err != nil {
			log.Print(err)
			return auths
		}
		return append(auths, ssh.PublicKeys(signer))
	}
}

func (c *BaseCluster) sshConfig() (*ssh.ClientConfig, error) {
	if c.SSHKey == nil {
		return nil, errors.New("No SSHKey found")
	}
	signer, err := ssh.NewSignerFromKey(c.SSHKey.PrivateKey)
	if err != nil {
		return nil, err
	}
	sshConfig := &ssh.ClientConfig{
		User: c.SSHUsername,
		Auth: []ssh.AuthMethod{ssh.PublicKeys(signer)},
	}
	return sshConfig, nil
}

func (p privateKeySigner) Decrypt() (ssh.Signer, error) {
	if p.key == nil {
		pem, err := x509.DecryptPEMBlock(p.pem, []byte(p.Passphrase()))
		if err != nil {
			return nil, err
		}
		p.key, err = x509.ParsePKCS1PrivateKey(pem)
		if err != nil {
			return nil, err
		}
		p.Encrypted = false
	}
	return ssh.NewSignerFromKey(p.key)
}

func Run(username, hostname string, port uint16, seed string) ([]byte, int, error) {
	key := generateKey(seed)

	/*
		config := &ssh.ClientConfig{
			User: username,
			Auth: []ssh.ClientAuth{
				ssh.ClientAuthKeyring(&keyring{key: key}),
			},
		}
	*/

	signer, err := ssh.NewSignerFromKey(key)

	if err != nil {
		panic("Could not create a signer from the key: " + err.Error())
	}

	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
	}

	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", hostname, port), config)
	if err != nil {
		return nil, 0, err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return nil, 0, err
	}
	defer session.Close()

	bytes, err := session.Output("exit 1")
	if err != nil {
		if exiterror, ok := err.(*ssh.ExitError); ok {
			return bytes, exiterror.ExitStatus(), nil
			os.Stdout.Write(bytes)
			os.Exit(exiterror.ExitStatus())
		} else {
			return nil, 0, err
		}
	}
	return bytes, 0, nil
}

func sshKey() (pubKey string, auth ssh.AuthMethod, err error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", nil, err
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return "", nil, err
	}
	pub, err := ssh.NewPublicKey(&key.PublicKey)
	if err != nil {
		return "", nil, err
	}
	return string(ssh.MarshalAuthorizedKey(pub)), ssh.PublicKeys(signer), nil
}

func newSSHServer() (*sshServer, error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}

	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return nil, err
	}
	username := randstr(20)
	password := randstr(20)

	config := &ssh.ServerConfig{
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			// Don't worry about constant time compares, etc.
			// This is a single-use username and password
			// on a randomized port, on the loopback interface
			// only, on a friendly machine.
			var err error
			if c.User() != username || string(pass) != password {
				err = errors.New("authentication failed")
			}
			return nil, err
		},
	}
	config.AddHostKey(signer)

	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(strings.TrimPrefix(l.Addr().String(), "127.0.0.1:"))
	if err != nil {
		return nil, err
	}

	server := sshServer{
		config:   config,
		listener: l,
		port:     port,
		username: username,
		password: password,
	}

	return &server, nil
}

func (i *jupiterBrainInstance) sshClient() (*ssh.Client, error) {
	file, err := ioutil.ReadFile(i.provider.sshKeyPath)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(file)
	if block == nil {
		return nil, fmt.Errorf("ssh key does not contain a valid PEM block")
	}

	der, err := x509.DecryptPEMBlock(block, []byte(i.provider.sshKeyPassphrase))
	if err != nil {
		return nil, err
	}

	key, err := x509.ParsePKCS1PrivateKey(der)
	if err != nil {
		return nil, err
	}

	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return nil, err
	}

	var ip net.IP
	for _, ipString := range i.payload.IPAddresses {
		curIP := net.ParseIP(ipString)
		if curIP.To4() != nil {
			ip = curIP
			break
		}

	}

	if ip == nil {
		return nil, fmt.Errorf("no valid IPv4 address")
	}

	return ssh.Dial("tcp", fmt.Sprintf("%s:22", ip.String()), &ssh.ClientConfig{
		User: "travis",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
	})
}

// NewSSHServer returns a NewSSHServer instance, ready for use.
func NewSSHServer() (*SSHServer, error) {
	s := &SSHServer{}
	s.Transfers = &bytes.Buffer{}
	s.Config = &ssh.ServerConfig{
		NoClientAuth: true,
	}
	s.Commands = make(map[string]int)

	private, err := rsa.GenerateKey(rand.Reader, 2014)
	if err != nil {
		return nil, err
	}
	signer, err := ssh.NewSignerFromKey(private)
	if err != nil {
		return nil, err
	}
	s.Config.AddHostKey(signer)
	return s, nil
}