diff --git a/example/README.md b/cmd/README.md
similarity index 57%
rename from example/README.md
rename to cmd/README.md
index dff0f66..d0efa12 100644
--- a/example/README.md
+++ b/cmd/README.md
@@ -2,5 +2,5 @@
 
 Run
 ```bash
-go run example/server/main.go
+go run cmd/tlsdiag.go server
 ```
diff --git a/cmd/tlsdiag.go b/cmd/tlsdiag.go
new file mode 100644
index 0000000..73aef5f
--- /dev/null
+++ b/cmd/tlsdiag.go
@@ -0,0 +1,33 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/libp2p/go-libp2p-tls/cmd/cmdimpl"
+)
+
+func main() {
+	if len(os.Args) <= 1 {
+		fmt.Println("missing argument: client / server")
+		return
+	}
+
+	role := os.Args[1]
+	// remove the role argument from os.Args
+	os.Args = append([]string{os.Args[0]}, os.Args[2:]...)
+
+	var err error
+	switch role {
+	case "client":
+		err = cmdimpl.StartClient()
+	case "server":
+		err = cmdimpl.StartServer()
+	default:
+		fmt.Println("invalid argument. Expected client / server")
+		return
+	}
+	if err != nil {
+		panic(err)
+	}
+}
diff --git a/example/client/main.go b/cmd/tlsdiag/client.go
similarity index 61%
rename from example/client/main.go
rename to cmd/tlsdiag/client.go
index 678ece6..9de3347 100644
--- a/example/client/main.go
+++ b/cmd/tlsdiag/client.go
@@ -1,49 +1,24 @@
-package main
+package cmdimpl
 
 import (
 	"context"
-	"crypto/rand"
 	"flag"
 	"fmt"
 	"io/ioutil"
 	"net"
 	"time"
 
-	ic "github.com/libp2p/go-libp2p-crypto"
 	peer "github.com/libp2p/go-libp2p-peer"
 	libp2ptls "github.com/libp2p/go-libp2p-tls"
 )
 
-func main() {
-	if err := startClient(); err != nil {
-		panic(err)
-	}
-}
-
-func startClient() error {
+func StartClient() error {
 	port := flag.Int("p", 5533, "port")
 	peerIDString := flag.String("id", "", "peer ID")
 	keyType := flag.String("key", "ecdsa", "rsa, ecdsa, ed25519 or secp256k1")
 	flag.Parse()
 
-	var priv ic.PrivKey
-	var err error
-	switch *keyType {
-	case "rsa":
-		fmt.Printf("Generated new peer with an RSA key.")
-		priv, _, err = ic.GenerateRSAKeyPair(2048, rand.Reader)
-	case "ecdsa":
-		fmt.Printf("Generated new peer with an ECDSA key.")
-		priv, _, err = ic.GenerateECDSAKeyPair(rand.Reader)
-	case "ed25519":
-		fmt.Printf("Generated new peer with an Ed25519 key.")
-		priv, _, err = ic.GenerateEd25519Key(rand.Reader)
-	case "secp256k1":
-		fmt.Printf("Generated new peer with an Secp256k1 key.")
-		priv, _, err = ic.GenerateSecp256k1Key(rand.Reader)
-	default:
-		return fmt.Errorf("unknown key type: %s", *keyType)
-	}
+	priv, err := generateKey(*keyType)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/tlsdiag/key.go b/cmd/tlsdiag/key.go
new file mode 100644
index 0000000..fc2ad8a
--- /dev/null
+++ b/cmd/tlsdiag/key.go
@@ -0,0 +1,28 @@
+package cmdimpl
+
+import (
+	"crypto/rand"
+	"fmt"
+
+	ic "github.com/libp2p/go-libp2p-crypto"
+)
+
+func generateKey(keyType string) (priv ic.PrivKey, err error) {
+	switch keyType {
+	case "rsa":
+		fmt.Printf("Generated new peer with an RSA key.")
+		priv, _, err = ic.GenerateRSAKeyPair(2048, rand.Reader)
+	case "ecdsa":
+		fmt.Printf("Generated new peer with an ECDSA key.")
+		priv, _, err = ic.GenerateECDSAKeyPair(rand.Reader)
+	case "ed25519":
+		fmt.Printf("Generated new peer with an Ed25519 key.")
+		priv, _, err = ic.GenerateEd25519Key(rand.Reader)
+	case "secp256k1":
+		fmt.Printf("Generated new peer with an Secp256k1 key.")
+		priv, _, err = ic.GenerateSecp256k1Key(rand.Reader)
+	default:
+		return nil, fmt.Errorf("unknown key type: %s", keyType)
+	}
+	return
+}
diff --git a/example/server/main.go b/cmd/tlsdiag/server.go
similarity index 62%
rename from example/server/main.go
rename to cmd/tlsdiag/server.go
index 9988ec0..b192c44 100644
--- a/example/server/main.go
+++ b/cmd/tlsdiag/server.go
@@ -1,47 +1,22 @@
-package main
+package cmdimpl
 
 import (
 	"context"
-	"crypto/rand"
 	"flag"
 	"fmt"
 	"net"
 	"time"
 
-	ic "github.com/libp2p/go-libp2p-crypto"
 	peer "github.com/libp2p/go-libp2p-peer"
 	libp2ptls "github.com/libp2p/go-libp2p-tls"
 )
 
-func main() {
-	if err := startServer(); err != nil {
-		panic(err)
-	}
-}
-
-func startServer() error {
+func StartServer() error {
 	port := flag.Int("p", 5533, "port")
 	keyType := flag.String("key", "ecdsa", "rsa, ecdsa, ed25519 or secp256k1")
 	flag.Parse()
 
-	var priv ic.PrivKey
-	var err error
-	switch *keyType {
-	case "rsa":
-		fmt.Printf("Generated new peer with an RSA key.")
-		priv, _, err = ic.GenerateRSAKeyPair(2048, rand.Reader)
-	case "ecdsa":
-		fmt.Printf("Generated new peer with an ECDSA key.")
-		priv, _, err = ic.GenerateECDSAKeyPair(rand.Reader)
-	case "ed25519":
-		fmt.Printf("Generated new peer with an Ed25519 key.")
-		priv, _, err = ic.GenerateEd25519Key(rand.Reader)
-	case "secp256k1":
-		fmt.Printf("Generated new peer with an Secp256k1 key.")
-		priv, _, err = ic.GenerateSecp256k1Key(rand.Reader)
-	default:
-		return fmt.Errorf("unknown key type: %s", *keyType)
-	}
+	priv, err := generateKey(*keyType)
 	if err != nil {
 		return err
 	}
@@ -62,7 +37,7 @@ func startServer() error {
 	}
 	fmt.Printf("Listening for new connections on %s\n", ln.Addr())
 	fmt.Printf("Now run the following command in a separate terminal:\n")
-	fmt.Printf("\tgo run example/client/main.go -p %d -id %s\n", *port, id.Pretty())
+	fmt.Printf("\tgo run cmd/tlsdiag.go client -p %d -id %s\n", *port, id.Pretty())
 
 	for {
 		conn, err := ln.Accept()