Building Command Line Tools with Go

blog-image

Using the cdr/cli package.

Estimated Reading Time : 7m

Options

Go is a great programming language for building command-line tools.

The most popular command line package in the eco-system being spf13’s Cobra.

The standard library is also a great option when you combine the os and flag packages.

You may be asking yourself, which one is right for me?

Choosing the right package

The package you choose should be based on the functionality you’re looking for.

If you’re just building a simple interface for application start up, you may not need something as robust as cobra.

You could probably get away with just using the standard library if all you need is some rudimentary flag parsing.

Now what if you’re in the middle?

What if you know you don’t need half the things cobra offers but you also need something more structured than a single switch statement for evaluating all your operating system arguments?

Enter cdr/cli

cdr/cli is a minimalistic go package for building command-line tools.

It’s simple, easy to use, and allows you to quickly put a CLI together.

Lets get a little more familiar with this package by building a port-scanner in go.

Project structure

Lets start by initializing our mod file and an idiomatic go application structure.

├── cmd
│   └── port-scanner
|       ├── scan.go
│       ├── root.go
│       └── main.go
├── go.mod
├── go.sum
└── port-scanner

Another common pattern is to keep each command in it’s own dir under cmd, including that commands subcommands.

You can find an example of that approach in the go standard library.

Our root command will only have one subcommand for our demo so I feel that justifies keeping them in the same directory for the purpose of this demo.

The port-scanner/cmd/port-scanner dir is our entry-point.

We can define our root command here.

Root command

port-scanner/cmd/port-scanner/root.go

package main

import (
    "github.com/spf13/pflag"
    "go.coder.com/cli"
)

// We define our command as a struct.
type root struct{}

// Our command struct needs to implement the Command interface as defined by cdr/cli.
// See: https://pkg.go.dev/go.coder.com/cli#Command for more details.
func (r *root) Run(fl *pflag.FlagSet) {
    fl.Usage()
}

// We'll also need to implement a method for returning a command spec to fully satisfy the
// the aforementioned Command interface.
func (r *root) Spec() cli.CommandSpec {
    return cli.CommandSpec{
        Name:  "port-scanner",
        Usage: "[subcommand] [flags]",
        Desc:  "A simple port-scanner.",
    }
}

// We can now attach subcommands to our command struct; implementing ParentCommand as a result as defined by cdr/cli.
// See https://pkg.go.dev/go.coder.com/cli#ParentCommand for more details.
func (r *root) Subcommands() []cli.Command {
    return []cli.Command{
        new(scanCmd),
    }
}

Great! Now that we’ve defined our root command, we can run it out of our main func.

port-scanner/cmd/port-scanner/main.go

package main

import "go.coder.com/cli"

func main() {
    cli.RunRoot(new(root))
}

Cool, lets compile and run our program and see what we have so far.

go build ./cmd/port-scanner && ./port-scanner

Usage: port-scanner [subcommand] [flags]

A great start no doubt but, we’ll need to add some subcommands if we wan’t our port-scanner to do anything.

Subcommand

I would normally implement the port-scanner under the internal package or the pkg package ( depending on whether or not I wanted to publish this as a public package for others to use ) but, to keep this short and sweet, lets implement the subcommand and the port scanner in the same file.

port-scanner/cmd/port-scanner/scan.go

package main

import (
    "context"
    "log"
    "net"
    "strconv"
    "sync"
    "time"

    "github.com/spf13/pflag"
    "go.coder.com/cli"
    "golang.org/x/xerrors"
)

const (
    wellKnownPorts = 1024
    allPorts       = 65535
)

var timeout = 3 * time.Second

// This time our command struct has a few fields, we can use these to store our flag values.
type scanCmd struct {
    host          string
    shouldScanAll bool
}

// cdr/cli supports subcommand aliases so lets define one in our command spec to provide
// our users with a more succinct input experience.
func (cmd *scanCmd) Spec() cli.CommandSpec {
    return cli.CommandSpec{
        Name:    "scan",
        Usage:   "[flags]",
        Aliases: []string{"s"},
        Desc:    "Scan a host for open ports.",
    }
}

// When adding flags, hang the following method-signature off your command struct to
// satisfy the FlaggedCommand interface as defined by cdr/cli.
// See https://pkg.go.dev/go.coder.com/cli#FlaggedCommand for more details.
func (cmd *scanCmd) RegisterFlags(fl *pflag.FlagSet) {
    fl.StringVar(&cmd.host, "host", "", "host to scan(ip address)")
    fl.BoolVarP(&cmd.shouldScanAll, "all", "a", false, "scan all ports(scans first 1024 if not enabled)")
}

func (cmd *scanCmd) Run(fl *pflag.FlagSet) {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    if cmd.host == "" {
        fl.Usage()
        log.Fatal("host not provided")
    }

    scanner, err := newScanner(cmd.host, cmd.shouldScanAll)
    if err != nil {
        fl.Usage()
        log.Fatalf("failed to initialize port scanner: %s", err)
    }

    log.Printf("scanning %s...", cmd.host)
    start := time.Now()
    openPorts := scanner.scan(ctx)
    log.Printf("scan completed in %s", time.Since(start))

    if len(openPorts) == 0 {
        log.Printf("%q has no open ports", cmd.host)
        return
    }
    log.Printf("found %d open ports", len(openPorts))
    log.Printf("open-ports: %v", openPorts)
}

// Now lets implement our port scanner.
type scanner struct {
    // We're going to wan't to scan each port concurrently so let's embed
    // a mutex lock into our scanner to make sure we do this in a thread-safe way.
    sync.Mutex
    host      string
    openPorts []int
    scanAll   bool
}

func newScanner(host string, scanAll bool) (*scanner, error) {
    if net.ParseIP(host) == nil {
        return nil, xerrors.Errorf("%q is an invalid ip address", host)
    }

    return &scanner{
        Mutex:   sync.Mutex{},
        host:    host,
        scanAll: scanAll,
    }, nil
}

func (s *scanner) add(port int) {
    // Since we'll be appending to the same slice from different goroutines,
    // lets make sure we're locking and unlocking between writes.
    s.Lock()
    s.openPorts = append(s.openPorts, port)
    s.Unlock()
}

func (s *scanner) scan(ctx context.Context) []int {
    // Lets use a wait group so we can wait for all of our
    // goroutines to exit before returning our results.
    var wg sync.WaitGroup
    for _, port := range portsToScan(s.scanAll) {
        wg.Add(1)
        // Because the value of the loop-variable 'port' changes on every iteration, we'll wan't
        // to pass a copy of its value into each new goroutine. We don't need to do this with the
        // 'host' variable because its not a loop-variable and it's value never changes.
        go func(p int) {
            defer wg.Done()
            if isOpen(s.host, p) {
                s.add(p)
            }
        }(port)
    }
    wg.Wait()
    return s.openPorts
}

func portsToScan(shouldScanAll bool) []int {
    max := wellKnownPorts
    if shouldScanAll {
        max = allPorts
    }

    var ports []int
    for port := 1; port < max; port++ {
        ports = append(ports, port)
    }
    return ports
}

func isOpen(host string, port int) bool {
    addr := net.JoinHostPort(host, strconv.Itoa(port))
    conn, err := net.DialTimeout("tcp", addr, timeout)
    if err != nil {
        return false
    }
    _ = conn.Close()
    return true
}

Great, now lets re-compile our program and run the root command to see what the output looks like.

Reviewing the command output

go build ./cmd/port-scanner && ./port-scanner

Usage: port-scanner [subcommand] [flags]

Description: A simple port-scanner.

Commands:
        s, scan  - Scan a host for open ports.

So far so good.

Let’s see what the help output for the scan subcommand looks like.

./port-scanner scan --help

Usage: port-scanner scan [flags]

Aliases: s

Description: Scan a host for open ports.

port-scanner scan flags:
-a, --all           scan all ports(scans first 1024 if not enabled)
    --host string   host to scan(ip address)

Beautiful. I think we’re ready to test this thing out.

Running the port scanner

Lets start by opening a new shell and firing up a local python web server.

python -m SimpleHTTPServer 8000 

Serving HTTP on 0.0.0.0 port 8000 ...

Let’s run our port-scanner in a different shell and while we’re at it, lets use the subcommand alias we defined for our scan command.

We’ll also wan’t to use our --all flag since our python server is running on port 8000(which is outside our default port-range).

./port-scanner s --host 127.0.0.1 --all
2021/03/18 21:52:49 scanning 127.0.0.1...
2021/03/18 21:52:50 scan completed in 886.199117ms
2021/03/18 21:52:50 found 10 open ports
2021/03/18 21:52:50 open-ports: [8000]

There is a discrepancy in the output because I’ve intentionally omitted certain ports for security reasons.

I hope you enjoyed this tutorial.

You can find this code on my github.

I encourage you to see what modifications you can make on your own.

More specifically, why don’t you try adding the capability to scan multiple hosts or an entire network.

That’s how simple it is to build command-line tools in go using the cdr/cli package.

Hope you found this useful.

Much love,

-Faris