141 lines
3.6 KiB
Go
141 lines
3.6 KiB
Go
package grpcclient
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/mrhid6/vantage/agent/internal/grpc/pb"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/encoding"
|
|
"google.golang.org/grpc/keepalive"
|
|
)
|
|
|
|
func init() {
|
|
encoding.RegisterCodec(JSONCodec{})
|
|
}
|
|
|
|
type Client struct {
|
|
conn *grpc.ClientConn
|
|
client pb.VantageClient
|
|
}
|
|
|
|
func New(serverURL string, useTLS bool) (*Client, error) {
|
|
serverURL = strings.TrimPrefix(serverURL, "https://")
|
|
serverURL = strings.TrimPrefix(serverURL, "http://")
|
|
|
|
// Send a ping every 30s so proxies with a 60s idle timeout don't kill the
|
|
// long-lived CommandStream when no commands are flowing.
|
|
dialOpts := []grpc.DialOption{
|
|
grpc.WithKeepaliveParams(keepalive.ClientParameters{
|
|
Time: 30 * time.Second,
|
|
Timeout: 10 * time.Second,
|
|
PermitWithoutStream: false,
|
|
}),
|
|
}
|
|
|
|
if useTLS {
|
|
tlsCfg := &tls.Config{
|
|
InsecureSkipVerify: false,
|
|
}
|
|
creds := credentials.NewTLS(tlsCfg)
|
|
dialOpts = append(dialOpts, grpc.WithTransportCredentials(creds))
|
|
} else {
|
|
dialOpts = append(dialOpts, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
conn, err := grpc.DialContext(ctx, serverURL, dialOpts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Client{
|
|
conn: conn,
|
|
client: pb.NewVantageClient(conn),
|
|
}, nil
|
|
}
|
|
|
|
func (c *Client) Close() error {
|
|
return c.conn.Close()
|
|
}
|
|
|
|
func (c *Client) Register(serverID, preRegToken, hostname, ipAddress, osInfo string) (string, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
resp, err := c.client.Register(ctx, &pb.RegisterRequest{
|
|
ServerId: serverID,
|
|
PreRegToken: preRegToken,
|
|
Hostname: hostname,
|
|
IpAddress: ipAddress,
|
|
OsInfo: osInfo,
|
|
})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return resp.AgentToken, nil
|
|
}
|
|
|
|
func (c *Client) SyncKeys(serverID, agentToken, version string) ([]string, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
resp, err := c.client.SyncKeys(ctx, &pb.SyncRequest{
|
|
ServerId: serverID,
|
|
AgentToken: agentToken,
|
|
AgentVersion: version,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return resp.PublicKeys, nil
|
|
}
|
|
|
|
func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, privateKey, label string) (string, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
resp, err := c.client.UploadGeneratedKey(ctx, &pb.UploadKeyRequest{
|
|
ServerId: serverID,
|
|
AgentToken: agentToken,
|
|
PublicKey: publicKey,
|
|
PrivateKey: privateKey,
|
|
Label: label,
|
|
})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return resp.KeyId, nil
|
|
}
|
|
|
|
func (c *Client) ReportUpdates(serverID, agentToken string, updates []pb.PackageUpdate) error {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
_, err := c.client.ReportUpdates(ctx, &pb.ReportUpdatesRequest{
|
|
ServerId: serverID,
|
|
AgentToken: agentToken,
|
|
Updates: updates,
|
|
})
|
|
return err
|
|
}
|
|
|
|
func (c *Client) ReportInventory(report *pb.InventoryReport) error {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
_, err := c.client.ReportInventory(ctx, report)
|
|
return err
|
|
}
|
|
|
|
// CommandStream opens a long-lived bidirectional stream for server-pushed commands.
|
|
// The caller controls the stream lifetime via ctx.
|
|
func (c *Client) CommandStream(ctx context.Context) (pb.Vantage_CommandStreamClient, error) {
|
|
return c.client.CommandStream(ctx)
|
|
}
|