// Copyright (c) 2017 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package gnmi

import (
	"bufio"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"slices"
	"strings"

	"github.com/aristanetworks/goarista/netns"
	pb "github.com/openconfig/gnmi/proto/gnmi"
	"github.com/openconfig/gnmi/proto/gnmi_ext"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/encoding/gzip"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"
)

const (
	defaultPort = "6030"
	// HostnameArg is the value to be replaced by the actual hostname
	HostnameArg = "HOSTNAME"
)

type tlsVersionMap map[string]uint16

func (m tlsVersionMap) String() string {
	r := make([]string, 0, len(m))
	for k := range m {
		r = append(r, k)
	}
	slices.Sort(r)
	return strings.Join(r, ", ")
}

// TLSVersions is a map from TLS version strings to the tls version
// constants in the crypto/tls package
var TLSVersions = getTLSVersions()

// PublishFunc is the method to publish responses
type PublishFunc func(addr string, message proto.Message)

// ParseHostnames parses a comma-separated list of names and replaces HOSTNAME with the current
// hostname in it
func ParseHostnames(list string) ([]string, error) {
	items := strings.Split(list, ",")
	hostname, err := os.Hostname()
	if err != nil {
		return nil, err
	}
	names := make([]string, len(items))
	for i, name := range items {
		if name == HostnameArg {
			name = hostname
		}
		names[i] = name
	}
	return names, nil
}

// Config is the gnmi.Client config
type Config struct {
	Addr string

	// File path to load data or raw cert data. Alternatively, raw data can be provided below.
	CAFile   string
	CertFile string
	KeyFile  string

	// Raw certificate data. If respective file is provided above, that is used instead.
	CAData   []byte
	CertData []byte
	KeyData  []byte

	Password      string
	Username      string
	TLS           bool
	TLSMinVersion string
	TLSMaxVersion string
	Compression   string
	BDP           bool
	DialOptions   []grpc.DialOption
	Token         string
	GRPCMetadata  map[string]string
	Proxy         string
}

// SubscribeOptions is the gNMI subscription request options
type SubscribeOptions struct {
	UpdatesOnly       bool
	Prefix            string
	Mode              string
	StreamMode        string
	SampleInterval    uint64
	SuppressRedundant bool
	HeartbeatInterval uint64
	Paths             [][]string
	Origin            string
	Target            string
	Extensions        []*gnmi_ext.Extension
}

// ParseFlags reads arguments from stdin and returns a populated Config object and a list of
// paths to subscribe to
func ParseFlags() (*Config, []string) {
	// flags
	var (
		addrsFlag = flag.String("addrs", "localhost:6030",
			"Comma-separated list of addresses of OpenConfig gRPC servers. The address 'HOSTNAME' "+
				"is replaced by the current hostname.")

		caFileFlag = flag.String("cafile", "",
			"Path to server TLS certificate file")

		certFileFlag = flag.String("certfile", "",
			"Path to client TLS certificate file")

		keyFileFlag = flag.String("keyfile", "",
			"Path to client TLS private key file")

		passwordFlag = flag.String("password", "",
			"Password to authenticate with")

		usernameFlag = flag.String("username", "",
			"Username to authenticate with")

		tlsFlag = flag.Bool("tls", false,
			"Enable TLS")
		tlsMinVersion = flag.String("tls-min-version", "",
			fmt.Sprintf("Set minimum TLS version for connection (%s)", TLSVersions))
		tlsMaxVersion = flag.String("tls-max-version", "",
			fmt.Sprintf("Set minimum TLS version for connection (%s)", TLSVersions))

		compressionFlag = flag.String("compression", "",
			"Type of compression to use")

		subscribeFlag = flag.String("subscribe", "",
			"Comma-separated list of paths to subscribe to upon connecting to the server")

		token = flag.String("token", "",
			"Authentication token")
	)
	flag.Parse()
	cfg := &Config{
		Addr:          *addrsFlag,
		CAFile:        *caFileFlag,
		CertFile:      *certFileFlag,
		KeyFile:       *keyFileFlag,
		Password:      *passwordFlag,
		Username:      *usernameFlag,
		TLS:           *tlsFlag,
		TLSMinVersion: *tlsMinVersion,
		TLSMaxVersion: *tlsMaxVersion,
		Compression:   *compressionFlag,
		Token:         *token,
	}
	subscriptions := strings.Split(*subscribeFlag, ",")
	return cfg, subscriptions

}

// accessTokenCred implements credentials.PerRPCCredentials, the gRPC
// interface for credentials that need to attach security information
// to every RPC.
type accessTokenCred struct {
	bearerToken string
}

// newAccessTokenCredential constructs a new per-RPC credential from a token.
func newAccessTokenCredential(token string) credentials.PerRPCCredentials {
	bearerFmt := "Bearer %s"
	return &accessTokenCred{bearerToken: fmt.Sprintf(bearerFmt, token)}
}

func (a *accessTokenCred) GetRequestMetadata(ctx context.Context,
	uri ...string) (map[string]string, error) {
	authHeader := "Authorization"
	return map[string]string{
		authHeader: a.bearerToken,
	}, nil
}

func (a *accessTokenCred) RequireTransportSecurity() bool { return true }

// DialContextConn connects to a gnmi service and return a client connection
func DialContextConn(ctx context.Context, cfg *Config) (*grpc.ClientConn, error) {
	opts := append([]grpc.DialOption(nil), cfg.DialOptions...)

	if !cfg.BDP {
		// By default, the client and server will dynamically adjust the connection's
		// window size using the Bandwidth Delay Product (BDP).
		// See: https://grpc.io/blog/grpc-go-perf-improvements/
		// The default values for InitialWindowSize and InitialConnWindowSize are 65535.
		// If values less than 65535 are used, then BDP and dynamic windows are enabled.
		// Here, we disable the BDP and dynamic windows by setting these values >= 65535.
		// We set these values to (1 << 20) * 16 as this is the largest window size that
		// the BDP estimator could ever use.
		// See: https://github.com/grpc/grpc-go/blob/master/internal/transport/bdp_estimator.go
		const maxWindowSize int32 = (1 << 20) * 16
		opts = append(opts,
			grpc.WithInitialWindowSize(maxWindowSize),
			grpc.WithInitialConnWindowSize(maxWindowSize),
		)
	}

	switch cfg.Compression {
	case "":
	case "gzip":
		opts = append(opts, grpc.WithDefaultCallOptions(grpc.UseCompressor(gzip.Name)))
	default:
		return nil, fmt.Errorf("unsupported compression option: %q", cfg.Compression)
	}

	var err error
	caData := cfg.CAData
	certData := cfg.CertData
	keyData := cfg.KeyData
	if cfg.CAFile != "" {
		if caData, err = os.ReadFile(cfg.CAFile); err != nil {
			return nil, err
		}
	}
	if cfg.CertFile != "" {
		if certData, err = os.ReadFile(cfg.CertFile); err != nil {
			return nil, err
		}
	}
	if cfg.KeyFile != "" {
		if keyData, err = os.ReadFile(cfg.KeyFile); err != nil {
			return nil, err
		}
	}

	if cfg.TLS || len(caData) > 0 || len(certData) > 0 || cfg.Token != "" {
		tlsConfig := &tls.Config{}
		if len(caData) > 0 {
			cp := x509.NewCertPool()
			if !cp.AppendCertsFromPEM(caData) {
				return nil, fmt.Errorf("credentials: failed to append certificates")
			}
			tlsConfig.RootCAs = cp
		} else {
			tlsConfig.InsecureSkipVerify = true
		}
		if len(certData) > 0 {
			if len(keyData) == 0 {
				return nil, fmt.Errorf("no key provided for client certificate")
			}
			cert, err := tls.X509KeyPair(certData, keyData)
			if err != nil {
				return nil, err
			}
			tlsConfig.Certificates = []tls.Certificate{cert}
		}
		if cfg.Token != "" {
			opts = append(opts,
				grpc.WithPerRPCCredentials(newAccessTokenCredential(cfg.Token)))
		}
		if cfg.TLSMaxVersion != "" {
			var ok bool
			tlsConfig.MaxVersion, ok = TLSVersions[cfg.TLSMaxVersion]
			if !ok {
				return nil, fmt.Errorf("unrecognised TLS max version."+
					" Supported TLS versions are %s", TLSVersions)
			}
		}
		if cfg.TLSMinVersion != "" {
			var ok bool
			tlsConfig.MinVersion, ok = TLSVersions[cfg.TLSMinVersion]
			if !ok {
				return nil, fmt.Errorf("unrecognised TLS min version."+
					" Supported TLS versions are %s", TLSVersions)
			}
		}
		if cfg.TLSMinVersion != "" && cfg.TLSMaxVersion != "" &&
			tlsConfig.MinVersion > tlsConfig.MaxVersion {
			return nil, fmt.Errorf(
				"TLS min version was greater than TLS max version")
		}

		opts = append(opts, grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)))
	} else {
		opts = append(opts, grpc.WithInsecure())
	}

	network, nsName, addr, err := ParseAddress(cfg.Addr)
	if err != nil {
		return nil, err
	}

	opts = append(opts,
		grpc.WithContextDialer(Dialer(&net.Dialer{}, network, nsName, WithProxy(cfg.Proxy))),

		// Allows received protobuf messages to be larger than 4MB
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(math.MaxInt32)),
	)

	return grpc.DialContext(ctx, addr, opts...)
}

// ParseAddress parses the network, namespace and address for TCP or
// Unix domain socket addresses. Examples:
//
//	"mgmt/127.0.0.1:6030" returns "tcp", "ns-mgmt", "127.0.0.1:6030"
//	"unix:///var/run/gnmiServer.sock" returns "unix", "", "/var/run/gnmiServer.sock".
func ParseAddress(addrIn string) (string, string, string, error) {
	// Parse "unix://" and "unix:" addresses as in
	// google.golang.org/grpc/internal/transport/http2_client.go
	if addr, ok := strings.CutPrefix(addrIn, "unix:"); ok {
		return "unix", "", strings.TrimPrefix(addr, "//"), nil
	}

	var network, nsName, addr string
	split := strings.Split(addrIn, "://")
	if l := len(split); l == 2 {
		network = split[0]
		addr = split[1]
	} else {
		network = "tcp"
		addr = split[0]
	}

	if !strings.ContainsRune(addr, ':') {
		addr += ":" + defaultPort
	}

	var err error
	nsName, addr, err = netns.ParseAddress(addr)
	if err != nil {
		return "", "", "", err
	}
	return network, nsName, addr, nil
}

// DialOption configures how we set up the connection.
type DialOption func(m *dialOptions)

type dialOptions struct {
	proxy string // example: http://[user:password@]ip:port
}

// WithProxy requires the connection to go through the specified proxy
func WithProxy(proxy string) DialOption {
	return func(d *dialOptions) {
		d.proxy = proxy
	}
}

// Dialer returns a dialer function with the given net.Dialer, network and namespace.
func Dialer(d *net.Dialer, network, nsName string, opts ...DialOption) func(
	context.Context, string) (net.Conn, error) {

	dialOpts := &dialOptions{}
	for _, opt := range opts {
		opt(dialOpts)
	}

	return func(ctx context.Context, addr string) (net.Conn, error) {
		var conn net.Conn
		err := netns.Do(nsName, func() error {
			var proxyUser, proxyPass string
			target := addr
			if dialOpts.proxy != "" {
				proxy, err := url.Parse(dialOpts.proxy)
				if err != nil {
					return fmt.Errorf("Can't parse proxy URL: %v", err)
				}
				if u := proxy.User; u != nil {
					proxyUser = u.Username()
					proxyPass, _ = u.Password()
				}
				addr = proxy.Host
			}

			c, err := d.DialContext(ctx, network, addr)
			if err != nil {
				return err
			}

			if dialOpts.proxy != "" {
				c, err = proxyConnect(ctx, c, target, proxyUser, proxyPass)
				if err != nil {
					return fmt.Errorf("can't connect through proxy: %w", err)
				}
			}
			conn = c
			return nil
		})
		return conn, err
	}
}

// DialContext connects to a gnmi service and returns a client
func DialContext(ctx context.Context, cfg *Config) (pb.GNMIClient, error) {
	grpcconn, err := DialContextConn(ctx, cfg)
	if err != nil {
		return nil, fmt.Errorf("failed to dial: %s", err)
	}
	return pb.NewGNMIClient(grpcconn), nil
}

// Dial connects to a gnmi service and returns a client
func Dial(cfg *Config) (pb.GNMIClient, error) {
	return DialContext(context.Background(), cfg)
}

// NewContext returns a new context with username and password
// metadata if they are set in cfg, as well as any other metadata
// provided.
func NewContext(ctx context.Context, cfg *Config) context.Context {
	md := map[string]string{}
	for k, v := range cfg.GRPCMetadata {
		md[k] = v
	}
	if cfg.Username != "" {
		md["username"] = cfg.Username
		md["password"] = cfg.Password
	}
	if len(md) > 0 {
		ctx = metadata.NewOutgoingContext(ctx, metadata.New(md))
	}
	return ctx
}

// NewGetRequest returns a GetRequest for the given paths
func NewGetRequest(paths [][]string, origin string) (*pb.GetRequest, error) {
	req := &pb.GetRequest{
		Path: make([]*pb.Path, len(paths)),
	}
	for i, p := range paths {
		gnmiPath, err := ParseGNMIElements(p)
		if err != nil {
			return nil, err
		}
		req.Path[i] = gnmiPath
		req.Path[i].Origin = origin
	}
	return req, nil
}

// NewSubscribeRequest returns a SubscribeRequest for the given paths
func NewSubscribeRequest(subscribeOptions *SubscribeOptions) (*pb.SubscribeRequest, error) {
	var mode pb.SubscriptionList_Mode
	switch subscribeOptions.Mode {
	case "once":
		mode = pb.SubscriptionList_ONCE
	case "poll":
		mode = pb.SubscriptionList_POLL
	case "":
		fallthrough
	case "stream":
		mode = pb.SubscriptionList_STREAM
	default:
		return nil, fmt.Errorf("subscribe mode (%s) invalid", subscribeOptions.Mode)
	}

	var streamMode pb.SubscriptionMode
	switch subscribeOptions.StreamMode {
	case "on_change":
		streamMode = pb.SubscriptionMode_ON_CHANGE
	case "sample":
		streamMode = pb.SubscriptionMode_SAMPLE
	case "":
		fallthrough
	case "target_defined":
		streamMode = pb.SubscriptionMode_TARGET_DEFINED
	default:
		return nil, fmt.Errorf("subscribe stream mode (%s) invalid", subscribeOptions.StreamMode)
	}

	prefixPath, err := ParseGNMIElements(SplitPath(subscribeOptions.Prefix))
	if err != nil {
		return nil, err
	}
	subList := &pb.SubscriptionList{
		Subscription: make([]*pb.Subscription, len(subscribeOptions.Paths)),
		Mode:         mode,
		UpdatesOnly:  subscribeOptions.UpdatesOnly,
		Prefix:       prefixPath,
	}
	if subscribeOptions.Target != "" {
		if subList.Prefix == nil {
			subList.Prefix = &pb.Path{}
		}
		subList.Prefix.Target = subscribeOptions.Target
	}
	for i, p := range subscribeOptions.Paths {
		gnmiPath, err := ParseGNMIElements(p)
		if err != nil {
			return nil, err
		}
		gnmiPath.Origin = subscribeOptions.Origin
		subList.Subscription[i] = &pb.Subscription{
			Path:              gnmiPath,
			Mode:              streamMode,
			SampleInterval:    subscribeOptions.SampleInterval,
			SuppressRedundant: subscribeOptions.SuppressRedundant,
			HeartbeatInterval: subscribeOptions.HeartbeatInterval,
		}
	}
	return &pb.SubscribeRequest{
		Extension: subscribeOptions.Extensions,
		Request: &pb.SubscribeRequest_Subscribe{
			Subscribe: subList,
		},
	}, nil
}

// HistorySnapshotExtension returns an Extension_History for the given
// time.
func HistorySnapshotExtension(t int64) *gnmi_ext.Extension_History {
	return &gnmi_ext.Extension_History{
		History: &gnmi_ext.History{
			Request: &gnmi_ext.History_SnapshotTime{
				SnapshotTime: t,
			},
		},
	}
}

// HistoryRangeExtension returns an Extension_History for the the
// specified start and end times.
func HistoryRangeExtension(s, e int64) *gnmi_ext.Extension_History {
	return &gnmi_ext.Extension_History{
		History: &gnmi_ext.History{
			Request: &gnmi_ext.History_Range{
				Range: &gnmi_ext.TimeRange{
					Start: s,
					End:   e,
				},
			},
		},
	}
}

// getTLSVersions generates a map of TLS version name to tls version, based on the versions
// available in the crypto/tls package
func getTLSVersions(testHook ...func(uint16, *regexp.Regexp)) tlsVersionMap {
	cipherSuites := tls.CipherSuites()
	allSupportedVersions := make(map[uint16]struct{})

	for _, cipherSuite := range cipherSuites {
		for _, version := range cipherSuite.SupportedVersions {
			allSupportedVersions[version] = struct{}{}
		}
	}

	// match TLS versions in dot format like X.Y or X.Y.Z etc (right now everything is X.Y)
	re := regexp.MustCompile(`[\d.]+`)

	nameToVersion := make(map[string]uint16, len(allSupportedVersions))
	for version := range allSupportedVersions {
		// tls.VersionName(version) will be something like "TLS 1.3"
		name := re.FindString(tls.VersionName(version))
		// check if the regex either failed to match, or if it is not specific enough
		// (matching something which was already found)
		if _, ok := nameToVersion[name]; ok || name == "" {
			// if we ever fail to match a regex we shouldn't do anything in production
			// but let's make a test fail so we can investigate and update the regex
			for _, f := range testHook {
				f(version, re)
			}
			continue
		}

		nameToVersion[name] = version

	}
	return nameToVersion
}

type proxyConn struct {
	net.Conn
	rdr io.Reader
}

func (c *proxyConn) Read(b []byte) (int, error) {
	return c.rdr.Read(b)
}

// proxyConnect connects through a proxy that supports HTTP tunnelling as described in RFC7231
// secion 4.3.6.
func proxyConnect(ctx context.Context, conn net.Conn, target, user, pass string) (net.Conn, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodConnect, "grpc://"+target, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create send request: %v", err)
	}
	if user != "" || pass != "" {
		auth := "Basic " + base64.StdEncoding.EncodeToString([]byte(user+":"+pass))
		req.Header.Add("Proxy-Authorization", auth)
	}

	if err := req.Write(conn); err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to send to proxy: %v", err)
	}

	rdr := bufio.NewReader(conn)
	resp, err := http.ReadResponse(rdr, req)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to read from proxy: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		conn.Close()
		return nil, fmt.Errorf("proxy returned status: %v ", resp.Status)
	}

	return &proxyConn{Conn: conn, rdr: rdr}, nil
}
