Merge pull request #10 from projectdiscovery/feature-cli-custom-message

adding custom cli notification format

internal/runner/config.go

@@ -34,6 +34,7 @@ type ConfigFile struct {
 	Interval    int    `yaml:"interval,omitempty"`
 	HTTPMessage string `yaml:"http_message,omitempty"`
 	DNSMessage  string `yaml:"dns_message,omitempty"`
+	CLIMessage  string `yaml:"cli_message,omitempty"`
 }
 
 // GetConfigDirectory from the system

internal/runner/options.go

@@ -32,6 +32,7 @@ type Options struct {
 	Interval                int
 	HTTPMessage             string
 	DNSMessage              string
+	CLIMessage              string
 }
 
 // ParseConfigFileOrOptions combining all settings
@@ -57,6 +58,7 @@ func ParseConfigFileOrOptions() *Options {
 	flag.IntVar(&options.Interval, "interval", 2, "Polling interval in seconds")
 	flag.StringVar(&options.HTTPMessage, "message-http", defaultHTTPMessage, "HTTP Message")
 	flag.StringVar(&options.DNSMessage, "message-dns", defaultDNSMessage, "DNS Message")
+	flag.StringVar(&options.CLIMessage, "message-cli", defaultCLIMessage, "CLI Message")
 
 	flag.Parse()
 
@@ -99,15 +101,17 @@ func (options *Options) configureOutput() {
 func (options *Options) writeDefaultConfig() {
 	configFile, err := getDefaultConfigFile()
 	if err != nil {
-		gologger.Warningf("Could not get default configuration file: %s\n", err)
+		gologger.Printf("Could not get default configuration file: %s\n", err)
 	}
 
 	if fileExists(configFile) {
+		gologger.Printf("Found existing config file: %s\n", configFile)
 		return
 	}
 
 	// Skip config file creation if run as root to avoid permission issues
 	if os.Getuid() == 0 {
+		gologger.Printf("Running as root, skipping config file write to avoid permissions issues: %s\n", configFile)
 		return
 	}
 
@@ -138,23 +142,24 @@ func (options *Options) writeDefaultConfig() {
 		"```\n" +
 		"{{request}}\n" +
 		"```"
+	dummyConfig.CLIMessage = "{{data}}"
 
 	err = dummyConfig.MarshalWrite(configFile)
 	if err != nil {
-		gologger.Warningf("Could not write configuration file to %s: %s\n", configFile, err)
+		gologger.Printf("Could not write configuration file to %s: %s\n", configFile, err)
 		return
 	}
 
 	// turn all lines into comments
 	origFile, err := os.Open(configFile)
 	if err != nil {
-		gologger.Warningf("Could not process temporary file: %s\n", err)
+		gologger.Printf("Could not process temporary file: %s\n", err)
 		return
 	}
 	tmpFile, err := ioutil.TempFile("", "")
 	if err != nil {
 		log.Println(err)
-		gologger.Warningf("Could not process temporary file: %s\n", err)
+		gologger.Printf("Could not process temporary file: %s\n", err)
 		return
 	}
 	sc := bufio.NewScanner(origFile)
@@ -170,15 +175,15 @@ func (options *Options) writeDefaultConfig() {
 	//nolint:errcheck // silent fail
 	os.Rename(tmpFileName, configFile)
 
-	gologger.Infof("Configuration file saved to %s\n", configFile)
+	gologger.Printf("Configuration file saved to %s\n", configFile)
 }
 
 // MergeFromConfig with existing options
 func (options *Options) MergeFromConfig(configFileName string, ignoreError bool) {
 	configFile, err := UnmarshalRead(configFileName)
 	if err != nil {
 		if ignoreError {
-			gologger.Warningf("Could not read configuration file %s: %s\n", configFileName, err)
+			gologger.Printf("Could not read configuration file %s - ignoring error: %s\n", configFileName, err)
 			return
 		}
 		gologger.Fatalf("Could not read configuration file %s: %s\n", configFileName, err)
@@ -226,6 +231,9 @@ func (options *Options) MergeFromConfig(configFileName string, ignoreError bool)
 	if configFile.DNSMessage != "" {
 		options.DNSMessage = configFile.DNSMessage
 	}
+	if configFile.CLIMessage != "" {
+		options.CLIMessage = configFile.CLIMessage
+	}
 	if configFile.Interval > 0 {
 		options.Interval = configFile.Interval
 	}

internal/runner/runner.go

@@ -16,6 +16,7 @@ import (
 const (
 	defaultHTTPMessage = "The collaborator server received an {{protocol}} request from {{from}} at {{time}}:\n```\n{{request}}\n{{response}}```"
 	defaultDNSMessage  = "The collaborator server received a DNS lookup of type {{type}} for the domain name {{domain}} from {{from}} at {{time}}:\n```{{request}}```"
+	defaultCLIMessage  = "{{data}}"
 )
 
 // Runner contains the internal logic of the program
@@ -56,6 +57,10 @@ func (r *Runner) Run() error {
 		br := bufio.NewScanner(os.Stdin)
 		for br.Scan() {
 			msg := br.Text()
+			rr := strings.NewReplacer(
+				"{{data}}", msg,
+			)
+			msg = rr.Replace(r.options.CLIMessage)
 			gologger.Printf(msg)
 			//nolint:errcheck // silent fail
 			r.notifier.SendNotification(msg)