add --wait flag to undo action after the time expires or users ctrl+c…

… input
guumaster · Apr 9, 2020 · 122cc0e · 122cc0e
1 parent 7925b01
commit 122cc0e
Show file tree

Hide file tree

Showing 10 changed files with 226 additions and 174 deletions.
diff --git a/cmd/add.go b/cmd/add.go
@@ -1,16 +1,15 @@
 package cmd
 
 import (
-	"fmt"
 	"os"
 
 	"github.com/spf13/cobra"
 
 	"github.com/guumaster/hostctl/pkg/host"
 )
 
-// addFromFileCmd represents the fromFile command
-var addFromFileCmd = &cobra.Command{
+// addCmd represents the fromFile command
+var addCmd = &cobra.Command{
 	Use:   "add",
 	Short: "Add content to a profile in your hosts file.",
 	Long: `
@@ -33,39 +32,32 @@ If the profile already exists it will be added to it.`,
 		from, _ := cmd.Flags().GetString("from")
 		profile, _ := cmd.Flags().GetString("profile")
 
-		var err error
 		if isPiped() {
-			fmt.Println("IS PIPED")
-			err = host.AddFromReader(os.Stdin, &host.AddFromFileOptions{
-				Dst:     src,
-				Profile: profile,
-				Reset:   false,
-			})
-		} else {
-			fmt.Println("FROM FILE")
-			err = host.AddFromFile(&host.AddFromFileOptions{
-				From:    from,
+			return host.AddFromReader(os.Stdin, &host.AddFromFileOptions{
 				Dst:     src,
 				Profile: profile,
 				Reset:   false,
 			})
 		}
-		if err != nil {
-			return err
-		}
 
-		return host.ListProfiles(src, &host.ListOptions{
+		return host.AddFromFile(&host.AddFromFileOptions{
+			From:    from,
+			Dst:     src,
 			Profile: profile,
+			Reset:   false,
 		})
 	},
+	PostRunE: func(cmd *cobra.Command, args []string) error {
+		return postActionCmd(cmd, args, removeCmd)
+	},
 }
 
 func init() {
-	rootCmd.AddCommand(addFromFileCmd)
-
-	addFromFileCmd.Flags().StringP("from", "f", "", "file to read")
+	rootCmd.AddCommand(addCmd)
+	addCmd.AddCommand(addDomainsCmd)
 
-	addFromFileCmd.AddCommand(addDomainsCmd)
+	addCmd.Flags().StringP("from", "f", "", "file to read")
+	addCmd.PersistentFlags().DurationP("wait", "w", -1, "Enables a profile for a specific amount of time")
 
 	addDomainsCmd.Flags().String("ip", "127.0.0.1", "domains ip")
 }
diff --git a/cmd/add_domains.go b/cmd/add_domains.go
@@ -29,7 +29,6 @@ If the profile already exists it will be added to it.`,
 		src, _ := cmd.Flags().GetString("host-file")
 		ip, _ := cmd.Flags().GetString("ip")
 		profile, _ := cmd.Flags().GetString("profile")
-		quiet, _ := cmd.Flags().GetBool("quiet")
 
 		err := host.AddFromArgs(&host.AddFromArgsOptions{
 			Domains: args,
@@ -42,17 +41,9 @@ If the profile already exists it will be added to it.`,
 			return err
 		}
 
-		err = host.Enable(src, profile)
-		if err != nil {
-			return err
-		}
-
-		if quiet {
-			return nil
-		}
-
-		return host.ListProfiles(src, &host.ListOptions{
-			Profile: profile,
-		})
+		return host.Enable(src, profile)
+	},
+	PostRunE: func(cmd *cobra.Command, args []string) error {
+		return postActionCmd(cmd, args, removeDomainsCmd)
 	},
 }
diff --git a/cmd/disable.go b/cmd/disable.go
@@ -6,54 +6,51 @@ import (
 	"github.com/guumaster/hostctl/pkg/host"
 )
 
-// disableCmd represents the disable command
-var disableCmd = &cobra.Command{
-	Use:   "disable",
-	Short: "Disable a profile from your hosts file.",
-	Long: `
+var disableCmd *cobra.Command
+
+func init() {
+
+	// disableCmd represents the disable command
+	disableCmd := &cobra.Command{
+		Use:   "disable",
+		Short: "Disable a profile from your hosts file.",
+		Long: `
 Disable a profile from your hosts file without removing it.
 It will be listed as "off" while it is disabled.
 `,
-	PreRunE: func(cmd *cobra.Command, args []string) error {
-		profile, _ := cmd.Flags().GetString("profile")
-		all, _ := cmd.Flags().GetBool("all")
-
-		if !all && profile == "" {
-			return host.MissingProfileError
-		}
-
-		if profile == "default" {
-			return host.DefaultProfileError
-		}
-		return nil
-	},
-	RunE: func(cmd *cobra.Command, args []string) error {
-		src, _ := cmd.Flags().GetString("host-file")
-		profile, _ := cmd.Flags().GetString("profile")
-		quiet, _ := cmd.Flags().GetBool("quiet")
-
-		all, _ := cmd.Flags().GetBool("all")
-
-		var err error
-		if all {
-			profile = ""
-		}
-		err = host.Disable(src, profile)
-		if err != nil {
-			return err
-		}
-
-		if quiet {
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			profile, _ := cmd.Flags().GetString("profile")
+			all, _ := cmd.Flags().GetBool("all")
+
+			if !all && profile == "" {
+				return host.MissingProfileError
+			}
+
+			if profile == "default" {
+				return host.DefaultProfileError
+			}
 			return nil
-		}
-		return host.ListProfiles(src, &host.ListOptions{
-			Profile: profile,
-		})
-	},
-}
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			src, _ := cmd.Flags().GetString("host-file")
+			profile, _ := cmd.Flags().GetString("profile")
+
+			all, _ := cmd.Flags().GetBool("all")
+
+			if all {
+				profile = ""
+			}
+			return host.Disable(src, profile)
+
+		},
+		PostRunE: func(cmd *cobra.Command, args []string) error {
+			return postActionCmd(cmd, args, enableCmd)
+		},
+	}
 
-func init() {
 	rootCmd.AddCommand(disableCmd)
 
 	disableCmd.Flags().BoolP("all", "", false, "Disable all profiles")
+	disableCmd.Flags().DurationP("wait", "w", -1, "Enables a profile for a specific amount of time")
+
 }
diff --git a/cmd/enable.go b/cmd/enable.go
@@ -6,61 +6,54 @@ import (
 	"github.com/guumaster/hostctl/pkg/host"
 )
 
-// enableCmd represents the enable command
-var enableCmd = &cobra.Command{
-	Use:   "enable",
-	Short: "Enable a profile on your hosts file.",
-	Long: `
+var enableCmd *cobra.Command
+
+func init() {
+	// enableCmd represents the enable command
+	enableCmd = &cobra.Command{
+		Use:   "enable",
+		Short: "Enable a profile on your hosts file.",
+		Long: `
 Enables an existing profile.
 It will be listed as "on" while it is enabled.
 `,
-	PreRunE: func(cmd *cobra.Command, args []string) error {
-		profile, _ := cmd.Flags().GetString("profile")
-		all, _ := cmd.Flags().GetBool("all")
-
-		if !all && profile == "" {
-			return host.MissingProfileError
-		}
-
-		if profile == "default" {
-			return host.DefaultProfileError
-		}
-		return nil
-	},
-	RunE: func(cmd *cobra.Command, args []string) error {
-		profile, _ := cmd.Flags().GetString("profile")
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			profile, _ := cmd.Flags().GetString("profile")
+			all, _ := cmd.Flags().GetBool("all")
 
-		all, _ := cmd.Flags().GetBool("all")
-		if all {
-			profile = ""
-		}
+			if !all && profile == "" {
+				return host.MissingProfileError
+			}
 
-		src, _ := cmd.Flags().GetString("host-file")
-		enableOnly, _ := cmd.Flags().GetBool("only")
-		quiet, _ := cmd.Flags().GetBool("quiet")
-
-		var err error
-		if enableOnly {
-			err = host.EnableOnly(src, profile)
-		} else {
-			err = host.Enable(src, profile)
-		}
-		if err != nil {
-			return err
-		}
-
-		if quiet {
+			if profile == "default" {
+				return host.DefaultProfileError
+			}
 			return nil
-		}
-		return host.ListProfiles(src, &host.ListOptions{
-			Profile: profile,
-		})
-	},
-}
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			profile, _ := cmd.Flags().GetString("profile")
+			src, _ := cmd.Flags().GetString("host-file")
+			enableOnly, _ := cmd.Flags().GetBool("only")
+
+			all, _ := cmd.Flags().GetBool("all")
+			if all {
+				profile = ""
+			}
+
+			if enableOnly && !all {
+				return host.EnableOnly(src, profile)
+			}
+			return host.Enable(src, profile)
+
+		},
+		PostRunE: func(cmd *cobra.Command, args []string) error {
+			return postActionCmd(cmd, args, disableCmd)
+		},
+	}
 
-func init() {
 	rootCmd.AddCommand(enableCmd)
 
 	enableCmd.Flags().BoolP("all", "", false, "Enable all profiles")
 	enableCmd.Flags().Bool("only", false, "Disable all other profiles")
+	enableCmd.Flags().DurationP("wait", "w", -1, "Enables a profile for a specific amount of time")
 }
diff --git a/cmd/list.go b/cmd/list.go
@@ -23,13 +23,11 @@ The "default" profile is all the content that is not handled by hostctl tool.
 		raw, _ := cmd.Flags().GetBool("raw")
 		cols, _ := cmd.Flags().GetStringSlice("column")
 
-		err := host.ListProfiles(src, &host.ListOptions{
+		return host.ListProfiles(src, &host.ListOptions{
 			Profile:  profile,
 			RawTable: raw,
 			Columns:  cols,
 		})
-
-		return err
 	},
 }
 
@@ -39,6 +37,4 @@ func init() {
 	listCmd.AddCommand(makeListStatusCmd(host.Enabled))
 	listCmd.AddCommand(makeListStatusCmd(host.Disabled))
 
-	listCmd.PersistentFlags().StringSliceP("column", "c", nil, "Columns to show on lists")
-	listCmd.PersistentFlags().Bool("raw", false, "Output without table borders")
 }