From 552c98c684790235ef0402ab0c5536cf29b4f34b Mon Sep 17 00:00:00 2001
From: perp <perp@tcp.direct>
Date: Sat, 12 Aug 2023 00:22:48 +0100
Subject: [PATCH] :sparkles: Fixed command and added new command

Fixed search command and added init command.
---
 cmd/shogo/commands/init.go   | 25 ++++++++++++++++++++
 cmd/shogo/commands/search.go | 45 +++++++-----------------------------
 2 files changed, 33 insertions(+), 37 deletions(-)
 create mode 100644 cmd/shogo/commands/init.go

diff --git a/cmd/shogo/commands/init.go b/cmd/shogo/commands/init.go
new file mode 100644
index 0000000..0486276
--- /dev/null
+++ b/cmd/shogo/commands/init.go
@@ -0,0 +1,25 @@
+package commands
+
+import (
+	"github.com/spf13/cobra"
+
+	"git.tcp.direct/perp/shogo/internal/shodan"
+	"git.tcp.direct/perp/shogo/internal/utils"
+)
+
+// Init command
+var initCmd = &cobra.Command{
+	Use:     "init",
+	Short:   "Initialize Shogo",
+	Example: "init <key>",
+	Args:    cobra.MinimumNArgs(1),
+	Run: func(_ *cobra.Command, args []string) {
+		utils.CheckColor()
+
+		// Create flags & pool
+		flags := &shodan.Flags{
+			Args: args,
+		}
+		shodan.Pool("init", flags, nil, nil)
+	},
+}
diff --git a/cmd/shogo/commands/search.go b/cmd/shogo/commands/search.go
index 728616a..036a8dd 100644
--- a/cmd/shogo/commands/search.go
+++ b/cmd/shogo/commands/search.go
@@ -1,10 +1,6 @@
 package commands
 
 import (
-	"fmt"
-	"os"
-
-	"github.com/panjf2000/ants/v2"
 	"github.com/spf13/cobra"
 
 	"git.tcp.direct/perp/shogo/internal/shodan"
@@ -27,40 +23,15 @@ var searchCmd = &cobra.Command{
 	Run: func(_ *cobra.Command, args []string) {
 		utils.CheckColor()
 
-		// Results channel
-		results := make(chan string)
-
-		// Goroutine pool
-		pool, err := ants.NewPool(threads)
-		if err != nil {
-			fmt.Printf("%s: %s\n", utils.Red("Error"), err.Error())
-			os.Exit(1)
+		// Create flags & pool
+		flags := &shodan.Flags{
+			Args:      args,
+			Fields:    fields,
+			Separator: separator,
+			Pages:     pages,
+			Page:      1,
 		}
-		defer pool.Release()
-
-		// Query each page
-		for p := 1; p <= pages; p++ {
-			search := &shodan.Search{
-				Query:     args[0],
-				Page:      p,
-				Fields:    fields,
-				Separator: separator,
-				Results:   results,
-			}
-
-			pool.Submit(func() {
-				go search.Search()
-			})
-		}
-
-		// Get page result
-		for j := 1; j <= pages*100; j++ {
-			result := <-results
-			if result != "" {
-				fmt.Println(result)
-			}
-		}
-		defer close(results)
+		shodan.Pool("search", flags, nil, nil)
 	},
 }