1 files changed, 74 insertions, 0 deletions

diff --git a/main.go b/main.go
new file mode 100644
index 0000000..75555e6
--- /dev/null
+++ b/main.go
@@ -0,0 +1,74 @@
+package main
+
+import (
+        "errors"
+        "fmt"
+        "log"
+        "net/http"
+        "os"
+        "strings"
+
+        "github.com/spf13/cobra"
+)
+
+var rootCmd = &cobra.Command{
+        Use:     "websocket-proxy",
+        Version: "0.1.0",
+        Short:   "Proxy TCP connections over a websocket",
+}
+
+var clientCmd = &cobra.Command{
+        Use:   "client [server host]",
+        Short: "Act as a client for a websocket-proxy server",
+        Args: func(cmd *cobra.Command, args []string) error {
+                if len(args) != 1 || args[0] == "" {
+                        return errors.New("Server host is a required argument")
+                }
+                if !strings.HasPrefix(args[0], "ws://") && !strings.HasPrefix(args[0], "wss://") {
+                        return errors.New("Server host format is ws[s]://host[:port]/[path]")
+                }
+                return nil
+        },
+        Run: func(cmd *cobra.Command, args []string) {
+                listenOn := cmd.Flag("listen").Value.String()
+
+                h := &ClientHandler{
+                        SocketListenOn:  listenOn,
+                        WebsocketServer: args[0],
+                }
+
+                log.Printf("Serving on %s", listenOn)
+                h.Run()
+        },
+}
+
+var serverCmd = &cobra.Command{
+        Use:   "server [next-hop host]",
+        Short: "Serve websocket proxy client",
+        Args: func(cmd *cobra.Command, args []string) error {
+                if len(args) != 1 || args[0] == "" {
+                        return errors.New("Next-hop host is a required argument")
+                }
+                return nil
+        },
+        Run: func(cmd *cobra.Command, args []string) {
+                listenOn := cmd.Flag("listen").Value.String()
+                log.Printf("Serving on %s", listenOn)
+
+                http.Handle("/", NewServerHandler(args[0]))
+                log.Fatal(http.ListenAndServe(listenOn, nil))
+        },
+}
+
+func main() {
+        rootCmd.AddCommand(clientCmd)
+        rootCmd.AddCommand(serverCmd)
+
+        clientCmd.Flags().StringP("listen", "l", ":9013", "[address]:port to bind for serving clients")
+        serverCmd.Flags().StringP("listen", "l", ":9012", "[address]:port to bind for serving clients")
+
+        if err := rootCmd.Execute(); err != nil {
+                fmt.Println(err)
+                os.Exit(1)
+        }
+}