Permalink
Find file Copy path
f76918d Sep 14, 2018
2 contributors

Users who have contributed to this file

@parsiya @nsa
executable file 522 lines (411 sloc) 14.7 KB

flag package

flag package is the Go equivalent of Python argparse. While not as powerful, it does what we expect it to do. It simplifies adding and parsing command line parameters, leaving us to concentrate on the tools. Most of our tools will need them to be actually useful (hardcoding URLs and IPs get old too fast).

Alternative community packages

Some community packages offer what flag does and more. In this guide I am trying to stick to the standard library. Some of these packages are:

  • Cobra: A Commander for modern Go CLI interactions
  • cli: A simple, fast, and fun package for building command line apps in Go

Basic flags

Declaring basic flags is easy. We can create basic types such as: string, bool and int.

A new flag is easy to add:

  • ipPtr := flag.String("ip", "127.0.0.1", "target IP")
    • String: Flag type.
    • ipPtr: Pointer to flag's value.
    • ip: Flag name, meaning flag can be called with -ip.
    • 127.0.0.1: Flag's default value if not provided.
    • target IP: Flag description, displayed with -h switch.

It's also possible to pass a pointer directly:

  • var port int
  • flag.IntVar(&port, "port", 8080, "Port")
// 03.1-01-flag1.go
package main

import (
    "flag"
    "fmt"
)

func main() {

    // Declare flags
    // Remember, flag methods return pointers
    ipPtr := flag.String("ip", "127.0.0.1", "target IP")

    var port int
    flag.IntVar(&port, "port", 8080, "Port")

    verbosePtr := flag.Bool("verbose", true, "verbosity")

    // Parse flags
    flag.Parse()

    // Hack IP:port
    fmt.Printf("Hacking %s:%d!\n", *ipPtr, port)

    // Display progress if verbose flag is set
    if *verbosePtr {
        fmt.Printf("Pew pew!\n")
    }
}

This program contains a mistake! Can you spot it? If not, don't worry.

-h/-help print usage:

$ go run 03.1-01-flag1.go -h
Usage of ... \_obj\exe\03.1-01-flag1.exe:
  -ip string
        target IP (default "127.0.0.1")
  -port int
        Port (default 8080)
  -verbose
        verbosity (default true)
exit status 2

Without any flags, default values are used:

$ go run 03.1-01-flag1.go
Hacking 127.0.0.1:8080!
Pew pew!

Flag use

Flag use is standard.

$ go run 03.1-01-flag1.go -ip 10.20.30.40 -port 12345
Hacking 10.20.30.40:12345!
Pew pew!

The problem is the default value of our boolean flag. A boolean flag is true if it occurs and false if it. We set the default value of verbose to true meaning with our current knowledge we cannot set verbose to false (we will see how below but it's not idiomatic).

Fix that line and run the program again:

$ go run 03.1-02-flag2.go -ip 10.20.30.40 -port 12345
Hacking 10.20.30.40:12345!

$ go run 03.1-02-flag2.go -ip 10.20.30.40 -port 12345 -verbose
Hacking 10.20.30.40:12345!
Pew pew!

= is allowed. Boolean flags can also be set this way (only way to set verbose to false in our previous program):

$ go run 03.1-02-flag2.go -ip=20.30.40.50 -port=54321 -verbose=true
Hacking 20.30.40.50:54321!
Pew pew!

$ go run 03.1-02-flag2.go -ip=20.30.40.50 -port=54321 -verbose=false
Hacking 20.30.40.50:54321!

--flag is also possible:

$ go run 03.1-02-flag2.go --ip 20.30.40.50 --port=12345 --verbose
Hacking 20.30.40.50:12345!
Pew pew!

Declaring flags in the init function

init function is a good location to declare flags. init function is executed after variable initialization values and before main. There's one little catch, variables declared in init are out of focus outside (and in main) hence we need to declare variables outside and use *Var methods:

package main

import (
    "flag"
    "fmt"
)

// Declare flag variables
var (
    ip      string
    port    int
    verbose bool
)

func init() {
    // Declare flags
    // Remember, flag methods return pointers
    flag.StringVar(&ip, "ip", "127.0.0.1", "target IP")

    flag.IntVar(&port, "port", 8080, "Port")

    flag.BoolVar(&verbose, "verbose", false, "verbosity")
}

func main() {

    // Parse flags
    flag.Parse()

    // Hack IP:port
    fmt.Printf("Hacking %s:%d!\n", ip, port)

    // Display progress if verbose flag is set
    if verbose {
        fmt.Printf("Pew pew!\n")
    }
}

Custom flag types and multiple values

Custom flag types are a bit more complicated. Each custom type needs to implement the flag.Value interface. This interface has two methods:

type Value interface {
        String() string
        Set(string) error
}

In simple words:

  1. Create a new type mytype.
  2. Create two methods with *mytype receivers named String() and Set().
    • String() casts the custom type to a string and returns it.
    • Set(string) has a string argument and populates the type and returns an error if applicable.
  3. Create a new flag without an initial value:
    • Call flag.NewFlagSet(&var, instead of flag.String(.
    • Call flag.Var( instead of flag.StringVar( or flag.IntVar(.

Now we can modify our previous example to accept multiple comma-separated IPs. Note, we are using the same structure of generateStrings and consumeString from section 02.6 - sync.WaitGroup. In short, we are going to generate all permutations of IP:ports and then "hack" each of them in one goroutine.

The permutation happens in its own goroutine and is results are sent to channel one by one. When all permutations are generated, channel is closed.

In main, we read from channel and spawn a new goroutine to hack each IP:port. When channel is closed, we wait for all goroutines to finish and then return.

package main

import (
    "errors"
    "flag"
    "fmt"
    "strings"
    "sync"
)

// 1. Create a custom type from a string slice
type strList []string

// 2.1 implement String()
func (str *strList) String() string {
    return fmt.Sprintf("%v", *str)
}

// 2.2 implement Set(*strList)
func (str *strList) Set(s string) error {
    // If input was empty, return an error
    if s == "" {
        return errors.New("nil input")
    }
    // Split input by ","
    *str = strings.Split(s, ",")
    // Do not return an error
    return nil
}

// Declare flag variables
var (
    ip      strList
    port    strList
    verbose bool
)

var wg sync.WaitGroup

func init() {
    // Declare flags
    // Remember, flag methods return pointers
    flag.Var(&ip, "ip", "target IP")

    flag.Var(&port, "port", "Port")

    flag.BoolVar(&verbose, "verbose", false, "verbosity")
}

// permutations creates all permutations of ip:port and sends them to a channel.
// This is preferable to returing a []string because we can spawn it in a
// goroutine and process items in the channel while it's running. Also save
// memory by not creating a large []string that contains all permutations.
func permutations(ips strList, ports strList, c chan<- string) {

    // Close channel when done
    defer close(c)
    for _, i := range ips {
        for _, p := range ports {
            c <- fmt.Sprintf("%s:%s", i, p)
        }
    }
}

// hack spawns a goroutine that "hacks" each target.
// Each goroutine prints a status and display progres if verbose is true
func hack(target string, verbose bool) {

    // Reduce waitgroups counter by one when hack finishes
    defer wg.Done()
    // Hack the planet!
    fmt.Printf("Hacking %s!\n", target)

    // Display progress if verbose flag is set
    if verbose {
        fmt.Printf("Pew pew!\n")
    }
}

func main() {

    // Parse flags
    flag.Parse()

    // Create channel for writing and reading IP:ports
    c := make(chan string)

    // Perform the permutation in a goroutine and send the results to a channel
    // This way we can start "hacking" during permutation generation and
    // not create a huge list of strings in memory
    go permutations(ip, port, c)

    for {
        select {
        // Read a string from channel
        case t, ok := <-c:
            // If channel is closed
            if !ok {
                // Wait until all goroutines are done
                wg.Wait()
                // Print hacking is finished and return
                fmt.Println("Hacking finished!")
                return
            }
            // Otherwise increase wg's counter by one
            wg.Add(1)
            // Spawn a goroutine to hack IP:port read from channel
            go hack(t, verbose)
        }
    }
}

Result:

$ go run 03.1-04-flag4.go -ip 10.20.30.40,50.60.70.80 -port 1234
Hacking 50.60.70.80:1234!
Hacking 10.20.30.40:1234!
Hacking finished!

$ go run 03.1-04-flag4.go -ip 10.20.30.40,50.60.70.80 -port 1234,4321 
Hacking 10.20.30.40:4321!
Hacking 10.20.30.40:1234!
Hacking 50.60.70.80:4321!
Hacking 50.60.70.80:1234!
Hacking finished!

$ go run 03.1-04-flag4.go -ip 10.20.30.40,50.60.70.80 -port 1234,4321 -verbose
Hacking 10.20.30.40:4321!
Pew pew!
Hacking 50.60.70.80:4321!
Pew pew!
Hacking 10.20.30.40:1234!
Pew pew!
Hacking 50.60.70.80:1234!
Pew pew!
Hacking finished!

Required flags

flag does not support this. In Python we can use parser.add_mutually_exclusive_group(). Instead we have to manually check if a flag is set. This can be done by comparing a flag with it's default value or the initial zero value of type in case it does not have a default value.

This can get complicated when the flag can contain the zero value. For example an int flag could be set with value 0 which is the same as the default value for ints. Something that can help is the number of flags after parsing available from flag.NFlag(). If number of flags is less than expected, we know something is wrong.

Alternate and shorthand flags

flag does not have support for shorthand or alternate flags. They need to be declared in a separate statement.

flag.BoolVar(&verbose, "verbose", false, "verbosity")
flag.BoolVar(&verbose, "v", false, "verbosity")

Non-flag arguments

After flag.Parse() it's possible to read other arguments passed to the application with flag.Args(). The number of them is available from flag.NArg() and they individually can be accessed by index using flag.Arg(i).

// 03.1-05-args.go
package main

import (
    "flag"
    "fmt"
)

func main() {
    // Set flag
    _ = flag.Int("flag1", 0, "flag1 description")
    // Parse all flags
    flag.Parse()
    // Enumererate flag.Args()
    for _, v := range flag.Args() {
        fmt.Println(v)
    }
    // Enumerate using flag.Arg(i)
    for i := 0; i < flag.NArg(); i++ {
        fmt.Println(flag.Arg(i))
    }
}

Running the program with non-flag arguments results in:

$ go run 03.1-05-flag5.go -flag1 12 one two 3
one
two
3
one
two
3

Subcommands

Subcommands are possible using flag.NewFlagSet.

  • func NewFlagSet(name string, errorHandling ErrorHandling) *FlagSet

We can decide what happens if parsing that subcommand fails with the second parameter:

const (
    ContinueOnError ErrorHandling = iota // Return a descriptive error.
    ExitOnError                          // Call os.Exit(2).
    PanicOnError                         // Call panic with a descriptive error.
)

After that we need to parse the subcommand. This is usually done by reading os.Args[1] (second argument after program name should be subcommand) and parsing the detected subcommand.

// 03.1-06-subcommand.go
package main

import (
    "flag"
    "fmt"
    "os"
)

var (
    sub1 *flag.FlagSet
    sub2 *flag.FlagSet

    sub1flag  *int
    sub2flag1 *string
    sub2flag2 int

    usage string
)

func init() {
    // Declare subcommand sub1
    sub1 = flag.NewFlagSet("sub1", flag.ExitOnError)
    // int flag for sub1
    sub1flag = sub1.Int("sub1flag", 0, "subcommand1 flag")

    // Declare subcommand sub2
    sub2 = flag.NewFlagSet("sub2", flag.ContinueOnError)
    // string flag for sub2
    sub2flag1 = sub2.String("sub2flag1", "", "subcommand2 flag1")
    // int flag for sub2
    sub2.IntVar(&sub2flag2, "sub2flag2", 0, "subcommand2 flag2")
    // Create usage
    usage = "sub1 -sub1flag (int)\nsub2 -sub2flag1 (string) -sub2flag2 (int)"
}

func main() {
    // If subcommand is not provided, print error, usage and return
    if len(os.Args) < 2 {
        fmt.Println("Not enough parameters")
        fmt.Println(usage)
        return
    }

    // Check the sub command
    switch os.Args[1] {

    // Parse sub1
    case "sub1":
        sub1.Parse(os.Args[2:])

    // Parse sub2
    case "sub2":
        sub2.Parse(os.Args[2:])

    // If subcommand is -h or --help
    case "-h":
        fallthrough
    case "--help":
        fmt.Printf(usage)
        return
    default:
        fmt.Printf("Invalid subcommand %v", os.Args[1])
        return
    }

    // If sub1 was provided and parse, print the flags
    if sub1.Parsed() {
        fmt.Printf("subcommand1 with flag %v\n", *sub1flag)
        return
    }

    // If sub2 was provided and parse, print the flags
    if sub2.Parsed() {
        fmt.Printf("subcommand2 with flags %v, %v\n", *sub2flag1, sub2flag2)
        return
    }
}

As you can see there's a lot of manual work in sub commands and they are not as elegant as normal flags.

Continue reading ⇒ 03.2 - log package