83 lines
1.7 KiB
Go
83 lines
1.7 KiB
Go
package client
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"sync"
|
|
|
|
"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/agent/internal/commander"
|
|
"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/proto/proto"
|
|
"golang.org/x/sync/errgroup"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials"
|
|
"google.golang.org/grpc/metadata"
|
|
)
|
|
|
|
type CommanderClient struct {
|
|
cmder *commander.Commander
|
|
wg *sync.WaitGroup
|
|
id string
|
|
}
|
|
|
|
func New(
|
|
cmder *commander.Commander,
|
|
wg *sync.WaitGroup,
|
|
id string,
|
|
) CommanderClient {
|
|
return CommanderClient{cmder, wg, id}
|
|
}
|
|
|
|
func (self *CommanderClient) HandleCommands(ctx context.Context, srvAddr string, tc credentials.TransportCredentials) error {
|
|
cli, err := grpc.NewClient(srvAddr, grpc.WithTransportCredentials(tc))
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to connect to gRPC: %w", err)
|
|
}
|
|
ccli := proto.NewCommanderClient(cli)
|
|
bidi, err := ccli.Stream(metadata.NewOutgoingContext(ctx, metadata.MD{"agentid": []string{self.id}}))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
wg := new(errgroup.Group)
|
|
wg.Go(self.recv(bidi))
|
|
// wg.Go(self.send(bidi))
|
|
err = wg.Wait()
|
|
self.wg.Wait()
|
|
return err
|
|
}
|
|
|
|
func (self *CommanderClient) recv(bidi grpc.BidiStreamingClient[proto.FinishedCommand, proto.Command]) func() error {
|
|
return func() error {
|
|
for {
|
|
msg, err := bidi.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
self.wg.Go(func() {
|
|
func() error {
|
|
fc, err := self.cmder.Execute(msg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return bidi.Send(fc)
|
|
}()
|
|
if err != nil {
|
|
log.Println(err)
|
|
}
|
|
})
|
|
|
|
}
|
|
}
|
|
}
|
|
|
|
// func (self *God) send(bidi grpc.BidiStreamingClient[proto.FinishedCommand, proto.Command]) func() error {
|
|
// return func() error {
|
|
// return nil
|
|
// }
|
|
// }
|