// see https://docs.velociraptor.app/docs/server_automation/server_api/
// for how to create the api config file.

// see https://github.com/Velocidex/pyvelociraptor/tree/master for python bindings

package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"flag"
	"io"
	"log"
	"os"

	"github.com/Velocidex/yaml/v2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
)

func main() {
	apiConfigPath := flag.String("api_config", "", "path to api config")
	flag.Parse()
	if *apiConfigPath == "" {
		flag.Usage()
		os.Exit(1)
	}

	y, err := os.ReadFile(*apiConfigPath)
	if err != nil {
		panic(err)
	}

	config := &config_proto.ApiClientConfig{}
	err = yaml.Unmarshal(y, config)
	if err != nil {
		panic(err)
	}

	cert, err := tls.X509KeyPair([]byte(config.ClientCert), []byte(config.ClientPrivateKey))
	if err != nil {
		log.Panic(err)
	}

	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM([]byte(config.CaCertificate))

	t := credentials.NewTLS(&tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs:      caCertPool,
		ServerName:   "VelociraptorServer",
	})

	conn, err := grpc.NewClient(config.ApiConnectionString, grpc.WithTransportCredentials(t))
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()

	request := &actions_proto.VQLCollectorArgs{
		Query: []*actions_proto.VQLRequest{
			{
				Name: "Query",
				VQL:  "SELECT * FROM info()",
			},
		},
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := api_proto.NewAPIClient(conn)
	queryClient, err := client.Query(ctx, request)
	if err != nil {
		log.Panic(err)
	}

	for {
		resp, err := queryClient.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Panic(err)
		}
		if resp.Response != "" {
			log.Print("Response: ", resp.Response)
		} else if resp.Log != "" {
			log.Print(resp.Log)
		}
	}

}