package protocol

import (
	"bytes"
	"strings"
)

func init() {
	// We can't match on SSH-?.? here, because the client or server may send a banner prior
	// to sending the SSH handshake.
	Register(Both, "", detectSSH)
}

// The required prefix for the SSH protocol identification line.
const (
	ssh199Prefix = "SSH-1.99-"
	ssh20Prefix  = "SSH-2.0-"
)

var (
	commonPort = map[int]bool{
		22:   true,
		2200: true,
		2222: true,
	}
	commonImplementations = []string{
		"OpenSSH_",
		"PuTTY",
		"libssh",
		"dropbear",
		"Go",
		"paramiko",
		"Cyberduck",
	}
)

func detectSSH(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64) {
	// The data must be at least as long as the prefix itself.
	if len(data) < len(ssh20Prefix) {
		return nil, 0
	}

	if commonPort[srcPort] || commonPort[dstPort] {
		confidence += .1
	}

	// The protocol allows for pre-banner text, so we have to check all lines.
	for _, line := range bytes.Split(data, []byte{'\n'}) {
		line = bytes.TrimSuffix(line, []byte{'\r'})

		// Match the most common SSH 2.0 protocol.
		if bytes.HasPrefix(line, []byte(ssh20Prefix)) {
			implementation := string(line[len(ssh20Prefix):])
			for _, prefix := range commonImplementations {
				if strings.HasPrefix(implementation, prefix) {
					confidence += .2
					break
				}
			}
			return &Protocol{
				Name: ProtocolSSH,
				Version: Version{
					Major: 2,
					Minor: 0,
					Patch: -1,
					Extra: implementation,
				},
			}, confidence + 0.65
		}

		// Match the (far) less common SSH 1.99 protocol.
		if bytes.HasPrefix(line, []byte(ssh199Prefix)) {
			implementation := string(line[len(ssh20Prefix):])
			for _, prefix := range commonImplementations {
				if strings.HasPrefix(implementation, prefix) {
					confidence += .2
					break
				}
			}
			return &Protocol{
				Name: ProtocolSSH,
				Version: Version{
					Major: 1,
					Minor: 99,
					Patch: -1,
					Extra: implementation,
				},
			}, confidence + 0.65
		}
	}

	return nil, 0
}