add --max, --before, --after switches to dtail and dgrep commands

7 files changed, 60 insertions, 16 deletions

diff --git a/internal/clients/args.go b/internal/clients/args.go
index 7f782f1..67d2044 100644
--- a/internal/clients/args.go
+++ b/internal/clients/args.go
@@ -6,14 +6,22 @@ import (
 	gossh "golang.org/x/crypto/ssh"
 )
 
+// LineContext is here to help filtering out only specific lines.
+type LineContext struct {
+	RegexStr      string
+	AfterContext  int
+	BeforeContext int
+	MaxCount      int
+}
+
 // Args is a helper struct to summarize common client arguments.
 type Args struct {
+	LineContext
 	Mode               omode.Mode
 	ServersStr         string
 	UserName           string
 	What               string
 	Arguments          []string
-	RegexStr           string
 	RegexInvert        bool
 	TrustAllHosts      bool
 	Discovery          string
@@ -22,5 +30,5 @@ type Args struct {
 	SSHAuthMethods     []gossh.AuthMethod
 	SSHHostKeyCallback gossh.HostKeyCallback
 	PrivateKeyPathFile string
-	Quiet            bool
+	Quiet              bool
 }
diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index f20156f..0bdd62e 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -2,6 +2,8 @@ package clients
 
 import (
 	"context"
+	"fmt"
+	"strings"
 	"sync"
 	"time"
 
@@ -118,10 +120,44 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con
 	}
 }
 
+func (c *baseClient) makeCommandOptions() map[string]string {
+	options := make(map[string]string)
+
+	if c.Args.Quiet {
+		options["quiet"] = fmt.Sprintf("%v", c.Args.Quiet)
+	}
+	if c.Args.LineContext.MaxCount != 0 {
+		options["max"] = fmt.Sprintf("%d", c.Args.LineContext.MaxCount)
+	}
+	if c.Args.LineContext.BeforeContext != 0 {
+		options["before"] = fmt.Sprintf("%d", c.Args.LineContext.BeforeContext)
+	}
+	if c.Args.LineContext.AfterContext != 0 {
+		options["after"] = fmt.Sprintf("%d", c.Args.LineContext.AfterContext)
+	}
+
+	return options
+}
+
+func (c *baseClient) commandOptionsToString(options map[string]string) string {
+	var sb strings.Builder
+
+	count := 0
+	for k, v := range options {
+		if count > 0 {
+			sb.WriteString(":")
+		}
+		sb.WriteString(fmt.Sprintf("%s=%s", k, v))
+		count++
+	}
+
+	return sb.String()
+}
+
 func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) *remote.Connection {
 	conn := remote.NewConnection(server, c.UserName, sshAuthMethods, hostKeyCallback)
 	conn.Handler = c.maker.makeHandler(server)
-	conn.Commands = c.maker.makeCommands()
+	conn.Commands = c.maker.makeCommands(c.makeCommandOptions())
 
 	return conn
 }
diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go
index b7b6131..db892f1 100644
--- a/internal/clients/catclient.go
+++ b/internal/clients/catclient.go
@@ -41,10 +41,10 @@ func (c CatClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c CatClient) makeCommands() (commands []string) {
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
+func (c CatClient) makeCommands(options map[string]string) (commands []string) {
+	optionsStr := c.commandOptionsToString(options)
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
 	}
 	return
 }
diff --git a/internal/clients/grepclient.go b/internal/clients/grepclient.go
index 652c31b..567193a 100644
--- a/internal/clients/grepclient.go
+++ b/internal/clients/grepclient.go
@@ -40,10 +40,10 @@ func (c GrepClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c GrepClient) makeCommands() (commands []string) {
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
+func (c GrepClient) makeCommands(options map[string]string) (commands []string) {
+	optionsStr := c.commandOptionsToString(options)
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
 	}
 
 	return
diff --git a/internal/clients/maker.go b/internal/clients/maker.go
index d5ffd8b..a1d6864 100644
--- a/internal/clients/maker.go
+++ b/internal/clients/maker.go
@@ -9,5 +9,5 @@ import (
 // and send different commands to the DTail server.
 type maker interface {
 	makeHandler(server string) handlers.Handler
-	makeCommands() (commands []string)
+	makeCommands(options map[string]string) (commands []string)
 }
diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go
index 1c0c2cc..6a06aaf 100644
--- a/internal/clients/maprclient.go
+++ b/internal/clients/maprclient.go
@@ -110,21 +110,21 @@ func (c MaprClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewMaprHandler(server, c.query, c.globalGroup)
 }
 
-func (c MaprClient) makeCommands() (commands []string) {
+func (c MaprClient) makeCommands(options map[string]string) (commands []string) {
 	commands = append(commands, fmt.Sprintf("map %s", c.query.RawQuery))
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
 
 	modeStr := "cat"
 	if c.Mode == omode.TailClient {
 		modeStr = "tail"
 	}
+	optionsStr := c.commandOptionsToString(options)
 
 	for _, file := range strings.Split(c.What, ",") {
 		if c.Timeout > 0 {
 			commands = append(commands, fmt.Sprintf("timeout %d %s %s %s", c.Timeout, modeStr, file, c.Regex.Serialize()))
 			continue
 		}
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", modeStr, options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", modeStr, optionsStr, file, c.Regex.Serialize()))
 	}
 
 	return
diff --git a/internal/clients/tailclient.go b/internal/clients/tailclient.go
index cefbaa7..853ef1d 100644
--- a/internal/clients/tailclient.go
+++ b/internal/clients/tailclient.go
@@ -37,10 +37,10 @@ func (c TailClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c TailClient) makeCommands() (commands []string) {
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
+func (c TailClient) makeCommands(options map[string]string) (commands []string) {
+	optionsStr := c.commandOptionsToString(options)
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
 	}
 	logger.Debug(commands)