82 lines
2.2 KiB
Go
82 lines
2.2 KiB
Go
|
package plugin
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"net"
|
||
|
|
|
||
|
|
"github.com/Neur0toxine/sshpoke/internal/config"
|
||
|
|
"github.com/Neur0toxine/sshpoke/internal/logger"
|
||
|
|
"github.com/Neur0toxine/sshpoke/internal/model"
|
||
|
|
"github.com/Neur0toxine/sshpoke/internal/server"
|
||
|
|
"github.com/Neur0toxine/sshpoke/internal/server/driver/plugin"
|
||
|
|
pb "github.com/Neur0toxine/sshpoke/pkg/plugin"
|
||
|
|
"google.golang.org/grpc"
|
||
|
|
"google.golang.org/grpc/metadata"
|
||
|
|
"google.golang.org/protobuf/types/known/emptypb"
|
||
|
|
)
|
||
|
|
|
||
|
|
var ErrUnauthorized = errors.New("unauthorized")
|
||
|
|
|
||
|
|
type pluginAPI struct {
|
||
|
|
pb.UnimplementedPluginServiceServer
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *pluginAPI) Event(stream pb.PluginService_EventServer) error {
|
||
|
|
pl := p.receiverForContext(stream.Context())
|
||
|
|
if pl == nil {
|
||
|
|
return ErrUnauthorized
|
||
|
|
}
|
||
|
|
logger.Sugar.Debugw("attached plugin event stream", "serverName", pl.Name())
|
||
|
|
err := pl.Listen(stream.Context(), &Stream{stream: stream})
|
||
|
|
if err != nil {
|
||
|
|
logger.Sugar.Debugw("detached plugin event stream", "serverName", pl.Name(), "error", err)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
logger.Sugar.Debugw("detached plugin event stream", "serverName", pl.Name())
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *pluginAPI) EventStatus(ctx context.Context, msg *pb.EventStatusMessage) (*emptypb.Empty, error) {
|
||
|
|
pl := p.receiverForContext(ctx)
|
||
|
|
if pl == nil {
|
||
|
|
return nil, ErrUnauthorized
|
||
|
|
}
|
||
|
|
pl.HandleStatus(model.EventRequest{
|
||
|
|
ID: msg.Id,
|
||
|
|
Error: msg.Error,
|
||
|
|
})
|
||
|
|
return &emptypb.Empty{}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *pluginAPI) receiverForContext(ctx context.Context) *plugin.Plugin {
|
||
|
|
md, ok := metadata.FromIncomingContext(ctx)
|
||
|
|
if !ok {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
tokens := md.Get("token")
|
||
|
|
if len(tokens) != 1 {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return server.DefaultManager.PluginByToken(tokens[0])
|
||
|
|
}
|
||
|
|
|
||
|
|
func StartAPIServer() {
|
||
|
|
port := config.Default.PluginAPIPort
|
||
|
|
if port == 0 {
|
||
|
|
port = 3000
|
||
|
|
}
|
||
|
|
socket, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
|
||
|
|
if err != nil {
|
||
|
|
logger.Sugar.Errorf("cannot start plugin API server on port %d: %s", port, err)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
s := grpc.NewServer()
|
||
|
|
pb.RegisterPluginServiceServer(s, &pluginAPI{})
|
||
|
|
logger.Sugar.Debugf("starting plugin server on :%d", port)
|
||
|
|
if err := s.Serve(socket); err != nil {
|
||
|
|
logger.Sugar.Fatalf("cannot start plugin server on :%d: %s", port, err)
|
||
|
|
}
|
||
|
|
}
|