include pastebin

.gitignore

@@ -16,3 +16,6 @@
 # vendor/
 id_ed25519
 id_ed25519.pub
+
+ussher
+pastedir/

cmd/ussher/main.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 
 	"github.com/barakmich/ussher"
+	"github.com/spf13/viper"
 )
 
 func main() {
@@ -17,12 +18,19 @@ func main() {
 		Keystore:    ussher.AcceptAllKeys(),
 		HTTPMux:     http.NewServeMux(),
 		SSHApps:     make(map[string]ussher.SSHApp),
+		BaseURL:     "http://localhost:8080/",
 	}
 	defer config.CloseApps()
 	err := ussher.RegisterEchoApp(config)
 	if err != nil {
 		fmt.Println(err)
 	}
+	v := viper.New()
+	v.Set("datadir", "./pastedir")
+	err = ussher.RegisterPastebin(config, v)
+	if err != nil {
+		fmt.Println(err)
+	}
 	go runHTTPServer(config)
 	err = ussher.RunSSHServer(config)
 	if err != nil {

config.go

@@ -15,6 +15,7 @@ type Config struct {
 	HTTPPort    int
 	SSHApps     map[string]SSHApp
 	HTTPMux     *http.ServeMux
+	BaseURL     string
 }
 
 func (c *Config) GetPrivateKey() (ssh.Signer, error) {

echo_app.go

@@ -12,9 +12,10 @@ type echoHandler struct {
 	count int
 }
 
-func (e *echoHandler) HandleExec(_ string, conn *ssh.ServerConn, channel ssh.Channel) {
+func (e *echoHandler) HandleExec(_ string, conn *ssh.ServerConn, channel ssh.Channel) error {
 	e.count += 1
-	io.Copy(channel, channel)
+	_, err := io.Copy(channel, channel)
+	return err
 }
 
 func (e *echoHandler) Close() error {

go.mod

@@ -3,6 +3,8 @@ module github.com/barakmich/ussher
 go 1.14
 
 require (
-	github.com/spf13/viper v1.7.1 // indirect
+	github.com/prometheus/common v0.4.0
+	github.com/spf13/viper v1.7.1
+	go.etcd.io/bbolt v1.3.2
 	golang.org/x/crypto v0.0.0-20200820211705-5c72a883971a
 )

go.sum

@@ -14,7 +14,9 @@ dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
+github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc h1:cAKDfWh5VpdgMhJosfJnn5/FoN2SRZ4p7fJNX58YPaU=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
+github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf h1:qet1QNfXsQxTZqLG4oE62mJzwPIB8+Tee4RNCL9ulrY=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=
 github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY=
@@ -131,6 +133,7 @@ github.com/prometheus/client_golang v0.9.3/go.mod h1:/TN21ttK/J9q6uSwhBd54HahCDf
 github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
 github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/prometheus/common v0.0.0-20181113130724-41aa239b4cce/go.mod h1:daVV7qP5qjZbuso7PdcryaAu0sAZbrN9i7WWcTMWvro=
+github.com/prometheus/common v0.4.0 h1:7etb9YClo3a6HjLzfl6rIQaU+FDfi0VSX39io3aQ+DM=
 github.com/prometheus/common v0.4.0/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
 github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
 github.com/prometheus/procfs v0.0.0-20190507164030-5867b95ac084/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
@@ -139,6 +142,7 @@ github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6So
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/ryanuber/columnize v0.0.0-20160712163229-9b3edd62028f/go.mod h1:sm1tb6uqfes/u+d4ooFouqFdy9/2g9QGwK3SQygK0Ts=
 github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc=
+github.com/sirupsen/logrus v1.2.0 h1:juTguoYk5qI21pwyTXY3B3Y5cOTH3ZUyZCg1v/mihuo=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
 github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1:OnSkiWE9lh6wB0YB77sQom3nweQdgAjqCqsofrRNTgc=
 github.com/smartystreets/goconvey v1.6.4/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA=
@@ -276,6 +280,7 @@ google.golang.org/genproto v0.0.0-20191108220845-16a3f7862a1a/go.mod h1:n3cpQtvx
 google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
 google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
 google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=
+gopkg.in/alecthomas/kingpin.v2 v2.2.6 h1:jMFz6MfLP0/4fUyZle81rXUoxOBFi19VUFKVDOQfozc=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

pastebin.go

@@ -1,23 +1,158 @@
-package pastebin
+package ussher
 
 import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/base64"
+	"errors"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net/http"
+	"net/url"
+	"os"
+	"path"
 	"path/filepath"
+	"strings"
+	"sync"
 
-	"github.com/barakmich/ussher"
+	"github.com/prometheus/common/log"
 	"github.com/spf13/viper"
 	bolt "go.etcd.io/bbolt"
 	"golang.org/x/crypto/ssh"
 )
 
+const DefaultBinLinkLen = 4
+
 type pastebin struct {
+	sync.Mutex
 	db      *bolt.DB
 	datadir string
+	config  *Config
 }
 
-func (p *pastebin) HandleExec(appstring string, conn *ssh.ServerConn, channel ssh.Channel) {
+func (p *pastebin) HandleExec(appstring string, conn *ssh.ServerConn, channel ssh.Channel) error {
+	commands := strings.Split(appstring, " ")
+	if len(commands) == 1 {
+		return p.handleNewFile(channel)
+	}
+	switch commands[1] {
+	case "create":
+		return p.handleNewFile(channel)
+	case "get":
+		return p.getFileSSH(commands[2:], channel)
+	}
+	return nil
+}
 
+func (p *pastebin) getFileSSH(commands []string, channel ssh.Channel) error {
+	if len(commands) == 0 {
+		return errors.New("Need to provide a shortlink string or URL of file to get")
+	}
+	key := commands[0]
+	if u, err := url.Parse(commands[0]); err == nil {
+		if strings.HasPrefix(u.Scheme, "http") {
+			id, err := p.getIDfromURL(u)
+			if err != nil {
+				return err
+			}
+			key = id
+		}
+	}
+	f, err := p.getFileFromID(key)
+	if err != nil {
+		return err
+	}
+	_, err = io.Copy(channel, f)
+	return err
+}
+
+func (p *pastebin) getFileFromID(id string) (*os.File, error) {
+	var shasum []byte
+	err := p.db.Update(func(tx *bolt.Tx) error {
+		bucket, err := tx.CreateBucketIfNotExists([]byte("pastebin"))
+		if err != nil {
+			return err
+		}
+		v := bucket.Get([]byte(id))
+		if v == nil {
+			return os.ErrNotExist
+		}
+		shasum = v
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	urlenc := base64.URLEncoding.EncodeToString(shasum)
+	filename := filepath.Join(p.datadir, urlenc)
+	return os.Open(filename)
+}
+
+func (p *pastebin) getIDfromURL(u *url.URL) (string, error) {
+	if u == nil {
+		panic("nil URL")
+	}
+	components := strings.Split(u.Path, "/")
+	if len(components) < 2 {
+		return "", errors.New("Incomplete URL")
+	}
+	return components[2], nil
+}
+
+func (p *pastebin) handleNewFile(channel ssh.Channel) error {
+	f, err := ioutil.TempFile(p.datadir, "pastebin_*")
+	if err != nil {
+		return err
+	}
+	tempname := f.Name()
+
+	hasher := sha256.New()
+	tee := io.TeeReader(channel, hasher)
+	_, err = io.Copy(f, tee)
+	if err != nil {
+		f.Close()
+		return err
+	}
+	f.Close()
+	shasum := hasher.Sum(nil)
+	urlenc := base64.URLEncoding.EncodeToString(shasum)
+
+	err = os.Rename(tempname, filepath.Join(p.datadir, urlenc))
+	if err != nil {
+		return err
+	}
+	shorturlLen := DefaultBinLinkLen
+	err = p.db.Update(func(tx *bolt.Tx) error {
+		bucket, err := tx.CreateBucketIfNotExists([]byte("pastebin"))
+		if err != nil {
+			return err
+		}
+		for {
+			v := bucket.Get([]byte(urlenc[:shorturlLen]))
+			if v == nil {
+				break
+			}
+			if bytes.Equal(v, shasum) {
+				// We already have the file! We just overwrote it.
+				// Return the same shortlink
+				break
+			}
+			shorturlLen += 1
+		}
+		return bucket.Put([]byte(urlenc[:shorturlLen]), shasum)
+	})
+	if err != nil {
+		return err
+	}
+	shortlink := path.Join(p.config.BaseURL, "p", urlenc[:shorturlLen])
+	_, err = fmt.Fprintln(channel, shortlink)
+	return err
+}
+
+func (p *pastebin) Close() error {
+	return p.db.Close()
 }
 
 func (p *pastebin) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -25,17 +160,31 @@ func (p *pastebin) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
-	id := req.URL.Path
+	id, err := p.getIDfromURL(req.URL)
-	fmt.Println("got id", id)
+	if err != nil {
+		w.WriteHeader(http.StatusBadRequest)
+		return
+	}
+	f, err := p.getFileFromID(id)
+	if err != nil {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	_, err = io.Copy(w, f)
+	if err != nil {
+		log.Errorln(err)
+	}
+	return
 }
 
-func RegisterPastebin(config *ussher.Config, configNS *viper.Viper) error {
+func RegisterPastebin(config *Config, configNS *viper.Viper) error {
 	path, err := filepath.Abs(configNS.GetString("datadir"))
 	if err != nil {
 		return err
 	}
 	p := &pastebin{
 		datadir: path,
+		config:  config,
 	}
 	db, err := bolt.Open(filepath.Join(p.datadir, "index.db"), 0600, nil)
 	if err != nil {
@@ -44,5 +193,6 @@ func RegisterPastebin(config *ussher.Config, configNS *viper.Viper) error {
 	p.db = db
 
 	config.HTTPMux.Handle("/p/", p)
+	config.SSHApps["pastebin"] = p
 	return nil
 }

ssh_app.go

@@ -3,6 +3,6 @@ package ussher
 import "golang.org/x/crypto/ssh"
 
 type SSHApp interface {
-	HandleExec(appstring string, conn *ssh.ServerConn, channel ssh.Channel)
+	HandleExec(appstring string, conn *ssh.ServerConn, channel ssh.Channel) error
 	Close() error
 }

ssh_server.go

@@ -86,16 +86,20 @@ func handleChannel(newChannel ssh.NewChannel, conn *ssh.ServerConn, config *Conf
 					exit = true
 				case "exec":
 					appstring := string(req.Payload[4:]) // skip the first four bytes, which are length of string
-					appspaced := strings.SplitN(appstring, " ", 1)
+					appspaced := strings.SplitN(appstring, " ", 2)
 					if len(appspaced) == 0 {
 						exit = true
 						break
 					}
-					v, ok := config.SSHApps[appspaced]
+					v, ok := config.SSHApps[appspaced[0]]
 					if ok {
-						v.HandleExec(appstring, conn, connection)
+						err := v.HandleExec(appstring, conn, connection)
+						if err != nil {
+							fmt.Fprintln(connection.Stderr(), err)
+							log.Printf("Got error while handling ssh session: %s", err)
+						}
 					} else {
-						log.Printf("Can't find app %s", appname)
+						log.Printf("Can't find app %s", appspaced[0])
 					}
 					exit = true
 				default: