package ntlm

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/oiweiwei/go-msrpc/ssp/credential"
	"github.com/oiweiwei/go-msrpc/ssp/crypto"
	"github.com/oiweiwei/go-msrpc/ssp/ntlm/internal"
	"github.com/oiweiwei/go-msrpc/text/encoding/utf16le"
)

// The LMv1Respons structure defines the NTLM v1 authentication
// LmChallengeResponse in the AuthenticateMessage. This response
// is used only when NTLMv1 authentication is configured.
type LMv1Response struct {
	// A 24-byte array of unsigned char that contains the client's
	// LmChallengeResponse
	Response []byte
}

func (m *LMv1Response) Marshal(ctx context.Context) ([]byte, error) {

	e := internal.NewCodec(ctx, nil)

	// response.
	e.WriteBytes(ctx, m.Response, 24)

	return e.Bytes(ctx)
}

func (m *LMv1Response) Unmarshal(ctx context.Context, b []byte) error {

	e := internal.NewCodec(ctx, b)

	// response.
	e.ReadBytes(ctx, &m.Response, 24)

	return e.ReadAll(ctx)
}

// The NTLMv1 Response.
type NTLMv1Response = LMv1Response

type V1 struct {
	*Config
	*SecurityParameters
}

func (v1 *V1) WithConfig(ctx context.Context, config *Config) NTLMVersion {
	v1.Config = config.clone()
	return v1
}

var (
	// interface guard.
	_ NTLMVersion = (*V1)(nil)
)

var (
	// The DES magic nonce for NTLMv1 LMOWF.
	magicLMv1Nonce = []byte{'K', 'G', 'S', '!', '@', '#', '$', '%'}
)

// LMOWF function is a NT LAN Manager (LM) one-way function used to create
// a hash based on the user's password.
func (v1 *V1) LMOWF(ctx context.Context, cred Credential) ([]byte, error) {

	if cred == nil {
		return nil, nil
	}

	if _, ok := cred.(credential.Password); !ok {
		return nil, nil
	}

	secret := make([]byte, 14)
	copy(secret, []byte(strings.ToUpper(cred.(credential.Password).Password())))

	buf := bytes.NewBuffer(nil)

	buf.Write(crypto.DES(secret[:7], magicLMv1Nonce))
	buf.Write(crypto.DES(secret[7:], magicLMv1Nonce))

	return buf.Bytes(), nil
}

// NTOWF function is a NT LAN Manager (NT) one-way function used to create
// a hash based on user's password.
func (v1 *V1) NTOWF(ctx context.Context, cred Credential) ([]byte, error) {

	if cred == nil {
		return nil, nil
	}

	if key, ok := GetCredentialNTHash(cred); ok {
		return key, nil
	}

	if _, ok := cred.(credential.Password); !ok {
		return nil, fmt.Errorf("v1: ntowf: encode password: invalid credential type %T", cred)
	}

	b, err := utf16le.Encode(cred.(credential.Password).Password())
	if err != nil {
		return nil, fmt.Errorf("v1: ntowf: encode password: %w", err)
	}

	b, err = crypto.MD4(b)
	if err != nil {
		return nil, fmt.Errorf("v1: ntowf: derive key: %w", err)
	}

	return b, nil
}

// ChallengeResponse function computes the response for the server challenge message
// based on client challenge `nonce` parameter and set of negotiated flags.
func (v1 *V1) ChallengeResponse(ctx context.Context, cred Credential,
	c *ChallengeMessage, nonce []byte) (*ChallengeResponse, error) {

	if cred == nil || (cred.UserName() == "" && IsCredentialEmpty(cred)) {
		// anonymous case.
		return &ChallengeResponse{
			LM:             []byte{0},
			IsAnonymous:    true,
			SessionBaseKey: make([]byte, 16),
		}, nil
	}

	var (
		resp = new(ChallengeResponse)
		err  error
	)

	if resp.KeyNT, err = v1.NTOWF(ctx, cred); err != nil {
		return nil, err
	}

	if resp.KeyLM, err = v1.LMOWF(ctx, cred); err != nil {
		return nil, err
	}

	respNT, respLM := new(NTLMv1Response), new(LMv1Response)

	if c.Negotiate.IsSet(NegotiateExtendedSessionSecurity) {
		if respNT.Response, err = crypto.MD5(c.ServerChallenge, nonce); err != nil {
			return nil, fmt.Errorf("v1: compute response: %w", err)
		}
		respNT.Response = crypto.DESL(resp.KeyNT, respNT.Response[:8])
		respLM.Response = nonce
	} else {
		respNT.Response = crypto.DESL(resp.KeyNT, c.ServerChallenge)
		if v1.NoLMResponse {
			respLM.Response = respNT.Response
		} else {
			respLM.Response = crypto.DESL(resp.KeyLM, c.ServerChallenge)
		}
	}

	if resp.NT, err = respNT.Marshal(ctx); err != nil {
		return nil, fmt.Errorf("v1: marshal nt response: %v", err)
	}

	if resp.LM, err = respLM.Marshal(ctx); err != nil {
		return nil, fmt.Errorf("v1: marshal lm response: %v", err)
	}

	if resp.SessionBaseKey, err = crypto.MD4(resp.KeyNT); err != nil {
		return nil, fmt.Errorf("v1: compute session base key: %v", err)
	}

	return resp, nil
}

// KeyExchangeKey function returns the key used to protect the session key that is
// generated by the client.
func (v1 *V1) KeyExchangeKey(ctx context.Context, c *ChallengeMessage, chal *ChallengeResponse) ([]byte, error) {

	if chal.IsAnonymous {
		return chal.SessionBaseKey, nil
	}

	if len(chal.LM) < 8 {
		return nil, errors.New("v1: key exchange key: challenge response lm is empty")
	}

	if c.Negotiate.IsSet(NegotiateExtendedSessionSecurity) {
		return crypto.HMACMD5(chal.SessionBaseKey, c.ServerChallenge, chal.LM[:8])
	}

	if len(chal.KeyLM) < 8 {
		return nil, errors.New("v1: key exchange key: challenge response key lm is empty")
	}

	if c.Negotiate.IsSet(NegotiateLMKey) {

		buf := bytes.NewBuffer(nil)

		if _, err := buf.Write(crypto.DES(chal.KeyLM[:7], chal.LM[:8])); err != nil {
			return nil, fmt.Errorf("v1: key exchange key: lm key: %v", err)
		}

		magic := []byte{chal.KeyLM[7], 0xBD, 0xBD, 0xBD, 0xBD, 0xBD, 0xBD}

		if _, err := buf.Write(crypto.DES(magic, chal.LM[:8])); err != nil {
			return nil, fmt.Errorf("v1: key exchange key: lm key: %v", err)
		}
		return buf.Bytes(), nil
	}

	if c.Negotiate.IsSet(RequestNonNTSessionKey) {
		return append(chal.KeyLM[:8], bytes.Repeat([]byte{0}, 8)...), nil
	}

	return chal.SessionBaseKey, nil
}