Skip to content

Commit 2aba94f

Browse files
NHASNHAS
committed
Improve input parsing
1 parent 1712dc7 commit 2aba94f

File tree

16 files changed

+392
-196
lines changed

16 files changed

+392
-196
lines changed

‎internal/server/clients/clients.go

+5-5
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import (
1212

1313
var lock sync.RWMutex
1414
var clients = map[string]*ssh.ServerConn{}
15-
var AutoCompleteIdentifiers = trie.NewTrie()
15+
var Autocomplete = trie.NewTrie()
1616

1717
var uniqueIdToAllAliases = map[string][]string{}
1818
var aliases = map[string]map[string]bool{}
@@ -44,9 +44,9 @@ func Add(conn *ssh.ServerConn) (string, error) {
4444

4545
clients[idString] = conn
4646

47-
AutoCompleteIdentifiers.Add(idString)
47+
Autocomplete.Add(idString)
4848
for _, v := range uniqueIdToAllAliases[idString] {
49-
AutoCompleteIdentifiers.Add(v)
49+
Autocomplete.Add(v)
5050
}
5151

5252
return idString, nil
@@ -107,7 +107,7 @@ func Remove(uniqueId string) {
107107
panic("Somehow a unqiue ID is being removed without being in the set, this is a programming issue guy")
108108
}
109109

110-
AutoCompleteIdentifiers.Remove(uniqueId)
110+
Autocomplete.Remove(uniqueId)
111111
delete(clients, uniqueId)
112112

113113
if currentAliases, ok := uniqueIdToAllAliases[uniqueId]; ok {
@@ -116,7 +116,7 @@ func Remove(uniqueId string) {
116116
delete(aliases[alias], uniqueId)
117117

118118
if len(aliases[alias]) <= 1 {
119-
AutoCompleteIdentifiers.Remove(alias)
119+
Autocomplete.Remove(alias)
120120
delete(aliases, alias)
121121
}
122122
}

‎internal/server/commands/connect.go

+4-8
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ type connect struct {
1818
user *internal.User
1919
}
2020

21-
func (c *connect) Run(tty io.ReadWriter, args ...string) error {
21+
func (c *connect) Run(tty io.ReadWriter, line terminal.ParsedLine) error {
2222
if c.user.Pty == nil {
2323
return fmt.Errorf("Connect requires a pty")
2424
}
@@ -28,11 +28,11 @@ func (c *connect) Run(tty io.ReadWriter, args ...string) error {
2828
return fmt.Errorf("connect can only be called from the terminal, if you want to connect to your clients without connecting to the terminal use jumphost syntax -J")
2929
}
3030

31-
if len(args) != 1 {
31+
if len(line.Leftovers) != 1 {
3232
return fmt.Errorf(c.Help(false))
3333
}
3434

35-
target, err := clients.Get(args[0])
35+
target, err := clients.Get(line.Leftovers[0].Value())
3636
if err != nil {
3737
return err
3838
}
@@ -65,11 +65,7 @@ func (c *connect) Run(tty io.ReadWriter, args ...string) error {
6565

6666
}
6767

68-
func (c *connect) Expect(sections []string) []string {
69-
70-
if len(sections) == 1 {
71-
return []string{autocomplete.RemoteId}
72-
}
68+
func (c *connect) Expect(line terminal.ParsedLine) []string {
7369

7470
return nil
7571
}

‎internal/server/commands/exit.go

+4-2
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,18 @@ package commands
22

33
import (
44
"io"
5+
6+
"github.com/NHAS/reverse_ssh/internal/terminal"
57
)
68

79
type exit struct {
810
}
911

10-
func (e *exit) Run(tty io.ReadWriter, args ...string) error {
12+
func (e *exit) Run(tty io.ReadWriter, line terminal.ParsedLine) error {
1113
return io.EOF
1214
}
1315

14-
func (e *exit) Expect(sections []string) []string {
16+
func (e *exit) Expect(line terminal.ParsedLine) []string {
1517
return nil
1618
}
1719

‎internal/server/commands/help.go

+7-9
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,15 @@ import (
55
"io"
66
"sort"
77

8-
"github.com/NHAS/reverse_ssh/internal/terminal/autocomplete"
8+
"github.com/NHAS/reverse_ssh/internal/terminal"
99
"github.com/NHAS/reverse_ssh/pkg/table"
1010
)
1111

1212
type help struct {
1313
}
1414

15-
func (h *help) Run(tty io.ReadWriter, args ...string) error {
16-
if len(args) < 1 {
15+
func (h *help) Run(tty io.ReadWriter, line terminal.ParsedLine) error {
16+
if len(line.Leftovers) < 1 {
1717

1818
t, err := table.NewTable("Commands", "Function", "Purpose")
1919
if err != nil {
@@ -41,9 +41,9 @@ func (h *help) Run(tty io.ReadWriter, args ...string) error {
4141
return nil
4242
}
4343

44-
l, ok := allCommands[args[0]]
44+
l, ok := allCommands[line.Leftovers[0].Value()]
4545
if !ok {
46-
return fmt.Errorf("Command %s not found", args[0])
46+
return fmt.Errorf("Command %s not found", line.Leftovers[0].Value())
4747
}
4848

4949
fmt.Fprintf(tty, "\ndescription:\n%s\n", l.Help(true))
@@ -53,10 +53,8 @@ func (h *help) Run(tty io.ReadWriter, args ...string) error {
5353
return nil
5454
}
5555

56-
func (h *help) Expect(sections []string) []string {
57-
if len(sections) == 1 {
58-
return []string{autocomplete.Functions}
59-
}
56+
func (h *help) Expect(line terminal.ParsedLine) []string {
57+
6058
return nil
6159
}
6260

‎internal/server/commands/kill.go

+6-10
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,21 @@ import (
55
"io"
66

77
"github.com/NHAS/reverse_ssh/internal/server/clients"
8-
"github.com/NHAS/reverse_ssh/internal/terminal/autocomplete"
8+
"github.com/NHAS/reverse_ssh/internal/terminal"
99
"github.com/NHAS/reverse_ssh/pkg/logger"
1010
)
1111

1212
type kill struct {
1313
log logger.Logger
1414
}
1515

16-
func (k *kill) Run(tty io.ReadWriter, args ...string) error {
16+
func (k *kill) Run(tty io.ReadWriter, line terminal.ParsedLine) error {
1717

18-
if len(args) != 1 {
18+
if len(line.Leftovers) != 1 {
1919
return fmt.Errorf(k.Help(false))
2020
}
2121

22-
if args[0] == "all" {
22+
if line.Leftovers[0].Value() == "all" {
2323
killedClients := 0
2424
allClients := clients.GetAll()
2525
for _, v := range allClients {
@@ -29,7 +29,7 @@ func (k *kill) Run(tty io.ReadWriter, args ...string) error {
2929
return fmt.Errorf("%d connections killed", killedClients)
3030
}
3131

32-
conn, err := clients.Get(args[0])
32+
conn, err := clients.Get(line.Leftovers[0].Value())
3333
if err != nil {
3434
return err
3535
}
@@ -39,11 +39,7 @@ func (k *kill) Run(tty io.ReadWriter, args ...string) error {
3939
return err
4040
}
4141

42-
func (k *kill) Expect(sections []string) []string {
43-
44-
if len(sections) == 1 {
45-
return []string{autocomplete.RemoteId}
46-
}
42+
func (k *kill) Expect(line terminal.ParsedLine) []string {
4743

4844
return nil
4945
}

‎internal/server/commands/link.go

+34-25
Original file line numberDiff line numberDiff line change
@@ -10,23 +10,24 @@ import (
1010
"time"
1111

1212
"github.com/NHAS/reverse_ssh/internal/server/webserver"
13+
"github.com/NHAS/reverse_ssh/internal/terminal"
14+
"github.com/NHAS/reverse_ssh/internal/terminal/autocomplete"
1315
"github.com/NHAS/reverse_ssh/pkg/table"
1416
)
1517

1618
type link struct {
1719
}
1820

19-
func (l *link) Run(tty io.ReadWriter, args ...string) error {
20-
flags, _ := parseFlags(args...)
21+
func (l *link) Run(tty io.ReadWriter, line terminal.ParsedLine) error {
2122

22-
if isSet("h", flags) {
23+
if terminal.IsSet("h", line.Flags) {
2324
return errors.New(l.Help(false))
2425
}
2526

26-
if toList, ok := flags["l"]; ok {
27+
if toList, ok := line.Flags["l"]; ok {
2728
t, _ := table.NewTable("Active Files", "ID", "GOOS", "GOARCH", "Expires")
2829

29-
files, err := webserver.List(strings.Join(toList, " "))
30+
files, err := webserver.List(strings.Join(toList.ArgValues(), " "))
3031
if err != nil {
3132
return err
3233
}
@@ -54,8 +55,8 @@ func (l *link) Run(tty io.ReadWriter, args ...string) error {
5455

5556
}
5657

57-
if toRemove, ok := flags["r"]; ok {
58-
for _, id := range toRemove {
58+
if toRemove, ok := line.Flags["r"]; ok {
59+
for _, id := range toRemove.ArgValues() {
5960
err := webserver.Delete(id)
6061
if err != nil {
6162
fmt.Fprintf(tty, "Unable to remove %s: %s\n", id, err)
@@ -69,46 +70,46 @@ func (l *link) Run(tty io.ReadWriter, args ...string) error {
6970
}
7071

7172
var e time.Duration
72-
if lifetime, ok := flags["t"]; ok {
73-
if len(lifetime) != 1 {
74-
return fmt.Errorf("Time supplied %d arguments, expected 1", len(lifetime))
73+
if lifetime, ok := line.Flags["t"]; ok {
74+
if len(lifetime.Args) != 1 {
75+
return fmt.Errorf("Time supplied %d arguments, expected 1", len(lifetime.Args))
7576
}
7677

77-
mins, err := strconv.Atoi(lifetime[0])
78+
mins, err := strconv.Atoi(lifetime.Args[0].Value())
7879
if err != nil {
79-
return fmt.Errorf("Unable to parse number of minutes (-t): %s", lifetime[0])
80+
return fmt.Errorf("Unable to parse number of minutes (-t): %s", lifetime.Args[0].Value())
8081
}
8182

8283
e = time.Duration(mins) * time.Minute
8384
}
8485

8586
var homeserver_address string
86-
if cb, ok := flags["s"]; ok {
87-
if len(cb) != 1 {
88-
return fmt.Errorf("Homeserver connect back address supplied %d arguments, expected 1", len(cb))
87+
if cb, ok := line.Flags["s"]; ok {
88+
if len(cb.Args) != 1 {
89+
return fmt.Errorf("Homeserver connect back address supplied %d arguments, expected 1", len(cb.Args))
8990
}
9091

91-
homeserver_address = cb[0]
92+
homeserver_address = cb.Args[0].Value()
9293

9394
}
9495

9596
var goos string
96-
if cb, ok := flags["goos"]; ok {
97-
if len(cb) != 1 {
98-
return fmt.Errorf("GOOS supplied %d arguments, expected 1", len(cb))
97+
if cb, ok := line.Flags["goos"]; ok {
98+
if len(cb.Args) != 1 {
99+
return fmt.Errorf("GOOS supplied %d arguments, expected 1", len(cb.Args))
99100
}
100101

101-
goos = cb[0]
102+
goos = cb.Args[0].Value()
102103

103104
}
104105

105106
var goarch string
106-
if cb, ok := flags["goarch"]; ok {
107-
if len(cb) != 1 {
108-
return fmt.Errorf("GOARCH supplied %d arguments, expected 1", len(cb))
107+
if cb, ok := line.Flags["goarch"]; ok {
108+
if len(cb.Args) != 1 {
109+
return fmt.Errorf("GOARCH supplied %d arguments, expected 1", len(cb.Args))
109110
}
110111

111-
goarch = cb[0]
112+
goarch = cb.Args[0].Value()
112113

113114
}
114115

@@ -122,7 +123,15 @@ func (l *link) Run(tty io.ReadWriter, args ...string) error {
122123
return nil
123124
}
124125

125-
func (l *link) Expect(sections []string) []string {
126+
func (l *link) Expect(line terminal.ParsedLine) []string {
127+
if line.Section != nil {
128+
fmt.Println(line.Section)
129+
switch line.Section.Value() {
130+
case "l", "r":
131+
return []string{autocomplete.WebServerFileIds}
132+
}
133+
}
134+
126135
return nil
127136
}
128137

‎internal/server/commands/list.go

+11-22
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import (
88
"strings"
99

1010
"github.com/NHAS/reverse_ssh/internal/server/clients"
11-
"github.com/NHAS/reverse_ssh/internal/terminal/autocomplete"
11+
"github.com/NHAS/reverse_ssh/internal/terminal"
1212
"github.com/NHAS/reverse_ssh/pkg/table"
1313
"golang.org/x/crypto/ssh"
1414
)
@@ -31,23 +31,15 @@ func fancyTable(tty io.ReadWriter, applicable []displayItem) {
3131
t.Fprint(tty)
3232
}
3333

34-
func (l *List) Run(tty io.ReadWriter, args ...string) error {
34+
func (l *List) Run(tty io.ReadWriter, line terminal.ParsedLine) error {
3535

36-
flags, leftover := parseFlags(args...)
37-
filter := strings.Join(leftover, " ")
36+
filter := strings.Join(line.LeftoversStrings(), " ")
3837

39-
if isSet("h", flags) {
38+
if terminal.IsSet("h", line.Flags) {
4039
fmt.Fprintf(tty, "%s", l.Help(false))
4140
return nil
4241
}
4342

44-
// If we have a single option e.g -a, it can capture the filter so make sure we put it in the right place
45-
for _, c := range "tlnai" {
46-
if len(flags[string(c)]) > 0 {
47-
filter += strings.Join(flags[string(c)], " ")
48-
}
49-
}
50-
5143
_, err := filepath.Match(filter, "")
5244
if err != nil {
5345
return fmt.Errorf("Filter is not well formed")
@@ -92,35 +84,35 @@ func (l *List) Run(tty io.ReadWriter, args ...string) error {
9284
}
9385
}
9486

95-
if isSet("t", flags) {
87+
if terminal.IsSet("t", line.Flags) {
9688
fancyTable(tty, toReturn)
9789
return nil
9890
}
9991

10092
sep := ", "
101-
if isSet("l", flags) {
93+
if terminal.IsSet("l", line.Flags) {
10294
sep = "\n"
10395
}
10496

10597
for i, tr := range toReturn {
10698

107-
if !isSet("n", flags) && !isSet("i", flags) && !isSet("a", flags) {
99+
if !terminal.IsSet("n", line.Flags) && !terminal.IsSet("i", line.Flags) && !terminal.IsSet("a", line.Flags) {
108100
fmt.Fprint(tty, tr.id)
109101
if i != len(toReturn)-1 {
110102
fmt.Fprint(tty, sep)
111103
}
112104
continue
113105
}
114106

115-
if isSet("a", flags) {
107+
if terminal.IsSet("a", line.Flags) {
116108
fmt.Fprint(tty, tr.id)
117109
}
118110

119-
if isSet("n", flags) || isSet("a", flags) {
111+
if terminal.IsSet("n", line.Flags) || terminal.IsSet("a", line.Flags) {
120112
fmt.Fprint(tty, " "+tr.sc.User())
121113
}
122114

123-
if isSet("i", flags) || isSet("a", flags) {
115+
if terminal.IsSet("i", line.Flags) || terminal.IsSet("a", line.Flags) {
124116
fmt.Fprint(tty, " "+tr.sc.RemoteAddr().String())
125117
}
126118

@@ -134,10 +126,7 @@ func (l *List) Run(tty io.ReadWriter, args ...string) error {
134126
return nil
135127
}
136128

137-
func (l *List) Expect(sections []string) []string {
138-
if len(sections) == 1 {
139-
return []string{autocomplete.RemoteId}
140-
}
129+
func (l *List) Expect(line terminal.ParsedLine) []string {
141130

142131
return nil
143132
}

0 commit comments

Comments
 (0)