From 6680c7419efb77a196271eacfd939ac4a725758e Mon Sep 17 00:00:00 2001 From: Andrey Petrov Date: Thu, 11 Dec 2014 01:29:18 -0800 Subject: [PATCH] Add commands. --- client.go | 61 +++++++++++++++++++++++++++--------- server.go | 92 +++++++++++++++++++++++++++++++++++++++++++------------ 2 files changed, 119 insertions(+), 34 deletions(-) diff --git a/client.go b/client.go index 0f756fa..eceae0f 100644 --- a/client.go +++ b/client.go @@ -2,6 +2,7 @@ package main import ( "fmt" + "strings" "golang.org/x/crypto/ssh" "golang.org/x/crypto/ssh/terminal" @@ -9,6 +10,24 @@ import ( const MSG_BUFFER = 10 +const HELP_TEXT = `-> Available commands: + /about + /exit + /help + /list + /nick $NAME +` + +const ABOUT_TEXT = `-> ssh-chat is made by @shazow. + + It is a custom ssh server built in Go to serve a chat experience + instead of a shell. + + Source: https://github.com/shazow/ssh-chat + + For more, visit shazow.net or follow at twitter.com/shazow +` + type Client struct { Server *Server Conn *ssh.ServerConn @@ -19,15 +38,11 @@ type Client struct { termHeight int } -func NewClient(server *Server, conn *ssh.ServerConn, name string) *Client { - if name == "" { - name = "Anonymoose" - } - +func NewClient(server *Server, conn *ssh.ServerConn) *Client { return &Client{ Server: server, Conn: conn, - Name: name, + Name: conn.User(), Msg: make(chan string, MSG_BUFFER), } } @@ -42,10 +57,9 @@ func (c *Client) Resize(width int, height int) error { return nil } -func (c *Client) sendWelcome() { - msg := fmt.Sprintf("Welcome to ssh-chat. Enter /help for more.\r\n") - c.Msg <- msg - +func (c *Client) Rename(name string) { + c.Name = name + c.term.SetPrompt(fmt.Sprintf("[%s] ", name)) } func (c *Client) handleShell(channel ssh.Channel) { @@ -63,12 +77,31 @@ func (c *Client) handleShell(channel ssh.Channel) { break } - switch line { - case "/exit": - channel.Close() + parts := strings.SplitN(line, " ", 2) + isCmd := strings.HasPrefix(parts[0], "/") + + if isCmd { + switch parts[0] { + case "/exit": + channel.Close() + case "/help": + c.Msg <- HELP_TEXT + case "/about": + c.Msg <- ABOUT_TEXT + case "/nick": + if len(parts) == 2 { + c.Server.Rename(c, parts[1]) + } else { + c.Msg <- fmt.Sprintf("-> Missing $NAME from: /nick $NAME\r\n") + } + case "/list": + c.Msg <- fmt.Sprintf("-> Connected: %s\r\n", strings.Join(c.Server.List(nil), ",")) + default: + c.Msg <- fmt.Sprintf("-> Invalid command: %s\r\n", line) + } + continue } - //c.term.Write(c.term.Escape.Reset) msg := fmt.Sprintf("%s: %s\r\n", c.Name, line) c.Server.Broadcast(msg, c) } diff --git a/server.go b/server.go index ba702c6..0f1524c 100644 --- a/server.go +++ b/server.go @@ -9,12 +9,15 @@ import ( "golang.org/x/crypto/ssh" ) +type Clients map[string]*Client + type Server struct { sshConfig *ssh.ServerConfig sshSigner *ssh.Signer done chan struct{} - clients map[*Client]struct{} + clients Clients lock sync.Mutex + count int } func NewServer(privateKey []byte) (*Server, error) { @@ -39,7 +42,8 @@ func NewServer(privateKey []byte) (*Server, error) { sshConfig: &config, sshSigner: &signer, done: make(chan struct{}), - clients: map[*Client]struct{}{}, + clients: Clients{}, + count: 0, } return &server, nil @@ -47,7 +51,8 @@ func NewServer(privateKey []byte) (*Server, error) { func (s *Server) Broadcast(msg string, except *Client) { logger.Debugf("Broadcast to %d: %s", len(s.clients), strings.TrimRight(msg, "\r\n")) - for client := range s.clients { + + for _, client := range s.clients { if except != nil && client == except { continue } @@ -55,6 +60,65 @@ func (s *Server) Broadcast(msg string, except *Client) { } } +func (s *Server) Add(client *Client) { + client.Msg <- fmt.Sprintf("-> Welcome to ssh-chat. Enter /help for more.\r\n") + + s.lock.Lock() + s.count++ + + _, collision := s.clients[client.Name] + if collision { + newName := fmt.Sprintf("Guest%d", s.count) + client.Msg <- fmt.Sprintf("-> Your name '%s' was taken, renamed to '%s'. Use /nick to change it.\r\n", client.Name, newName) + client.Name = newName + } + + s.clients[client.Name] = client + num := len(s.clients) + s.lock.Unlock() + + s.Broadcast(fmt.Sprintf("* %s joined. (Total connected: %d)\r\n", client.Name, num), nil) +} + +func (s *Server) Remove(client *Client) { + s.lock.Lock() + delete(s.clients, client.Name) + s.lock.Unlock() + + s.Broadcast(fmt.Sprintf("* %s left.\r\n", client.Name), nil) +} + +func (s *Server) Rename(client *Client, newName string) { + s.lock.Lock() + + _, collision := s.clients[newName] + if collision { + client.Msg <- fmt.Sprintf("-> %s is not available.\r\n", newName) + s.lock.Unlock() + return + } + delete(s.clients, client.Name) + oldName := client.Name + client.Rename(newName) + s.clients[client.Name] = client + s.lock.Unlock() + + s.Broadcast(fmt.Sprintf("* %s is now known as %s.\r\n", oldName, newName), nil) +} + +func (s *Server) List(prefix *string) []string { + r := []string{} + + for name, _ := range s.clients { + if prefix != nil && !strings.HasPrefix(name, *prefix) { + continue + } + r = append(r, name) + } + + return r +} + func (s *Server) Start(laddr string) error { // Once a ServerConfig has been configured, connections can be // accepted. @@ -88,25 +152,13 @@ func (s *Server) Start(laddr string) error { go ssh.DiscardRequests(requests) - client := NewClient(s, sshConn, sshConn.User()) - // TODO: mutex this - - s.lock.Lock() - s.clients[client] = struct{}{} - num := len(s.clients) - s.lock.Unlock() - - client.sendWelcome() - - s.Broadcast(fmt.Sprintf("* %s joined. (Total connected: %d)\r\n", client.Name, num), nil) + client := NewClient(s, sshConn) + s.Add(client) go func() { + // Block until done, then remove. sshConn.Wait() - s.lock.Lock() - delete(s.clients, client) - s.lock.Unlock() - - s.Broadcast(fmt.Sprintf("* %s left.\r\n", client.Name), nil) + s.Remove(client) }() go client.handleChannels(channels) @@ -123,7 +175,7 @@ func (s *Server) Start(laddr string) error { } func (s *Server) Stop() { - for client := range s.clients { + for _, client := range s.clients { client.Conn.Close() }