Skip to content

Commit

Permalink
feat: support server banners (#210)
Browse files Browse the repository at this point in the history
* feat: support server banners

closes #205
needs charmbracelet/ssh#10

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

* test: banner

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

* chore: dep

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

* feat: banner handler

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

* go mod tidy

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

* fix: improve example

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

---------

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
  • Loading branch information
caarlos0 authored Jan 4, 2024
1 parent 8aefea8 commit ddb47df
Show file tree
Hide file tree
Showing 7 changed files with 142 additions and 6 deletions.
14 changes: 14 additions & 0 deletions examples/pwd-banner/banner.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@

Hello %s, welcome to a SSH server powered by
__ ___ _ _ _ _
\ \ / (_)___| |__ | | | |
\ \ /\ / /| / __| '_ \| | | |
\ V V / | \__ \ | | |_|_|_|
\_/\_/ |_|___/_| |_(_|_|_)


PS: The password is "asd123".
PPS: Visit https://charm.sh to learn more!

---

69 changes: 69 additions & 0 deletions examples/pwd-banner/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
package main

import (
"context"
"errors"
"fmt"
"os"
"os/signal"
"syscall"
"time"

_ "embed"

"github.com/charmbracelet/log"
"github.com/charmbracelet/ssh"
"github.com/charmbracelet/wish"
"github.com/charmbracelet/wish/logging"
)

const (
host = "localhost"
port = 23234
)

//go:embed banner.txt
var banner string

func main() {
s, err := wish.NewServer(
wish.WithAddress(fmt.Sprintf("%s:%d", host, port)),
wish.WithHostKeyPath(".ssh/term_info_ed25519"),
wish.WithBannerHandler(func(ctx ssh.Context) string {
return fmt.Sprintf(banner, ctx.User())
}),
wish.WithPasswordAuth(func(ctx ssh.Context, password string) bool {
return password == "asd123"
}),
wish.WithMiddleware(
func(h ssh.Handler) ssh.Handler {
return func(s ssh.Session) {
wish.Println(s, "Hello, world!")
h(s)
}
},
logging.Middleware(),
),
)
if err != nil {
log.Error("could not start server", "error", err)
}

done := make(chan os.Signal, 1)
signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
log.Info("Starting SSH server", "host", host, "port", port)
go func() {
if err = s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
log.Error("could not start server", "error", err)
done <- nil
}
}()

<-done
log.Info("Stopping SSH server")
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer func() { cancel() }()
if err := s.Shutdown(ctx); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
log.Error("could not stop server", "error", err)
}
}
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ require (
github.com/charmbracelet/keygen v0.5.0
github.com/charmbracelet/lipgloss v0.9.1
github.com/charmbracelet/log v0.3.1
github.com/charmbracelet/ssh v0.0.0-20221117183211-483d43d97103
github.com/charmbracelet/ssh v0.0.0-20240104172912-e11ae277b249
github.com/go-git/go-git/v5 v5.11.0
github.com/google/go-cmp v0.6.0
github.com/hashicorp/golang-lru/v2 v2.0.7
Expand Down
7 changes: 2 additions & 5 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@ github.com/charmbracelet/lipgloss v0.9.1 h1:PNyd3jvaJbg4jRHKWXnCj1akQm4rh8dbEzN1
github.com/charmbracelet/lipgloss v0.9.1/go.mod h1:1mPmG4cxScwUQALAAnacHaigiiHB9Pmr+v1VEawJl6I=
github.com/charmbracelet/log v0.3.1 h1:TjuY4OBNbxmHWSwO3tosgqs5I3biyY8sQPny/eCMTYw=
github.com/charmbracelet/log v0.3.1/go.mod h1:OR4E1hutLsax3ZKpXbgUqPtTjQfrh1pG3zwHGWuuq8g=
github.com/charmbracelet/ssh v0.0.0-20221117183211-483d43d97103 h1:wpHMERIN0pQZE635jWwT1dISgfjbpUcEma+fbPKSMCU=
github.com/charmbracelet/ssh v0.0.0-20221117183211-483d43d97103/go.mod h1:0Vm2/8yBljiLDnGJHU8ehswfawrEybGk33j5ssqKQVM=
github.com/charmbracelet/ssh v0.0.0-20240104172912-e11ae277b249 h1:M1Q/UIbi9Cfla0HK3A9puhQhb8ZPkA5HQxeSYcVrGpo=
github.com/charmbracelet/ssh v0.0.0-20240104172912-e11ae277b249/go.mod h1:A1H384KV/cJcSKofWjdSIb+dfbikXiW6449EluL3qJI=
github.com/cloudflare/circl v1.3.3 h1:fE/Qz0QdIGqeWfnwq0RE0R7MI51s0M2E4Ga9kq5AEMs=
github.com/cloudflare/circl v1.3.3/go.mod h1:5XYMA4rFBvNIrhs50XuiBJ15vF2pZn4nnUKZrLbUZFA=
github.com/containerd/console v1.0.4-0.20230313162750-1ae8d489ac81 h1:q2hJAaP1k2wIvVRd/hEHD7lacgqrCPS+k8g1MndzfWY=
Expand Down Expand Up @@ -103,7 +103,6 @@ github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5t
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
golang.org/x/crypto v0.0.0-20220826181053-bd7e27e6170d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
golang.org/x/crypto v0.3.1-0.20221117191849-2c476679df9a/go.mod h1:hebNnKkNXi2UzZN1eVRvBB7co0a+JxK6XbPiWVs/3J4=
golang.org/x/crypto v0.7.0/go.mod h1:pYwdfH91IfpZVANVyUOhSIPZaFoJGxTFbZhFTx+dXZU=
golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k=
Expand Down Expand Up @@ -137,7 +136,6 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220825204002-c680a09ffe64/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.3.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
Expand All @@ -147,7 +145,6 @@ golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc=
golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
golang.org/x/term v0.0.0-20220722155259-a9ba230a4035/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
golang.org/x/term v0.2.0/go.mod h1:TVmDHMZPmdnySmBfhjOoOdhjzdE1h4u1VwSiw2l1Nuc=
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U=
Expand Down
17 changes: 17 additions & 0 deletions options.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,23 @@ func WithVersion(version string) ssh.Option {
}
}

// WithBanner return an ssh.Option that sets the server banner.
func WithBanner(banner string) ssh.Option {
return func(s *ssh.Server) error {
s.Banner = banner
return nil
}
}

// WithBannerHandler return an ssh.Option that sets the server banner handler,
// overriding WithBanner.
func WithBannerHandler(h ssh.BannerHandler) ssh.Option {
return func(s *ssh.Server) error {
s.BannerHandler = h
return nil
}
}

// WithMiddleware composes the provided Middleware and returns an ssh.Option.
// This is useful if you manually create an ssh.Server and want to set the
// Server.Handler.
Expand Down
38 changes: 38 additions & 0 deletions options_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,44 @@ import (
gossh "golang.org/x/crypto/ssh"
)

func TestWithBanner(t *testing.T) {
const banner = "a banner"
var got string

srv := &ssh.Server{
Handler: func(s ssh.Session) {},
}
requireNoError(t, WithBanner(banner)(srv))

requireNoError(t, testsession.New(t, srv, &gossh.ClientConfig{
BannerCallback: func(message string) error {
got = message
return nil
},
}).Run(""))
requireEqual(t, banner, got)
}

func TestWithBannerHandler(t *testing.T) {
var got string

srv := &ssh.Server{
Handler: func(s ssh.Session) {},
}
requireNoError(t, WithBannerHandler(func(ctx ssh.Context) string {
return fmt.Sprintf("banner for %s", ctx.User())
})(srv))

requireNoError(t, testsession.New(t, srv, &gossh.ClientConfig{
User: "fulano",
BannerCallback: func(message string) error {
got = message
return nil
},
}).Run(""))
requireEqual(t, "banner for fulano", got)
}

func TestWithIdleTimeout(t *testing.T) {
s := ssh.Server{}
requireNoError(t, WithIdleTimeout(time.Second)(&s))
Expand Down
1 change: 1 addition & 0 deletions wish_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ func TestNewServerWithOptions(t *testing.T) {
if _, err := NewServer(
WithHostKeyPath(fp),
WithMaxTimeout(time.Second),
WithBanner("welcome"),
WithAddress(":2222"),
); err != nil {
t.Fatal(err)
Expand Down

0 comments on commit ddb47df

Please sign in to comment.