whisper/client.go

package main

import (
	"bufio"
	"code.google.com/p/go.crypto/ssh/terminal"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"unicode"
)

type Auth struct {
	Nick string
	Key  rsa.PublicKey
}

func (a *Auth) Kind() string {
	return "auth"
}

type ReadWriter struct {
	io.Reader
	io.Writer
}

type Client struct {
	key    *rsa.PrivateKey
	host   string
	port   int
	nick   string
	conn   net.Conn
	done   chan interface{}
	mu     sync.Mutex
	prompt string
	line   []rune
	prev   *terminal.State
}

// establishes a connection to the server
func (c *Client) dial() error {
	addr := fmt.Sprintf("%s:%d", c.host, c.port)
	c.info("dialing %s", addr)
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return fmt.Errorf("client unable to connect: %v", err)
	}
	c.info("connected to %s", addr)
	c.conn = conn
	c.prompt = fmt.Sprintf("%s> ", addr)
	go c.handleMessages()
	return nil
}

// handles messages received from the current server
func (c *Client) handleMessages() {
	messages := make(chan Envelope)
	errors := make(chan error)
	done := make(chan interface{})
	go stream(c.conn, messages, errors, done)
	for {
		select {
		case message := <-messages:
			if err := c.handleMessage(message); err != nil {
				c.err("error handling message from server: %v", err)
			}
		case err := <-errors:
			c.err("server error: %v", err)
		case <-done:
			return
		}
	}
}

// handle a message received from the server
func (c *Client) handleMessage(m Envelope) error {
	switch m.Kind {
	case "meta":
		return c.handleMeta(m.Body)
	case "note":
		return c.handleNote(m.Body)
	default:
		return fmt.Errorf("received message of unsupported type: %v", m.Kind)
	}
}

// handles a meta message; that is, a message that is shown to the user
func (c *Client) handleMeta(body json.RawMessage) error {
	var meta Meta
	if err := json.Unmarshal(body, &meta); err != nil {
		return fmt.Errorf("unable to unmarshal meta message: %v", err)
	}
	c.info("message from server: %v", meta)
	return nil
}

func (c *Client) handleNote(body json.RawMessage) error {
	var ctxt []byte
	if err := json.Unmarshal(body, &ctxt); err != nil {
		return fmt.Errorf("unable to read note response: %v", err)
	}
	ptxt, err := rsa.DecryptPKCS1v15(rand.Reader, c.key, ctxt)
	if err != nil {
		return fmt.Errorf("unable to decrypt note response: %v", err)
	}
	var note NoteData
	if err := json.Unmarshal(ptxt, &note); err != nil {
		return fmt.Errorf("unable to unmarshal note response: %v", err)
	}
	c.info("title: %s", note.Title)
	c.info("body: %s", string(note.Body))
	return nil
}

func (c *Client) handshake() error {
	r := &Auth{Nick: c.nick, Key: c.key.PublicKey}
	c.info("authenticating as %s", c.nick)
	return c.sendRequest(r)
}

func (c *Client) sendRequest(r request) error {
	return writeRequest(c.conn, r)
}

func (c *Client) info(template string, args ...interface{}) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.trunc()
	fmt.Print("\033[90m# ")
	fmt.Printf(template, args...)
	if !strings.HasSuffix(template, "\n") {
		fmt.Print("\n")
	}
	fmt.Printf("\033[0m")
	c.renderLine()
}

func (c *Client) trunc() {
	fmt.Print("\033[1K") // clear to beginning of the line
	fmt.Print("\r")      // move to beginning of the line
}

func (c *Client) err(template string, args ...interface{}) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.trunc()
	fmt.Print("\033[31m# ")
	fmt.Printf(template, args...)
	if !strings.HasSuffix(template, "\n") {
		fmt.Print("\n")
	}
	fmt.Printf("\033[0m")
	c.renderLine()
}

func (c *Client) run() {
	go c.term()

	if err := c.dial(); err != nil {
		exit(1, "%v", err)
	}
	defer c.conn.Close()
	if err := c.handshake(); err != nil {
		exit(1, "%v", err)
	}
	<-c.done
	if c.prev != nil {
		terminal.Restore(0, c.prev)
	}
}

func (c *Client) renderLine() {
	fmt.Printf("\033[1K")                        // clear to beginning of current line
	fmt.Printf("\r")                             // move to beginning of current line
	fmt.Printf("%s%s", c.prompt, string(c.line)) // print the line with prompt
}

func (c *Client) control(r rune) {
	switch r {
	case 3: // ctrl+c
		c.eof()
	case 4: // EOF
		c.eof()
	case 12: // ctrl+l
		c.clear()
	case 13: // enter
		c.enter()
	case 21: // ctrl+u
		c.clearLine()
	case 27: // up
	case 127: // backspace
		c.backspace()
	default:
		c.info("undefined control sequence: %v %d %c", r, r, r)
	}
}

func (c *Client) enter() {
	fmt.Print("\n")
	line := string(c.line)
	c.line = make([]rune, 0, 32)
	c.exec(line)
}

func (c *Client) exec(line string) {
	parts := strings.Split(line, " ")
	if len(parts) == 0 || strings.TrimSpace(parts[0]) == "" {
		c.renderLine()
		return
	}
	switch parts[0] {
	case "notes/create":
		c.createNote(parts[1:])
	case "notes/get":
		c.getNote(parts[1:])
	default:
		c.err("unrecognized client command: %s", parts[0])
	}
}

func (c *Client) createNote(args []string) {
	if len(args) < 1 {
		c.err("yeah you need to specify a title.")
		return
	}
	title := strings.Join(args, " ")
	c.info("creating new note: %s", title)
	msg, err := c.readTextBlock()
	if err != nil {
		c.err("%v", err)
		return
	}
	note, err := c.encryptNote(title, msg)
	if err != nil {
		c.err("%v", err)
		return
	}
	if err := c.sendRequest(note); err != nil {
		c.err("error sending note: %v", err)
	}
}

func (c *Client) getNote(args []string) {
	if len(args) != 1 {
		c.err("ok notes/get takes exactly 1 argument")
		return
	}
	id, err := strconv.Atoi(args[0])
	if err != nil {
		c.err("that doesn't look like an int: %v", err)
		return
	}
	if err := c.sendRequest(GetNoteRequest(id)); err != nil {
		c.err("couldn't request note: %v", err)
		return
	}
}

func (c *Client) encryptNote(title string, note []rune) (NoteRequest, error) {
	obj := &NoteData{
		Title: title,
		Body:  []byte(string(note)), // lol, nooo, stahp
	}
	b, err := json.Marshal(obj)
	if err != nil {
		return nil, fmt.Errorf("unable to marshal note: %v", err)
	}
	ctxt, err := rsa.EncryptPKCS1v15(rand.Reader, &c.key.PublicKey, b)
	if err != nil {
		return nil, fmt.Errorf("unable to encrypt note: %v", err)
	}
	return ctxt, nil
}

func (c *Client) readTextBlock() ([]rune, error) {
	// god dammit what have i gotten myself into
	msg := make([]rune, 0, 400)
	fmt.Print("\033[1K") // clear to beginning of current line
	fmt.Print("\r")      // move to beginning of current line
	fmt.Print("\033[s")  // save the cursor position
	renderMsg := func() {
		fmt.Print("\033[u")           // restore cursor position
		fmt.Print("\033[0J")          // clear to screen end
		fmt.Printf("%s", string(msg)) // write message out
	}
	in := bufio.NewReader(os.Stdin)
	for {
		r, _, err := in.ReadRune()
		switch err {
		case io.EOF:
			return msg, nil
		case nil:
		default:
			return nil, fmt.Errorf("error reading textblock: %v", err)
		}
		if unicode.IsGraphic(r) {
			msg = append(msg, r)
			renderMsg()
			continue
		}
		switch r {
		case 13: // enter
			msg = append(msg, '\n')
			renderMsg()
		case 127: // backspace
			if len(msg) == 0 {
				break
			}
			msg = msg[:len(msg)-1]
			renderMsg()
		case 4: // ctrl+d
			return msg, nil
		}
	}
}

func (c *Client) eof() {
	fmt.Print("\033[1K") // clear to beginning of current line
	fmt.Print("\r")      // move to beginning of current line
	c.done <- 1
}

func (c *Client) clear() {
	fmt.Print("\033[2J")   // clear the screen
	fmt.Print("\033[0;0f") // move to 0, 0
	c.renderLine()
}

func (c *Client) clearLine() {
	c.line = make([]rune, 0, 32)
	c.renderLine()
}

func (c *Client) backspace() {
	if len(c.line) == 0 {
		return
	}
	c.line = c.line[:len(c.line)-1]
	c.renderLine()
}

func (c *Client) term() {
	old, err := terminal.MakeRaw(0)
	if err != nil {
		panic(err)
	}
	c.prev = old
	defer close(c.done)

	in := bufio.NewReader(os.Stdin)
	for {
		r, _, err := in.ReadRune()
		switch err {
		case io.EOF:
			return
		case nil:
		default:
			c.err("error reading rune: %v", err)
		}

		if unicode.IsGraphic(r) {
			c.line = append(c.line, r)
			c.renderLine()
		} else {
			c.control(r)
		}
	}
}

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,
		done: make(chan interface{}),
		line: make([]rune, 0, 32),
	}
	client.run()
}