// ParseHostKeys parses the host key files.
//
// By default it looks in /etc/ssh for host keys of the patterh ssh_host_{{TYPE}}_key.
//
// Params:
// 	- keytypes ([]string): Key types to parse. Defaults to []string{rsa, dsa, ecdsa}
// 	- enableV1 (bool): Allow V1 keys. By default this is disabled.
// 	- path (string): Override the lookup pattern. If %s, it will be replaced with the keytype.
//
// Returns:
// 	[]ssh.Signer
func ParseHostKeys(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	log.Debugf(c, "Parsing ssh host keys")
	hostKeyTypes := p.Get("keytypes", []string{"rsa", "dsa", "ecdsa"}).([]string)
	pathTpl := p.Get("path", "/etc/ssh/ssh_host_%s_key").(string)
	hostKeys := make([]ssh.Signer, 0, len(hostKeyTypes))
	for _, t := range hostKeyTypes {
		path := fmt.Sprintf(pathTpl, t)

		if key, err := ioutil.ReadFile(path); err == nil {
			if hk, err := ssh.ParsePrivateKey(key); err == nil {
				log.Infof(c, "Parsed host key %s.", path)
				hostKeys = append(hostKeys, hk)
			} else {
				log.Errf(c, "Failed to parse host key %s (skipping): %s", path, err)
			}
		}
	}
	if c.Get("enableV1", false).(bool) {
		path := "/etc/ssh/ssh_host_key"
		if key, err := ioutil.ReadFile(path); err != nil {
			log.Errf(c, "Failed to read ssh_host_key")
		} else if hk, err := ssh.ParsePrivateKey(key); err == nil {
			log.Infof(c, "Parsed host key %s.", path)
			hostKeys = append(hostKeys, hk)
		} else {
			log.Errf(c, "Failed to parse host key %s: %s", path, err)
		}
	}
	return hostKeys, nil
}

// SSHConfig returns a function that can be used for the SSH communicator
// config for connecting to the instance created over SSH using the provided
// private key
func SSHConfig(username string) func(multistep.StateBag) (*ssh.ClientConfig, error) {
	return func(state multistep.StateBag) (*ssh.ClientConfig, error) {
		privateKey := state.Get("privateKey").(string) // ad hoc key
		privateKeyBastion := state.Get("privateKeyBastion").(string)

		signer, err := ssh.ParsePrivateKey([]byte(privateKey))
		if err != nil {
			return nil, fmt.Errorf("error setting up SSH config: %s", err)
		}

		authMethods := []ssh.AuthMethod{ssh.PublicKeys(signer)}

		// if we have a Bastion, that key should also be supported
		if privateKeyBastion != "" {
			signerBastion, err := ssh.ParsePrivateKey([]byte(privateKeyBastion))
			if err != nil {
				return nil, fmt.Errorf("error setting up SSH config for Bastion: %s", err)
			}

			authMethods = append(authMethods, ssh.PublicKeys(signerBastion))

		}

		return &ssh.ClientConfig{
			User: username,
			Auth: authMethods,
		}, nil
	}
}

func StartServer() (io.Closer, error) {
	hostPrivateKey, err := ioutil.ReadFile(config.KeyPath)
	if err != nil {
		return nil, err
	}

	hostPrivateKeySigner, err := ssh.ParsePrivateKey(hostPrivateKey)
	if err != nil {
		return nil, err
	}

	sshServer = &ssh.ServerConfig{
		PasswordCallback:  auth.PassAuth(),
		PublicKeyCallback: auth.KeyAuth(),
	}

	sshServer.AddHostKey(hostPrivateKeySigner)

	serverSocket, err := net.Listen("tcp", config.SshListenAddress)
	if err != nil {
		return nil, err
	}

	go func() {
		for {
			conn, err := serverSocket.Accept()
			if err != nil {
				return
			}
			go handleConnection(conn)
		}
	}()
	config.Log.Info("SSH listening on %s", config.SshListenAddress)
	return serverSocket, nil
}

func NewNativeConfig(user string, auth *Auth) (ssh.ClientConfig, error) {
	var (
		authMethods []ssh.AuthMethod
	)

	for _, k := range auth.Keys {
		key, err := ioutil.ReadFile(k)
		if err != nil {
			return ssh.ClientConfig{}, err
		}

		privateKey, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return ssh.ClientConfig{}, err
		}

		authMethods = append(authMethods, ssh.PublicKeys(privateKey))
	}

	for _, p := range auth.Passwords {
		authMethods = append(authMethods, ssh.Password(p))
	}

	return ssh.ClientConfig{
		User: user,
		Auth: authMethods,
	}, nil
}

func (suite *ServerSuite) TestClientConnection() {

	// Get signer
	signer, err := ssh.ParsePrivateKey([]byte(clientPrivateKey))
	if err != nil {
		suite.Fail("Private key could not be parsed" + err.Error())
	}

	// Configure client connection
	config := &ssh.ClientConfig{
		User: "admin",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
	}

	// Create client connection
	client, err := ssh.Dial("tcp", "127.0.0.1:9022", config)
	if err != nil {
		suite.Fail(err.Error())
		return
	}
	defer client.Close()

	// Open channel
	channel, requests, err := client.OpenChannel("/echo", []byte{})
	if err != nil {
		suite.Fail(err.Error())
		return
	}
	go ssh.DiscardRequests(requests)
	defer channel.Close()
}

//NewVagrantNode intializes a node in vagrant testbed
func NewVagrantNode(name, port, privKeyFile string) (*VagrantNode, error) {
	var (
		vnode      *VagrantNode
		err        error
		signer     ssh.Signer
		privateKey []byte
	)

	if privateKey, err = ioutil.ReadFile(privKeyFile); err != nil {
		return nil, err
	}

	if signer, err = ssh.ParsePrivateKey(privateKey); err != nil {
		return nil, err
	}

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

	vnode = &VagrantNode{Name: name}
	if vnode.client, err = ssh.Dial("tcp", fmt.Sprintf("127.0.0.1:%s", port), config); err != nil {
		return nil, err
	}

	return vnode, nil
}

func loadDefaultKeys() (auths []ssh.AuthMethod, err error) {
	k := ""
	currentUser, err := user.Current()
	defaultKeyPathA := filepath.FromSlash(currentUser.HomeDir + "/.ssh/id_rsa")
	defaultKeyPathB := filepath.FromSlash(currentUser.HomeDir + "/ssh/id_rsa")
	if fileExists(defaultKeyPathA) {
		k = defaultKeyPathA
	} else if fileExists(defaultKeyPathB) {
		k = defaultKeyPathB
	}
	if len(k) == 0 {
		err = errors.New("No key specified")
		return
	}
	pemBytes, err := ioutil.ReadFile(k)
	if err != nil {
		return
	}
	signer, err := ssh.ParsePrivateKey(pemBytes)
	if err != nil {
		return
	}
	auths = []ssh.AuthMethod{ssh.PublicKeys(signer)}
	return
}

// Listen starts a SSH server listens on given port.
func Listen(port int) {
	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			pkey, err := models.SearchPublicKeyByContent(strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key))))
			if err != nil {
				log.Error(3, "SearchPublicKeyByContent: %v", err)
				return nil, err
			}
			return &ssh.Permissions{Extensions: map[string]string{"key-id": com.ToStr(pkey.ID)}}, nil
		},
	}

	keyPath := filepath.Join(setting.AppDataPath, "ssh/gogs.rsa")
	if !com.IsExist(keyPath) {
		os.MkdirAll(filepath.Dir(keyPath), os.ModePerm)
		_, stderr, err := com.ExecCmd("ssh-keygen", "-f", keyPath, "-t", "rsa", "-N", "")
		if err != nil {
			panic(fmt.Sprintf("Fail to generate private key: %v - %s", err, stderr))
		}
		log.Trace("New private key is generateed: %s", keyPath)
	}

	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		panic("Fail to load private key")
	}
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		panic("Fail to parse private key")
	}
	config.AddHostKey(private)

	go listen(config, port)
}

func setSshValues(c *Config) error {
	if c.Comm.SSHTimeout == 0 {
		c.Comm.SSHTimeout = 20 * time.Minute
	}

	if c.Comm.SSHPrivateKey != "" {
		privateKeyBytes, err := ioutil.ReadFile(c.Comm.SSHPrivateKey)
		if err != nil {
			panic(err)
		}
		signer, err := ssh.ParsePrivateKey(privateKeyBytes)
		if err != nil {
			panic(err)
		}

		publicKey := signer.PublicKey()
		c.sshAuthorizedKey = fmt.Sprintf("%s %s packer Azure Deployment%s",
			publicKey.Type(),
			base64.StdEncoding.EncodeToString(publicKey.Marshal()),
			time.Now().Format(time.RFC3339))
		c.sshPrivateKey = string(privateKeyBytes)

	} else {
		sshKeyPair, err := NewOpenSshKeyPair()
		if err != nil {
			return err
		}

		c.sshAuthorizedKey = sshKeyPair.AuthorizedKey()
		c.sshPrivateKey = sshKeyPair.PrivateKey()
	}

	return nil
}

func (r *Runner) ConnectAndRun(host, command string, options *ConnectionOptions) (string, error) {
	signer, err := ssh.ParsePrivateKey(options.PrivateKeyPEM)
	if err != nil {
		return "", err
	}

	config := &ssh.ClientConfig{
		User: options.Username,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
	}
	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", host, options.Port), config)
	if err != nil {
		return "", fmt.Errorf("failed to dial: %s", err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("failed to create session: ", err)
	}
	defer session.Close()

	var stdoutBytes bytes.Buffer
	session.Stdout = &stdoutBytes
	session.Stderr = os.Stderr
	if err := session.Run(command); err != nil {
		return "", fmt.Errorf("failed while running command: %s", err)
	}
	return stdoutBytes.String(), nil
}

// SSHFileSigner returns an ssh.Signer for a key file.
func SSHFileSigner(path string) (ssh.Signer, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	keyBytes, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	// We parse the private key on our own first so that we can
	// show a nicer error if the private key has a password.
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return nil, fmt.Errorf(
			"Failed to read key '%s': no key found", path)
	}
	if block.Headers["Proc-Type"] == "4,ENCRYPTED" {
		return nil, fmt.Errorf(
			"Failed to read key '%s': password protected keys are\n"+
				"not supported. Please decrypt the key prior to use.", path)
	}

	signer, err := ssh.ParsePrivateKey(keyBytes)
	if err != nil {
		return nil, fmt.Errorf("Error setting up SSH config: %s", err)
	}

	return signer, nil
}

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

	keypath = expandPath(keypath)

	// 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) {
		log.Printf("warning: ignoring encrypted key '%s'", keypath)
		return auths
	} else {
		signer, err := ssh.ParsePrivateKey(pemBytes)
		if err != nil {
			log.Print(err)
			return auths
		}
		return append(auths, ssh.PublicKeys(signer))
	}
}

// Run commands on the remote host
func (rs *Rsync) run(keys *drone.Key, host string) error {

	// join the host and port if necessary
	addr := net.JoinHostPort(host, strconv.Itoa(rs.Port))

	// trace command used for debugging in the build logs
	fmt.Printf("$ ssh %[email protected]%s -p %d\n", rs.User, addr, rs.Port)

	signer, err := ssh.ParsePrivateKey([]byte(keys.Private))
	if err != nil {
		return fmt.Errorf("Error parsing private key. %s.", err)
	}

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

	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return fmt.Errorf("Error dialing server. %s.", err)
	}

	session, err := client.NewSession()
	if err != nil {
		return fmt.Errorf("Error starting ssh session. %s.", err)
	}
	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return session.Run(strings.Join(rs.Commands, "\n"))
}

func (b *backend) pathKeysWrite(req *logical.Request, d *framework.FieldData) (*logical.Response, error) {
	keyName := d.Get("key_name").(string)
	if keyName == "" {
		return logical.ErrorResponse("Missing key_name"), nil
	}

	keyString := d.Get("key").(string)

	// Check if the key provided is infact a private key
	signer, err := ssh.ParsePrivateKey([]byte(keyString))
	if err != nil || signer == nil {
		return logical.ErrorResponse("Invalid key"), nil
	}

	if keyString == "" {
		return logical.ErrorResponse("Missing key"), nil
	}

	keyPath := fmt.Sprintf("keys/%s", keyName)

	// Store the key
	entry, err := logical.StorageEntryJSON(keyPath, map[string]interface{}{
		"key": keyString,
	})
	if err != nil {
		return nil, err
	}
	if err := req.Storage.Put(entry); err != nil {
		return nil, err
	}
	return nil, nil
}

func readPrivateKey(pk string) (ssh.AuthMethod, error) {
	key, _, err := pathorcontents.Read(pk)
	if err != nil {
		return nil, fmt.Errorf("Failed to read private key %q: %s", pk, err)
	}

	// We parse the private key on our own first so that we can
	// show a nicer error if the private key has a password.
	block, _ := pem.Decode([]byte(key))
	if block == nil {
		return nil, fmt.Errorf("Failed to read key %q: no key found", pk)
	}
	if block.Headers["Proc-Type"] == "4,ENCRYPTED" {
		return nil, fmt.Errorf(
			"Failed to read key %q: password protected keys are\n"+
				"not supported. Please decrypt the key prior to use.", pk)
	}

	signer, err := ssh.ParsePrivateKey([]byte(key))
	if err != nil {
		return nil, fmt.Errorf("Failed to parse key file %q: %s", pk, err)
	}

	return ssh.PublicKeys(signer), nil
}

func init() {
	signer, err := ssh.ParsePrivateKey([]byte(TestPrivateKey))
	if err != nil {
		panic(err)
	}
	testPublicKey = signer.PublicKey().Marshal()
}

// Configure creates a new SSH configuration object.
//
// Config sets a PublicKeyCallback handler that forwards public key auth
// requests to the route named "pubkeyAuth".
//
// This assumes certain details about our environment, like the location of the
// host keys. It also provides only key-based authentication.
// ConfigureServerSshConfig
//
// Returns:
//  An *ssh.ServerConfig
func Configure() (*ssh.ServerConfig, error) {
	cfg := &ssh.ServerConfig{
		PublicKeyCallback: func(m ssh.ConnMetadata, k ssh.PublicKey) (*ssh.Permissions, error) {
			return AuthKey(k)
		},
	}
	hostKeyTypes := []string{"rsa", "dsa", "ecdsa"}
	pathTpl := "/var/run/secrets/deis/builder/ssh/ssh-host-%s-key"
	for _, t := range hostKeyTypes {
		path := fmt.Sprintf(pathTpl, t)

		key, err := ioutil.ReadFile(path)
		if err != nil {
			log.Debug("Failed to read key %s (skipping): %s", path, err)
			return nil, err
		}
		hk, err := ssh.ParsePrivateKey(key)
		if err != nil {
			log.Debug("Failed to parse host key %s (skipping): %s", path, err)
			return nil, err
		}
		log.Debug("Parsed host key %s.", path)
		cfg.AddHostKey(hk)
	}
	return cfg, nil
}

func init() {
	var err error
	hostPrivateKeySigner, err = ssh.ParsePrivateKey(privKey)
	if err != nil {
		panic(err)
	}
}

func MakePrivateKeySignerFromBytes(buffer []byte) (ssh.Signer, error) {
	signer, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil, fmt.Errorf("error parsing SSH key %s: '%v'", buffer, err)
	}
	return signer, nil
}

// sshKeyAuth is a helper function to get the ssh key auth struct needed
func (obj *Remotes) sshKeyAuth() (ssh.AuthMethod, error) {
	if obj.sshPrivIdRsa == "" {
		return nil, fmt.Errorf("Empty path specified!")
	}
	p := ""
	// TODO: this doesn't match strings of the form: ~james/.ssh/id_rsa
	if strings.HasPrefix(obj.sshPrivIdRsa, "~/") {
		usr, err := user.Current()
		if err != nil {
			log.Printf("Remote: Can't find home directory automatically.")
			return nil, err
		}
		p = path.Join(usr.HomeDir, obj.sshPrivIdRsa[len("~/"):])
	}
	if p == "" {
		return nil, fmt.Errorf("Empty path specified!")
	}
	// A public key may be used to authenticate against the server by using
	// an unencrypted PEM-encoded private key file. If you have an encrypted
	// private key, the crypto/x509 package can be used to decrypt it.
	key, err := ioutil.ReadFile(p)
	if err != nil {
		return nil, err
	}

	// Create the Signer for this private key.
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, err
	}

	return ssh.PublicKeys(signer), nil
}