package ssh

import (
	"encoding/asn1"
	"fmt"
	"testing"
)

func TestParseGSSAPIPayload(t *testing.T) {
	payload := []byte{0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x0b, 0x06, 0x09,
		0x2a, 0x86, 0x48, 0x86, 0xf7, 0x12, 0x01, 0x02, 0x02}
	res, err := parseGSSAPIPayload(payload)
	if err != nil {
		t.Fatal(err)
	}
	if ok := res.OIDS[0].Equal(krb5Mesh); !ok {
		t.Fatalf("got %v, want %v", res, krb5Mesh)
	}
}

func TestParseGSSAPIPayloadMultipleMechs(t *testing.T) {
	// A payload offering two mechanisms: Kerberos V5 (1.2.840.113554.1.2.2)
	// followed by SPNEGO (1.3.6.1.5.5.2). The parser must track the
	// remaining payload across loop iterations, not the trailing bytes
	// returned by asn1.Unmarshal.
	payload := []byte{
		0x00, 0x00, 0x00, 0x02, // two mechanisms
		0x00, 0x00, 0x00, 0x0b, // length of the first OID
		0x06, 0x09, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x12, 0x01, 0x02, 0x02,
		0x00, 0x00, 0x00, 0x08, // length of the second OID
		0x06, 0x06, 0x2b, 0x06, 0x01, 0x05, 0x05, 0x02,
	}
	res, err := parseGSSAPIPayload(payload)
	if err != nil {
		t.Fatal(err)
	}
	if len(res.OIDS) != 2 {
		t.Fatalf("got %d OIDs, want 2", len(res.OIDS))
	}
	if !res.OIDS[0].Equal(krb5Mesh) {
		t.Errorf("OIDS[0]: got %v, want %v", res.OIDS[0], krb5Mesh)
	}
	spnego := asn1.ObjectIdentifier{1, 3, 6, 1, 5, 5, 2}
	if !res.OIDS[1].Equal(spnego) {
		t.Errorf("OIDS[1]: got %v, want %v", res.OIDS[1], spnego)
	}
}

func TestParseDubiousGSSAPIPayload(t *testing.T) {
	for _, tc := range []struct {
		name    string
		payload []byte
		wanterr bool
	}{
		{
			"num mechanisms is unrealistic",
			[]byte{0xFF, 0x00, 0x00, 0xFF,
				0x00, 0x00, 0x00, 0x0b, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x12, 0x01, 0x02, 0x02},
			true,
		},
		{
			"num mechanisms greater than payload",
			[]byte{0x00, 0x00, 0x00, 0x40, // 64, |rest| too small
				0x00, 0x00, 0x00, 0x0b, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x12, 0x01, 0x02, 0x02},
			true,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			_, err := parseGSSAPIPayload(tc.payload)
			if tc.wanterr && err == nil {
				t.Errorf("got nil, want error")
			}
			if !tc.wanterr && err != nil {
				t.Errorf("got %v, want nil", err)
			}
		})
	}
}

func TestBuildMIC(t *testing.T) {
	sessionID := []byte{134, 180, 134, 194, 62, 145, 171, 82, 119, 149, 254, 196, 125, 173, 177, 145, 187, 85, 53,
		183, 44, 150, 219, 129, 166, 195, 19, 33, 209, 246, 175, 121}
	username := "testuser"
	service := "ssh-connection"
	authMethod := "gssapi-with-mic"
	expected := []byte{0, 0, 0, 32, 134, 180, 134, 194, 62, 145, 171, 82, 119, 149, 254, 196, 125, 173, 177, 145, 187, 85, 53, 183, 44, 150, 219, 129, 166, 195, 19, 33, 209, 246, 175, 121, 50, 0, 0, 0, 8, 116, 101, 115, 116, 117, 115, 101, 114, 0, 0, 0, 14, 115, 115, 104, 45, 99, 111, 110, 110, 101, 99, 116, 105, 111, 110, 0, 0, 0, 15, 103, 115, 115, 97, 112, 105, 45, 119, 105, 116, 104, 45, 109, 105, 99}
	result := buildMIC(string(sessionID), username, service, authMethod)
	if string(result) != string(expected) {
		t.Fatalf("buildMic: got %v, want %v", result, expected)
	}
}

type exchange struct {
	outToken      string
	expectedToken string
}

type FakeClient struct {
	exchanges []*exchange
	round     int
	mic       []byte
	maxRound  int
}

func (f *FakeClient) InitSecContext(target string, token []byte, isGSSDelegCreds bool) (outputToken []byte, needContinue bool, err error) {
	if token == nil {
		if f.exchanges[f.round].expectedToken != "" {
			err = fmt.Errorf("got empty token, want %q", f.exchanges[f.round].expectedToken)
		} else {
			outputToken = []byte(f.exchanges[f.round].outToken)
		}
	} else {
		if string(token) != string(f.exchanges[f.round].expectedToken) {
			err = fmt.Errorf("got %q, want token %q", token, f.exchanges[f.round].expectedToken)
		} else {
			outputToken = []byte(f.exchanges[f.round].outToken)
		}
	}
	f.round++
	needContinue = f.round < f.maxRound
	return
}

func (f *FakeClient) GetMIC(micField []byte) ([]byte, error) {
	return f.mic, nil
}

func (f *FakeClient) DeleteSecContext() error {
	return nil
}

type FakeServer struct {
	exchanges   []*exchange
	round       int
	expectedMIC []byte
	srcName     string
	maxRound    int
}

func (f *FakeServer) AcceptSecContext(token []byte) (outputToken []byte, srcName string, needContinue bool, err error) {
	if token == nil {
		if f.exchanges[f.round].expectedToken != "" {
			err = fmt.Errorf("got empty token, want %q", f.exchanges[f.round].expectedToken)
		} else {
			outputToken = []byte(f.exchanges[f.round].outToken)
		}
	} else {
		if string(token) != string(f.exchanges[f.round].expectedToken) {
			err = fmt.Errorf("got %q, want token %q", token, f.exchanges[f.round].expectedToken)
		} else {
			outputToken = []byte(f.exchanges[f.round].outToken)
		}
	}
	f.round++
	needContinue = f.round < f.maxRound
	srcName = f.srcName
	return
}

func (f *FakeServer) VerifyMIC(micField []byte, micToken []byte) error {
	if string(micToken) != string(f.expectedMIC) {
		return fmt.Errorf("got MICToken %q, want %q", micToken, f.expectedMIC)
	}
	return nil
}

func (f *FakeServer) DeleteSecContext() error {
	return nil
}