-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
5 changed files
with
182 additions
and
55 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,118 @@ | ||
| // Package main provides an example of how to use the promwish package. | ||
| // | ||
| // You can test with: | ||
| // | ||
| // go run main.go | ||
| // ssh -o UserKnownHostsFile=/dev/null -p 2222 localhost | ||
| // curl -s localhost:9222/metrics | grep wish_ | ||
| package main | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "os" | ||
| "os/signal" | ||
| "syscall" | ||
| "time" | ||
|
|
||
| tea "github.com/charmbracelet/bubbletea" | ||
| "github.com/charmbracelet/log" | ||
| "github.com/charmbracelet/promwish" | ||
| "github.com/charmbracelet/ssh" | ||
| "github.com/charmbracelet/wish" | ||
| bm "github.com/charmbracelet/wish/bubbletea" | ||
| "github.com/prometheus/client_golang/prometheus" | ||
| "github.com/prometheus/client_golang/prometheus/promauto" | ||
| "github.com/prometheus/client_golang/prometheus/promhttp" | ||
| ) | ||
|
|
||
| var keyTypeCounter = promauto.NewCounterVec(prometheus.CounterOpts{ | ||
| Name: "wish_auth_by_type_total", | ||
| Help: "Custom metric example", | ||
| }, []string{"type"}) | ||
|
|
||
| func main() { | ||
| registry := prometheus.NewRegistry() | ||
|
|
||
| s, err := wish.NewServer( | ||
| wish.WithAddress("localhost:2223"), | ||
| wish.WithPublicKeyAuth(func(ssh.Context, ssh.PublicKey) bool { return true }), | ||
| wish.WithMiddleware( | ||
| bm.Middleware(func(s ssh.Session) (tea.Model, []tea.ProgramOption) { | ||
| keyTypeCounter.WithLabelValues(s.PublicKey().Type()).Inc() | ||
| pty, _, active := s.Pty() | ||
| if !active { | ||
| fmt.Println("no active terminal, skipping") | ||
| return nil, nil | ||
| } | ||
| m := model{ | ||
| term: pty.Term, | ||
| user: s.User(), | ||
| } | ||
| return m, []tea.ProgramOption{} | ||
| }), | ||
| promwish.MiddlewareRegistry( | ||
| registry, | ||
| prometheus.Labels{ | ||
| "app": "test", | ||
| }, | ||
| promwish.DefaultCommandFn, | ||
| ), | ||
| ), | ||
| ) | ||
| if err != nil { | ||
| log.Fatal("Fail to start SSH server", "error", err) | ||
| } | ||
|
|
||
| log.Info("Starting server", "metrics", "http://localhost:9222/metrics", "ssh", "localhost:2223") | ||
| done := make(chan os.Signal, 1) | ||
| metrics := promwish.NewServer( | ||
| "localhost:9222", | ||
| promhttp.InstrumentMetricHandler( | ||
| registry, promhttp.HandlerFor(registry, promhttp.HandlerOpts{}), | ||
| ), | ||
| ) | ||
| signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM) | ||
| go func() { | ||
| if err = s.ListenAndServe(); err != nil { | ||
| log.Fatal("Fail to start HTTP server", "error", err) | ||
| } | ||
| }() | ||
| go func() { | ||
| if err = metrics.ListenAndServe(); err != nil { | ||
| log.Fatal("Fail to start metrics server", "error", err) | ||
| } | ||
| }() | ||
| <-done | ||
| if err := s.Close(); err != nil { | ||
| log.Fatal("Fail to close SSH server", "error", err) | ||
| } | ||
|
|
||
| ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) | ||
| defer func() { cancel() }() | ||
| if err := metrics.Shutdown(ctx); err != nil { | ||
| log.Fatal("Fail to shut down metrics server", "error", err) | ||
| } | ||
| } | ||
|
|
||
| // Just a generic tea.Model to demo terminal information of ssh. | ||
| type model struct { | ||
| user, term string | ||
| quitting bool | ||
| } | ||
|
|
||
| func (m model) Init() tea.Cmd { | ||
| return nil | ||
| } | ||
|
|
||
| func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { | ||
| if !m.quitting { | ||
| m.quitting = true | ||
| return m, func() tea.Msg { return nil } | ||
| } | ||
| return m, tea.Quit | ||
| } | ||
|
|
||
| func (m model) View() string { | ||
| return fmt.Sprintf("\n\nHello, %s. Your terminal is %s!\n\n\n", m.user, m.term) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.