domrichardson
111 lines
2.7 KiB
Go
111 lines
2.7 KiB
Go
package grpcclient
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/mrhid6/keymanager/agent/internal/grpc/pb"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/encoding"
|
|
)
|
|
|
|
func init() {
|
|
encoding.RegisterCodec(JSONCodec{})
|
|
}
|
|
|
|
type Client struct {
|
|
conn *grpc.ClientConn
|
|
client pb.KeyManagerClient
|
|
}
|
|
|
|
func New(serverURL string, useTLS bool) (*Client, error) {
|
|
serverURL = strings.TrimPrefix(serverURL, "https://")
|
|
serverURL = strings.TrimPrefix(serverURL, "http://")
|
|
|
|
var dialOpts []grpc.DialOption
|
|
|
|
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.NewKeyManagerClient(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 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,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return resp.PublicKeys, nil
|
|
}
|
|
|
|
func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, 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,
|
|
Label: label,
|
|
})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return resp.KeyId, nil
|
|
}
|
|
|
|
// 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.KeyManager_CommandStreamClient, error) {
|
|
return c.client.CommandStream(ctx)
|
|
}
|