1// Copyright 2011 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package terminal
6
7import (
8	"io"
9	"sync"
10)
11
12// EscapeCodes contains escape sequences that can be written to the terminal in
13// order to achieve different styles of text.
14type EscapeCodes struct {
15	// Foreground colors
16	Black, Red, Green, Yellow, Blue, Magenta, Cyan, White []byte
17
18	// Reset all attributes
19	Reset []byte
20}
21
22var vt100EscapeCodes = EscapeCodes{
23	Black:   []byte{keyEscape, '[', '3', '0', 'm'},
24	Red:     []byte{keyEscape, '[', '3', '1', 'm'},
25	Green:   []byte{keyEscape, '[', '3', '2', 'm'},
26	Yellow:  []byte{keyEscape, '[', '3', '3', 'm'},
27	Blue:    []byte{keyEscape, '[', '3', '4', 'm'},
28	Magenta: []byte{keyEscape, '[', '3', '5', 'm'},
29	Cyan:    []byte{keyEscape, '[', '3', '6', 'm'},
30	White:   []byte{keyEscape, '[', '3', '7', 'm'},
31
32	Reset: []byte{keyEscape, '[', '0', 'm'},
33}
34
35// Terminal contains the state for running a VT100 terminal that is capable of
36// reading lines of input.
37type Terminal struct {
38	// AutoCompleteCallback, if non-null, is called for each keypress
39	// with the full input line and the current position of the cursor.
40	// If it returns a nil newLine, the key press is processed normally.
41	// Otherwise it returns a replacement line and the new cursor position.
42	AutoCompleteCallback func(line []byte, pos, key int) (newLine []byte, newPos int)
43
44	// Escape contains a pointer to the escape codes for this terminal.
45	// It's always a valid pointer, although the escape codes themselves
46	// may be empty if the terminal doesn't support them.
47	Escape *EscapeCodes
48
49	// lock protects the terminal and the state in this object from
50	// concurrent processing of a key press and a Write() call.
51	lock sync.Mutex
52
53	c      io.ReadWriter
54	prompt string
55
56	// line is the current line being entered.
57	line []byte
58	// pos is the logical position of the cursor in line
59	pos int
60	// echo is true if local echo is enabled
61	echo bool
62
63	// cursorX contains the current X value of the cursor where the left
64	// edge is 0. cursorY contains the row number where the first row of
65	// the current line is 0.
66	cursorX, cursorY int
67	// maxLine is the greatest value of cursorY so far.
68	maxLine int
69
70	termWidth, termHeight int
71
72	// outBuf contains the terminal data to be sent.
73	outBuf []byte
74	// remainder contains the remainder of any partial key sequences after
75	// a read. It aliases into inBuf.
76	remainder []byte
77	inBuf     [256]byte
78}
79
80// NewTerminal runs a VT100 terminal on the given ReadWriter. If the ReadWriter is
81// a local terminal, that terminal must first have been put into raw mode.
82// prompt is a string that is written at the start of each input line (i.e.
83// "> ").
84func NewTerminal(c io.ReadWriter, prompt string) *Terminal {
85	return &Terminal{
86		Escape:     &vt100EscapeCodes,
87		c:          c,
88		prompt:     prompt,
89		termWidth:  80,
90		termHeight: 24,
91		echo:       true,
92	}
93}
94
95const (
96	keyCtrlD     = 4
97	keyEnter     = '\r'
98	keyEscape    = 27
99	keyBackspace = 127
100	keyUnknown   = 256 + iota
101	keyUp
102	keyDown
103	keyLeft
104	keyRight
105	keyAltLeft
106	keyAltRight
107)
108
109// bytesToKey tries to parse a key sequence from b. If successful, it returns
110// the key and the remainder of the input. Otherwise it returns -1.
111func bytesToKey(b []byte) (int, []byte) {
112	if len(b) == 0 {
113		return -1, nil
114	}
115
116	if b[0] != keyEscape {
117		return int(b[0]), b[1:]
118	}
119
120	if len(b) >= 3 && b[0] == keyEscape && b[1] == '[' {
121		switch b[2] {
122		case 'A':
123			return keyUp, b[3:]
124		case 'B':
125			return keyDown, b[3:]
126		case 'C':
127			return keyRight, b[3:]
128		case 'D':
129			return keyLeft, b[3:]
130		}
131	}
132
133	if len(b) >= 6 && b[0] == keyEscape && b[1] == '[' && b[2] == '1' && b[3] == ';' && b[4] == '3' {
134		switch b[5] {
135		case 'C':
136			return keyAltRight, b[6:]
137		case 'D':
138			return keyAltLeft, b[6:]
139		}
140	}
141
142	// If we get here then we have a key that we don't recognise, or a
143	// partial sequence. It's not clear how one should find the end of a
144	// sequence without knowing them all, but it seems that [a-zA-Z] only
145	// appears at the end of a sequence.
146	for i, c := range b[0:] {
147		if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' {
148			return keyUnknown, b[i+1:]
149		}
150	}
151
152	return -1, b
153}
154
155// queue appends data to the end of t.outBuf
156func (t *Terminal) queue(data []byte) {
157	t.outBuf = append(t.outBuf, data...)
158}
159
160var eraseUnderCursor = []byte{' ', keyEscape, '[', 'D'}
161var space = []byte{' '}
162
163func isPrintable(key int) bool {
164	return key >= 32 && key < 127
165}
166
167// moveCursorToPos appends data to t.outBuf which will move the cursor to the
168// given, logical position in the text.
169func (t *Terminal) moveCursorToPos(pos int) {
170	if !t.echo {
171		return
172	}
173
174	x := len(t.prompt) + pos
175	y := x / t.termWidth
176	x = x % t.termWidth
177
178	up := 0
179	if y < t.cursorY {
180		up = t.cursorY - y
181	}
182
183	down := 0
184	if y > t.cursorY {
185		down = y - t.cursorY
186	}
187
188	left := 0
189	if x < t.cursorX {
190		left = t.cursorX - x
191	}
192
193	right := 0
194	if x > t.cursorX {
195		right = x - t.cursorX
196	}
197
198	t.cursorX = x
199	t.cursorY = y
200	t.move(up, down, left, right)
201}
202
203func (t *Terminal) move(up, down, left, right int) {
204	movement := make([]byte, 3*(up+down+left+right))
205	m := movement
206	for i := 0; i < up; i++ {
207		m[0] = keyEscape
208		m[1] = '['
209		m[2] = 'A'
210		m = m[3:]
211	}
212	for i := 0; i < down; i++ {
213		m[0] = keyEscape
214		m[1] = '['
215		m[2] = 'B'
216		m = m[3:]
217	}
218	for i := 0; i < left; i++ {
219		m[0] = keyEscape
220		m[1] = '['
221		m[2] = 'D'
222		m = m[3:]
223	}
224	for i := 0; i < right; i++ {
225		m[0] = keyEscape
226		m[1] = '['
227		m[2] = 'C'
228		m = m[3:]
229	}
230
231	t.queue(movement)
232}
233
234func (t *Terminal) clearLineToRight() {
235	op := []byte{keyEscape, '[', 'K'}
236	t.queue(op)
237}
238
239const maxLineLength = 4096
240
241// handleKey processes the given key and, optionally, returns a line of text
242// that the user has entered.
243func (t *Terminal) handleKey(key int) (line string, ok bool) {
244	switch key {
245	case keyBackspace:
246		if t.pos == 0 {
247			return
248		}
249		t.pos--
250		t.moveCursorToPos(t.pos)
251
252		copy(t.line[t.pos:], t.line[1+t.pos:])
253		t.line = t.line[:len(t.line)-1]
254		if t.echo {
255			t.writeLine(t.line[t.pos:])
256		}
257		t.queue(eraseUnderCursor)
258		t.moveCursorToPos(t.pos)
259	case keyAltLeft:
260		// move left by a word.
261		if t.pos == 0 {
262			return
263		}
264		t.pos--
265		for t.pos > 0 {
266			if t.line[t.pos] != ' ' {
267				break
268			}
269			t.pos--
270		}
271		for t.pos > 0 {
272			if t.line[t.pos] == ' ' {
273				t.pos++
274				break
275			}
276			t.pos--
277		}
278		t.moveCursorToPos(t.pos)
279	case keyAltRight:
280		// move right by a word.
281		for t.pos < len(t.line) {
282			if t.line[t.pos] == ' ' {
283				break
284			}
285			t.pos++
286		}
287		for t.pos < len(t.line) {
288			if t.line[t.pos] != ' ' {
289				break
290			}
291			t.pos++
292		}
293		t.moveCursorToPos(t.pos)
294	case keyLeft:
295		if t.pos == 0 {
296			return
297		}
298		t.pos--
299		t.moveCursorToPos(t.pos)
300	case keyRight:
301		if t.pos == len(t.line) {
302			return
303		}
304		t.pos++
305		t.moveCursorToPos(t.pos)
306	case keyEnter:
307		t.moveCursorToPos(len(t.line))
308		t.queue([]byte("\r\n"))
309		line = string(t.line)
310		ok = true
311		t.line = t.line[:0]
312		t.pos = 0
313		t.cursorX = 0
314		t.cursorY = 0
315		t.maxLine = 0
316	default:
317		if t.AutoCompleteCallback != nil {
318			t.lock.Unlock()
319			newLine, newPos := t.AutoCompleteCallback(t.line, t.pos, key)
320			t.lock.Lock()
321
322			if newLine != nil {
323				if t.echo {
324					t.moveCursorToPos(0)
325					t.writeLine(newLine)
326					for i := len(newLine); i < len(t.line); i++ {
327						t.writeLine(space)
328					}
329					t.moveCursorToPos(newPos)
330				}
331				t.line = newLine
332				t.pos = newPos
333				return
334			}
335		}
336		if !isPrintable(key) {
337			return
338		}
339		if len(t.line) == maxLineLength {
340			return
341		}
342		if len(t.line) == cap(t.line) {
343			newLine := make([]byte, len(t.line), 2*(1+len(t.line)))
344			copy(newLine, t.line)
345			t.line = newLine
346		}
347		t.line = t.line[:len(t.line)+1]
348		copy(t.line[t.pos+1:], t.line[t.pos:])
349		t.line[t.pos] = byte(key)
350		if t.echo {
351			t.writeLine(t.line[t.pos:])
352		}
353		t.pos++
354		t.moveCursorToPos(t.pos)
355	}
356	return
357}
358
359func (t *Terminal) writeLine(line []byte) {
360	for len(line) != 0 {
361		remainingOnLine := t.termWidth - t.cursorX
362		todo := len(line)
363		if todo > remainingOnLine {
364			todo = remainingOnLine
365		}
366		t.queue(line[:todo])
367		t.cursorX += todo
368		line = line[todo:]
369
370		if t.cursorX == t.termWidth {
371			t.cursorX = 0
372			t.cursorY++
373			if t.cursorY > t.maxLine {
374				t.maxLine = t.cursorY
375			}
376		}
377	}
378}
379
380func (t *Terminal) Write(buf []byte) (n int, err error) {
381	t.lock.Lock()
382	defer t.lock.Unlock()
383
384	if t.cursorX == 0 && t.cursorY == 0 {
385		// This is the easy case: there's nothing on the screen that we
386		// have to move out of the way.
387		return t.c.Write(buf)
388	}
389
390	// We have a prompt and possibly user input on the screen. We
391	// have to clear it first.
392	t.move(0 /* up */, 0 /* down */, t.cursorX /* left */, 0 /* right */)
393	t.cursorX = 0
394	t.clearLineToRight()
395
396	for t.cursorY > 0 {
397		t.move(1 /* up */, 0, 0, 0)
398		t.cursorY--
399		t.clearLineToRight()
400	}
401
402	if _, err = t.c.Write(t.outBuf); err != nil {
403		return
404	}
405	t.outBuf = t.outBuf[:0]
406
407	if n, err = t.c.Write(buf); err != nil {
408		return
409	}
410
411	t.queue([]byte(t.prompt))
412	chars := len(t.prompt)
413	if t.echo {
414		t.queue(t.line)
415		chars += len(t.line)
416	}
417	t.cursorX = chars % t.termWidth
418	t.cursorY = chars / t.termWidth
419	t.moveCursorToPos(t.pos)
420
421	if _, err = t.c.Write(t.outBuf); err != nil {
422		return
423	}
424	t.outBuf = t.outBuf[:0]
425	return
426}
427
428// ReadPassword temporarily changes the prompt and reads a password, without
429// echo, from the terminal.
430func (t *Terminal) ReadPassword(prompt string) (line string, err error) {
431	t.lock.Lock()
432	defer t.lock.Unlock()
433
434	oldPrompt := t.prompt
435	t.prompt = prompt
436	t.echo = false
437
438	line, err = t.readLine()
439
440	t.prompt = oldPrompt
441	t.echo = true
442
443	return
444}
445
446// ReadLine returns a line of input from the terminal.
447func (t *Terminal) ReadLine() (line string, err error) {
448	t.lock.Lock()
449	defer t.lock.Unlock()
450
451	return t.readLine()
452}
453
454func (t *Terminal) readLine() (line string, err error) {
455	// t.lock must be held at this point
456
457	if t.cursorX == 0 && t.cursorY == 0 {
458		t.writeLine([]byte(t.prompt))
459		t.c.Write(t.outBuf)
460		t.outBuf = t.outBuf[:0]
461	}
462
463	for {
464		rest := t.remainder
465		lineOk := false
466		for !lineOk {
467			var key int
468			key, rest = bytesToKey(rest)
469			if key < 0 {
470				break
471			}
472			if key == keyCtrlD {
473				return "", io.EOF
474			}
475			line, lineOk = t.handleKey(key)
476		}
477		if len(rest) > 0 {
478			n := copy(t.inBuf[:], rest)
479			t.remainder = t.inBuf[:n]
480		} else {
481			t.remainder = nil
482		}
483		t.c.Write(t.outBuf)
484		t.outBuf = t.outBuf[:0]
485		if lineOk {
486			return
487		}
488
489		// t.remainder is a slice at the beginning of t.inBuf
490		// containing a partial key sequence
491		readBuf := t.inBuf[len(t.remainder):]
492		var n int
493
494		t.lock.Unlock()
495		n, err = t.c.Read(readBuf)
496		t.lock.Lock()
497
498		if err != nil {
499			return
500		}
501
502		t.remainder = t.inBuf[:n+len(t.remainder)]
503	}
504	panic("unreachable")
505}
506
507// SetPrompt sets the prompt to be used when reading subsequent lines.
508func (t *Terminal) SetPrompt(prompt string) {
509	t.lock.Lock()
510	defer t.lock.Unlock()
511
512	t.prompt = prompt
513}
514
515func (t *Terminal) SetSize(width, height int) {
516	t.lock.Lock()
517	defer t.lock.Unlock()
518
519	t.termWidth, t.termHeight = width, height
520}
521