package ssh

import (
	"bufio"
	"bytes"
	"strings"

	osm "git.loafle.net/overflow/service_matcher-go"
)

type SSHMatcher struct {
	osm.Matchers
}

func (m *SSHMatcher) Key(matchCtx *osm.MatchCtx) string {
	return "SSH"
}

func (m *SSHMatcher) Type(matchCtx *osm.MatchCtx) string {
	return "NETWORK"
}

func (m *SSHMatcher) Vendor(matchCtx *osm.MatchCtx) string {
	return "UNKNOWN"
}

func (m *SSHMatcher) Version(matchCtx *osm.MatchCtx) string {
	return "UNKNOWN"
}

func (m *SSHMatcher) OsType(matchCtx *osm.MatchCtx) string {
	return "UNKNOWN"
}

func (m *SSHMatcher) OsVersion(matchCtx *osm.MatchCtx) string {
	return "UNKNOWN"
}

func (m *SSHMatcher) Name(matchCtx *osm.MatchCtx) string {
	name := "SSH"
	if v, ok := matchCtx.GetAttribute("softwareversion"); ok {
		name = name + " (" + v + ")"
	}
	return name
}

func (m *SSHMatcher) IsPrePacket() bool {
	return true
}

func (m *SSHMatcher) HasResponse(matchCtx *osm.MatchCtx, index int) bool {
	return true
}

func (m *SSHMatcher) IsError(matchCtx *osm.MatchCtx, index int, packet *osm.Packet) bool {
	return false
}

func (m *SSHMatcher) Match(matchCtx *osm.MatchCtx, index int, packet *osm.Packet) error {

	if packet == nil || !packet.Valid() {
		return osm.NoPacketReceivedError()
	}

	// SSH-protoversion-softwareversion SP comments CR LF
	// e.g. ) SSH-2.0-OpenSSH_7.5p1 Ubuntu-10ubuntu0.1\n
	scanner := bufio.NewScanner(bytes.NewReader(packet.Buffer))
	for scanner.Scan() {
		exchange := scanner.Text()

		if !strings.HasPrefix(exchange, "SSH") {
			return osm.NotMatchedError()
		}

		temp := strings.Split(exchange, " ")
		versions := strings.Split(temp[0], "-")

		protoversion := versions[1]
		softwareversion := versions[2]

		if strings.HasPrefix(protoversion, "1") || strings.HasPrefix(protoversion, "2") {
			matchCtx.SetAttribute("protoversion", protoversion)
			matchCtx.SetAttribute("softwareversion", softwareversion)

			if len(temp) > 1 {
				matchCtx.SetAttribute("comments", temp[1])
			}
			return nil
		}
		break
	}

	return osm.NotMatchedError()
}

func NewMatcher() osm.Matcher {
	m := &SSHMatcher{}

	return m
}