Skip to content

Commit

Permalink
extract grpc compoenent
Browse files Browse the repository at this point in the history
  • Loading branch information
@noboruma
noboruma committed Jul 8, 2024
1 parent 8ca24b0 commit bf90277
Show file tree
Hide file tree
Showing 3 changed files with 163 additions and 86 deletions.
15 changes: 14 additions & 1 deletion main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,12 @@ import (

"github.com/deepfence/YaraHunter/pkg/config"
"github.com/deepfence/YaraHunter/pkg/runner"
"github.com/deepfence/YaraHunter/pkg/server"
cfg "github.com/deepfence/match-scanner/pkg/config"
log "github.com/sirupsen/logrus"
"google.golang.org/grpc"

pb "github.com/deepfence/agent-plugins-grpc/srcgo"
)

// Read the regex signatures from config file, options etc.
Expand Down Expand Up @@ -91,5 +95,14 @@ func main() {
go runner.ScheduleYaraHunterUpdater(ctx, runnerOpts)
}

runner.StartYaraHunter(ctx, runnerOpts, config)
runner.StartYaraHunter(ctx, runnerOpts, config,
func(base *server.GRPCScannerServer) server.MalwareRPCServer {
return server.MalwareRPCServer{
GRPCScannerServer: base,
UnimplementedMalwareScannerServer: pb.UnimplementedMalwareScannerServer{},
}
},
func(s grpc.ServiceRegistrar, impl any) {
pb.RegisterMalwareScannerServer(s, impl.(pb.MalwareScannerServer))
})
}
27 changes: 22 additions & 5 deletions pkg/runner/runner.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/deepfence/golang_deepfence_sdk/utils/tasks"
cfg "github.com/deepfence/match-scanner/pkg/config"
log "github.com/sirupsen/logrus"
"google.golang.org/grpc"
)

type RunnerOptions struct {
Expand All @@ -34,20 +35,36 @@ type RunnerOptions struct {
InactiveThreshold int
}

func StartYaraHunter(ctx context.Context, opts RunnerOptions, config cfg.Config) {
func StartYaraHunter[T any](ctx context.Context,
opts RunnerOptions,
config cfg.Config,
constructServer func(srv *server.GRPCScannerServer) T,
attachRegistrar func(s grpc.ServiceRegistrar, impl any)) {

if opts.SocketPath == "" {
runOnce(ctx, opts, config)
return
}

base, err := server.NewGRPCScannerServer(
opts.HostMountPath,
opts.SocketPath,
opts.RulesPath,
opts.InactiveThreshold,
opts.FailOnCompileWarning, config, constants.PluginName,
)
if err != nil {
log.Fatal("Cannot init grpc")
}
go func() {

svc := constructServer(base)

if err := server.RunGrpcServer(ctx,
opts.HostMountPath,
opts.SocketPath,
opts.RulesPath,
opts.InactiveThreshold,
opts.FailOnCompileWarning, config, constants.PluginName); err != nil {
&svc,
attachRegistrar,
); err != nil {
log.Panicf("main: failed to serve: %v", err)
}
}()
Expand Down
207 changes: 127 additions & 80 deletions pkg/server/grpc.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
yararules "github.com/deepfence/YaraHunter/pkg/yararules"
pb "github.com/deepfence/agent-plugins-grpc/srcgo"
tasks "github.com/deepfence/golang_deepfence_sdk/utils/tasks"
"github.com/hillu/go-yara/v4"
log "github.com/sirupsen/logrus"
"google.golang.org/grpc"

Expand All @@ -38,26 +39,29 @@ func init() {
cntnrPathPrefixRegexObj = regexp.MustCompile(cntnrPathPrefixRegex)
}

type gRPCServer struct {
type GRPCScannerServer struct {
HostMountPath, SocketPath string
InactiveThreshold int
extractorConfig cfg.Config
yaraRules *yararules.YaraRules
pluginName string
pb.UnimplementedMalwareScannerServer
pb.UnimplementedAgentPluginServer
pb.UnimplementedScannersServer

scanMap sync.Map
}

func (s *gRPCServer) ReportJobsStatus(context.Context, *pb.Empty) (*pb.JobReports, error) {
type MalwareRPCServer struct {
*GRPCScannerServer
pb.UnimplementedMalwareScannerServer
}

func (s *GRPCScannerServer) ReportJobsStatus(context.Context, *pb.Empty) (*pb.JobReports, error) {
return &pb.JobReports{
RunningJobs: jobs.GetRunningJobCount(),
}, nil
}

func (s *gRPCServer) StopScan(c context.Context, req *pb.StopScanRequest) (*pb.StopScanResult, error) {
func (s *GRPCScannerServer) StopScan(c context.Context, req *pb.StopScanRequest) (*pb.StopScanResult, error) {
scanID := req.ScanId
result := &pb.StopScanResult{
Success: true,
Expand Down Expand Up @@ -85,113 +89,156 @@ func (s *gRPCServer) StopScan(c context.Context, req *pb.StopScanRequest) (*pb.S
return result, nil
}

func (s *gRPCServer) GetName(context.Context, *pb.Empty) (*pb.Name, error) {
func (s *GRPCScannerServer) GetName(context.Context, *pb.Empty) (*pb.Name, error) {
return &pb.Name{Str: s.pluginName}, nil
}

func (s *gRPCServer) GetUID(context.Context, *pb.Empty) (*pb.Uid, error) {
func (s *GRPCScannerServer) GetUID(context.Context, *pb.Empty) (*pb.Uid, error) {
return &pb.Uid{Str: fmt.Sprintf("%s-%s", s.pluginName, s.SocketPath)}, nil
}

func (s *gRPCServer) FindMalwareInfo(c context.Context, r *pb.MalwareRequest) (*pb.MalwareResult, error) {
func (s *MalwareRPCServer) FindMalwareInfo(c context.Context, r *pb.MalwareRequest) (*pb.MalwareResult, error) {
yaraScanner, err := s.yaraRules.NewScanner()
if err != nil {
return &pb.MalwareResult{}, err
}

go func() {
log.Infof("request to scan %+v", r)

jobs.StartScanJob()
defer jobs.StopScanJob()

yaraScanner, err := s.yaraRules.NewScanner()
scanner := scan.New(s.HostMountPath, s.extractorConfig, yaraScanner, r.ScanId)
res, ctx := tasks.StartStatusReporter(
r.ScanId,
func(status tasks.ScanStatus) error {
output.WriteScanStatus(status.ScanStatus, status.ScanId, status.ScanMessage)
return nil
},
tasks.StatusValues{
IN_PROGRESS: "IN_PROGRESS",
CANCELLED: "CANCELLED",
FAILED: "ERROR",
SUCCESS: "COMPLETE",
},
time.Duration(s.InactiveThreshold)*time.Second)
s.scanMap.Store(scanner.ScanID, ctx)
defer func() {
s.scanMap.Delete(scanner.ScanID)
res <- err
close(res)
}()

// Check for error only after the StartStatusReporter as we have to report
// the error if we failed to create the yara scanner
if err != nil {
log.Error("Failed to create Yara Scanner, error:", err)
return
}
writeToFile := func(res output.IOCFound, scanID string) {
output.WriteScanData([]*pb.MalwareInfo{output.MalwaresToMalwareInfo(res)}, scanID)
}
namespace := ""
container := ""
image := ""
path := ""
switch {
case r.GetPath() != "":
log.Infof("scan for malwares in path %s", r.GetPath())
err = scanner.Scan(ctx, scan.DirScan, "", r.GetPath(), r.GetScanId(), writeToFile)
case r.GetImage() != nil && r.GetImage().Name != "":
log.Infof("scan for malwares in image %s", r.GetImage())
err = scanner.Scan(ctx, scan.ImageScan, "", r.GetImage().Name, r.GetScanId(), writeToFile)
case r.GetContainer() != nil && r.GetContainer().Id != "":
log.Infof("scan for malwares in container %s", r.GetContainer())
err = scanner.Scan(ctx, scan.ContainerScan, r.GetContainer().Namespace, r.GetContainer().Id, r.GetScanId(), writeToFile)
case r.GetContainer() != nil:
namespace = r.GetContainer().GetNamespace()
container = r.GetContainer().GetId()
case r.GetImage() != nil:
image = r.GetImage().GetName()
default:
err = fmt.Errorf("invalid request")
path = r.GetPath()
}

if err != nil {
println(err.Error())
}
DoScan(
r.ScanId,
s.HostMountPath,
s.extractorConfig,
s.InactiveThreshold,
&s.scanMap,
namespace,
path,
image,
container,
yaraScanner,
func(res output.IOCFound, scanID string) {
output.WriteScanData([]*pb.MalwareInfo{output.MalwaresToMalwareInfo(res)}, scanID)
},
)
}()
return &pb.MalwareResult{}, nil
}

func RunGrpcServer(ctx context.Context,
func DoScan(
scanID string,
hostMountPath string,
extractorConfig cfg.Config,
inactiveThreshold int,
scanMap *sync.Map,
namespace,
path, image, container string,
yaraScanner *yara.Scanner, writeToFile func(res output.IOCFound, scanID string)) {

jobs.StartScanJob()
defer jobs.StopScanJob()

scanner := scan.New(hostMountPath, extractorConfig, yaraScanner, scanID)
res, ctx := tasks.StartStatusReporter(
scanID,
func(status tasks.ScanStatus) error {
output.WriteScanStatus(status.ScanStatus, status.ScanId, status.ScanMessage)
return nil
},
tasks.StatusValues{
IN_PROGRESS: "IN_PROGRESS",
CANCELLED: "CANCELLED",
FAILED: "ERROR",
SUCCESS: "COMPLETE",
},
time.Duration(inactiveThreshold)*time.Second)
scanMap.Store(scanner.ScanID, ctx)
var err error
defer func() {
scanMap.Delete(scanner.ScanID)
res <- err
close(res)
}()

switch {
case path != "":
log.Infof("scan for malwares in path %s", path)
err = scanner.Scan(ctx, scan.DirScan, "", path, scanID, writeToFile)
case image != "":
log.Infof("scan for malwares in image %s", image)
err = scanner.Scan(ctx, scan.ImageScan, "", image, scanID, writeToFile)
case container != "":
log.Infof("scan for malwares in container %s", container)
err = scanner.Scan(ctx, scan.ContainerScan, namespace, container, scanID, writeToFile)
default:
err = fmt.Errorf("invalid request")
}

if err != nil {
log.Errorf("err: %v", err)
}
}

func NewGRPCScannerServer(
hostMoundPath, socketPath, rulesPath string,
InactiveThreshold int,

Check failure on line 198 in pkg/server/grpc.go

View workflow job for this annotation

GitHub Actions / lint 

captLocal: `InactiveThreshold' should not be capitalized (gocritic)
failOnCompileWarning bool,
config cfg.Config, pluginName string) error {

lis, err := net.Listen("unix", socketPath)
config cfg.Config, pluginName string,
) (*GRPCScannerServer, error) {
yaraRules := yararules.New(rulesPath)
err := yaraRules.Compile(constants.Filescan, failOnCompileWarning)
if err != nil {
return err
return nil, err
}
s := grpc.NewServer()

impl := &gRPCServer{
HostMountPath: hostMoundPath,
SocketPath: socketPath,
InactiveThreshold: InactiveThreshold,
extractorConfig: config,
yaraRules: &yararules.YaraRules{},
pluginName: pluginName,
UnimplementedMalwareScannerServer: pb.UnimplementedMalwareScannerServer{},
UnimplementedAgentPluginServer: pb.UnimplementedAgentPluginServer{},
UnimplementedScannersServer: pb.UnimplementedScannersServer{},
scanMap: sync.Map{},
res := &GRPCScannerServer{
HostMountPath: hostMoundPath,
SocketPath: socketPath,
InactiveThreshold: InactiveThreshold,
extractorConfig: config,
yaraRules: yaraRules,
pluginName: pluginName,
UnimplementedAgentPluginServer: pb.UnimplementedAgentPluginServer{},
UnimplementedScannersServer: pb.UnimplementedScannersServer{},
scanMap: sync.Map{},
}

return res, nil
}

func RunGrpcServer(ctx context.Context,
socketPath string,
impl any,
customImpl func(s grpc.ServiceRegistrar, impl any)) error {

lis, err := net.Listen("unix", socketPath)
if err != nil {
return err
}

impl.scanMap = sync.Map{}
s := grpc.NewServer()

// compile yara rules
impl.yaraRules = yararules.New(rulesPath)
err = impl.yaraRules.Compile(constants.Filescan, failOnCompileWarning)
if err != nil {
return err
}

pb.RegisterAgentPluginServer(s, impl)
pb.RegisterMalwareScannerServer(s, impl)
pb.RegisterScannersServer(s, impl)
pb.RegisterAgentPluginServer(s, impl.(pb.AgentPluginServer))
pb.RegisterScannersServer(s, impl.(pb.ScannersServer))
customImpl(s, impl)

log.Info("main: server listening at ", lis.Addr())
go func() {
if err := s.Serve(lis); err != nil {
Expand Down

0 comments on commit bf90277

Please sign in to comment.