From e72b11c2d8a286da0da0a390820242e8101a2642 Mon Sep 17 00:00:00 2001
From: Ernesto ponce <ernestoponce27@gmail.com>
Date: Tue, 13 Feb 2024 17:49:01 -0300
Subject: [PATCH] refactor flags,  add donwloads implementations

---
 main.go     | 53 ++++++++++++++++++-----------------------------------
 protocol.go | 30 ++++++++++++++++++++++++++----
 2 files changed, 44 insertions(+), 39 deletions(-)

diff --git a/main.go b/main.go
index 8f9f565..990f06b 100644
--- a/main.go
+++ b/main.go
@@ -55,33 +55,30 @@ var styleMessage = lipgloss.NewStyle().
 	Background(lipgloss.Color("#606683")).Padding(1).MarginBottom(1).MarginTop(1).MarginLeft(1)
 
 func main() {
-	typeParam := flag.String("t", "http", "type connection (ws, gq, http, grpc)")
-	urlParam := flag.String("u", "", "url to connect")
-	messageParam := flag.String("p", "", "data send to server")
-	queryParam := flag.String("q", "", "query params")
-	verboseParam := flag.Bool("v", false, "show response server headers")
-	headerP := flag.String("h", "", "header params")
-	method := flag.String("m", "GET", "method request")
-	file := flag.String("f", "", "file path")
-	download := flag.Bool("d", false, "download content")
+	params := Params{}
+
+	flag.StringVar(&params.typeP, "t", "http", "type connection (ws, gq, http, grpc)")
+	flag.StringVar(&params.url, "u", "", "url to connect")
+	flag.StringVar(&params.message, "p", "", "data send to server")
+	flag.StringVar(&params.query, "q", "", "query params")
+	flag.BoolVar(&params.verbose, "v", false, "show response server headers")
+	flag.StringVar(&params.header, "h", "", "header params")
+	flag.StringVar(&params.method, "m", "GET", "method request")
+	flag.StringVar(&params.file, "f", "", "file path")
 
 	// GRPC exclusive flags
-	importPath := flag.String("import-path", "", "The path to a directory from which proto sources can be imported, for use with -proto flags")
-	proto := flag.String("proto", "", "The proto file")
-	methodName := flag.String("method", "", "The service method to call")
+	flag.StringVar(&params.importPath, "import-path", "", "The path to a directory from which proto sources can be imported, for use with -proto flags")
+	flag.StringVar(&params.proto, "proto", "", "The proto file")
+	flag.StringVar(&params.methodName, "method", "", "The service method to call")
 
+	download := flag.Bool("d", false, "download content")
 	flag.Parse()
 
-	if *urlParam == "" {
-		flag.PrintDefaults()
-		return
-	}
-
 	validProtocos := []string{"ws", "gq", "http", "grpc"}
 
 	valid := false
 	for _, v := range validProtocos {
-		if *typeParam == v {
+		if params.typeP == v {
 			valid = true
 			break
 		}
@@ -92,24 +89,10 @@ func main() {
 		return
 	}
 
-	params := Params{
-		typeP:      *typeParam,
-		url:        *urlParam,
-		query:      *queryParam,
-		header:     *headerP,
-		message:    *messageParam,
-		method:     *method,
-		file:       *file,
-		importPath: *importPath,
-		proto:      *proto,
-		methodName: *methodName,
-		verbose:    *verboseParam,
-	}
-
 	p := getProtocol(params)
 	defer p.Close()
 
-	if *messageParam == "" {
+	if params.message == "" {
 		p.OnMessageReceived()
 	}
 
@@ -144,7 +127,7 @@ func main() {
 			return
 		}
 		fmt.Print(dr)
-		if *verboseParam {
+		if params.verbose {
 			p.PrintHeaderResponse()
 		}
 		return
@@ -157,7 +140,7 @@ func main() {
 	}
 	fmt.Print(out)
 
-	if *verboseParam {
+	if params.verbose {
 		p.PrintHeaderResponse()
 	}
 }
diff --git a/protocol.go b/protocol.go
index 0fe6328..e32e189 100644
--- a/protocol.go
+++ b/protocol.go
@@ -38,6 +38,7 @@ type GraphQL struct {
 	url      string
 	query    string
 	httpResp *http.Response
+	response string
 }
 
 func NewGrapQL(params Params) *GraphQL {
@@ -69,11 +70,22 @@ func (g *GraphQL) RequestResponse() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	return string(body), nil
+
+	g.response = string(body)
+	return g.response, nil
 }
 
 func (g *GraphQL) OnMessageReceived() {}
-func (g *GraphQL) Download() error    { return nil }
+
+func (g *GraphQL) Download() error {
+	err := saveToFile(g.response, g.url)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (g *GraphQL) PrintHeaderResponse() {
 	printHttpResponse(g.httpResp)
 }
@@ -348,6 +360,7 @@ type GRPC struct {
 	message    string
 	methodName string
 	verbose    bool
+	response   string
 }
 
 func NewGRPC(params Params) *GRPC {
@@ -444,12 +457,21 @@ func (g *GRPC) RequestResponse() (string, error) {
 		grpcurl.PrintStatus(os.Stderr, h.Status, formatter)
 	}
 
-	return buf.String(), nil
+	g.response = buf.String()
+
+	return g.response, nil
 }
 
 func (g *GRPC) OnMessageReceived()   {}
 func (g *GRPC) PrintHeaderResponse() {}
-func (g *GRPC) Download() error      { return nil }
+func (g *GRPC) Download() error {
+	err := saveToFile(g.response, g.url)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
 
 func (g *GRPC) Close() {
 	if g.cc != nil {