From 3fac9f5caffecda141cefb9239d91b0f0b0d0d69 Mon Sep 17 00:00:00 2001
From: Jordan Orelli <jordanorelli@gmail.com>
Date: Sun, 7 Dec 2014 10:07:21 -0500
Subject: [PATCH] cleanup

---
 client.go  | 72 +++++++++++++++++++++++++++++++++++++++---------------
 request.go | 12 ++++++---
 2 files changed, 60 insertions(+), 24 deletions(-)

diff --git a/client.go b/client.go
index 4e22767..e951c1f 100644
--- a/client.go
+++ b/client.go
@@ -3,7 +3,6 @@ package main
 import (
 	"code.google.com/p/go.crypto/ssh/terminal"
 	"crypto/rsa"
-	"encoding/json"
 	"fmt"
 	"io"
 	"net"
@@ -24,36 +23,50 @@ type ReadWriter struct {
 	io.Writer
 }
 
-func connect() {
-	if !terminal.IsTerminal(0) {
-		exit(1, "yeah you have to run this from a tty")
-	}
-	f, err := os.Open(options.key)
+type Client struct {
+	key  *rsa.PrivateKey
+	host string
+	port int
+	nick string
+	conn net.Conn
+}
+
+func (c *Client) dial() error {
+	conn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", c.host, c.port))
 	if err != nil {
-		exit(1, "unable to open private key file at %s: %v", options.key, err)
+		return fmt.Errorf("client unable to connect: %v", err)
 	}
-	defer f.Close()
+	c.conn = conn
+	return nil
+}
 
-	d1 := json.NewDecoder(f)
-	var key rsa.PrivateKey
-	if err := d1.Decode(&key); err != nil {
-		exit(1, "unable to decode key: %v", err)
-	}
+func (c *Client) handshake() error {
+	r := &Auth{Nick: c.nick, Key: c.key.PublicKey}
+	return c.sendRequest(r)
+}
 
-	conn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", options.host, options.port))
-	if err != nil {
-		exit(1, "unable to connect to server at %s:%d: %v", options.host, options.port, err)
+func (c *Client) sendRequest(r request) error {
+	return writeRequest(c.conn, r)
+}
+
+func (c *Client) run() {
+	if err := c.dial(); err != nil {
+		exit(1, "%v", err)
 	}
-	auth := Auth{
-		Nick: options.nick,
-		Key:  key.PublicKey,
+	defer c.conn.Close()
+	if err := c.handshake(); err != nil {
+		exit(1, "%v", err)
 	}
-	encodeRequest(conn, &auth)
+	c.term()
+}
+
+func (c *Client) term() {
 	old, err := terminal.MakeRaw(0)
 	if err != nil {
 		panic(err)
 	}
 	defer terminal.Restore(0, old)
+
 	r := &ReadWriter{Reader: os.Stdin, Writer: os.Stdout}
 	term := terminal.NewTerminal(r, "> ")
 
@@ -67,3 +80,22 @@ func connect() {
 		exit(1, "error on line read: %v", err)
 	}
 }
+
+func connect() {
+	if !terminal.IsTerminal(0) {
+		exit(1, "yeah, this only works from a TTY for now, sry.")
+	}
+
+	key, err := privateKey()
+	if err != nil {
+		exit(1, "unable to open private key file: %v", err)
+	}
+
+	client := &Client{
+		key:  key,
+		host: options.host,
+		port: options.port,
+		nick: options.nick,
+	}
+	client.run()
+}
diff --git a/request.go b/request.go
index f98b3b1..80c495f 100644
--- a/request.go
+++ b/request.go
@@ -3,6 +3,7 @@ package main
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"net"
 )
 
@@ -15,10 +16,10 @@ type request interface {
 	Kind() string
 }
 
-func encodeRequest(conn net.Conn, r request) {
+func writeRequest(w io.Writer, r request) error {
 	b, err := json.Marshal(r)
 	if err != nil {
-		exit(1, "unable to encode client request body: %v", err)
+		return fmt.Errorf("unable to marshal request: %v", err)
 	}
 	msg := json.RawMessage(b)
 	e := &Envelope{
@@ -27,9 +28,12 @@ func encodeRequest(conn net.Conn, r request) {
 	}
 	raw, err := json.Marshal(e)
 	if err != nil {
-		exit(1, "unable to encode client request: %v", err)
+		return fmt.Errorf("unable to marshal request envelope: %v", err)
 	}
-	conn.Write(raw)
+	if _, err := w.Write(raw); err != nil {
+		return fmt.Errorf("unable to write request: %v", err)
+	}
+	return nil
 }
 
 func decodeRequest(conn net.Conn) (request, error) {