1package terminal
2
3import (
4	"bytes"
5	"fmt"
6	"io"
7	"strconv"
8	"strings"
9	"syscall"
10	"unsafe"
11
12	"github.com/mattn/go-isatty"
13)
14
15var (
16	cursorFunctions = map[rune]func(c *Cursor) func(int){
17		'A': func(c *Cursor) func(int) { return c.Up },
18		'B': func(c *Cursor) func(int) { return c.Down },
19		'C': func(c *Cursor) func(int) { return c.Forward },
20		'D': func(c *Cursor) func(int) { return c.Back },
21		'E': func(c *Cursor) func(int) { return c.NextLine },
22		'F': func(c *Cursor) func(int) { return c.PreviousLine },
23		'G': func(c *Cursor) func(int) { return c.HorizontalAbsolute },
24	}
25)
26
27const (
28	foregroundBlue      = 0x1
29	foregroundGreen     = 0x2
30	foregroundRed       = 0x4
31	foregroundIntensity = 0x8
32	foregroundMask      = (foregroundRed | foregroundBlue | foregroundGreen | foregroundIntensity)
33	backgroundBlue      = 0x10
34	backgroundGreen     = 0x20
35	backgroundRed       = 0x40
36	backgroundIntensity = 0x80
37	backgroundMask      = (backgroundRed | backgroundBlue | backgroundGreen | backgroundIntensity)
38)
39
40type Writer struct {
41	out     FileWriter
42	handle  syscall.Handle
43	orgAttr word
44}
45
46func NewAnsiStdout(out FileWriter) io.Writer {
47	var csbi consoleScreenBufferInfo
48	if !isatty.IsTerminal(out.Fd()) {
49		return out
50	}
51	handle := syscall.Handle(out.Fd())
52	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
53	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
54}
55
56func NewAnsiStderr(out FileWriter) io.Writer {
57	var csbi consoleScreenBufferInfo
58	if !isatty.IsTerminal(out.Fd()) {
59		return out
60	}
61	handle := syscall.Handle(out.Fd())
62	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
63	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
64}
65
66func (w *Writer) Write(data []byte) (n int, err error) {
67	r := bytes.NewReader(data)
68
69	for {
70		ch, size, err := r.ReadRune()
71		if err != nil {
72			break
73		}
74		n += size
75
76		switch ch {
77		case '\x1b':
78			size, err = w.handleEscape(r)
79			n += size
80			if err != nil {
81				break
82			}
83		default:
84			fmt.Fprint(w.out, string(ch))
85		}
86	}
87	return
88}
89
90func (w *Writer) handleEscape(r *bytes.Reader) (n int, err error) {
91	buf := make([]byte, 0, 10)
92	buf = append(buf, "\x1b"...)
93
94	// Check '[' continues after \x1b
95	ch, size, err := r.ReadRune()
96	if err != nil {
97		fmt.Fprint(w.out, string(buf))
98		return
99	}
100	n += size
101	if ch != '[' {
102		fmt.Fprint(w.out, string(buf))
103		return
104	}
105
106	// Parse escape code
107	var code rune
108	argBuf := make([]byte, 0, 10)
109	for {
110		ch, size, err = r.ReadRune()
111		if err != nil {
112			fmt.Fprint(w.out, string(buf))
113			return
114		}
115		n += size
116		if ('a' <= ch && ch <= 'z') || ('A' <= ch && ch <= 'Z') {
117			code = ch
118			break
119		}
120		argBuf = append(argBuf, string(ch)...)
121	}
122
123	w.applyEscapeCode(buf, string(argBuf), code)
124	return
125}
126
127func (w *Writer) applyEscapeCode(buf []byte, arg string, code rune) {
128	c := &Cursor{Out: w.out}
129
130	switch arg + string(code) {
131	case "?25h":
132		c.Show()
133		return
134	case "?25l":
135		c.Hide()
136		return
137	}
138
139	if f, ok := cursorFunctions[code]; ok {
140		if n, err := strconv.Atoi(arg); err == nil {
141			f(c)(n)
142			return
143		}
144	}
145
146	switch code {
147	case 'm':
148		w.applySelectGraphicRendition(arg)
149	default:
150		buf = append(buf, string(code)...)
151		fmt.Fprint(w.out, string(buf))
152	}
153}
154
155// Original implementation: https://github.com/mattn/go-colorable
156func (w *Writer) applySelectGraphicRendition(arg string) {
157	if arg == "" {
158		procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(w.orgAttr))
159		return
160	}
161
162	var csbi consoleScreenBufferInfo
163	procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
164	attr := csbi.attributes
165
166	for _, param := range strings.Split(arg, ";") {
167		n, err := strconv.Atoi(param)
168		if err != nil {
169			continue
170		}
171
172		switch {
173		case n == 0 || n == 100:
174			attr = w.orgAttr
175		case 1 <= n && n <= 5:
176			attr |= foregroundIntensity
177		case 30 <= n && n <= 37:
178			attr = (attr & backgroundMask)
179			if (n-30)&1 != 0 {
180				attr |= foregroundRed
181			}
182			if (n-30)&2 != 0 {
183				attr |= foregroundGreen
184			}
185			if (n-30)&4 != 0 {
186				attr |= foregroundBlue
187			}
188		case 40 <= n && n <= 47:
189			attr = (attr & foregroundMask)
190			if (n-40)&1 != 0 {
191				attr |= backgroundRed
192			}
193			if (n-40)&2 != 0 {
194				attr |= backgroundGreen
195			}
196			if (n-40)&4 != 0 {
197				attr |= backgroundBlue
198			}
199		case 90 <= n && n <= 97:
200			attr = (attr & backgroundMask)
201			attr |= foregroundIntensity
202			if (n-90)&1 != 0 {
203				attr |= foregroundRed
204			}
205			if (n-90)&2 != 0 {
206				attr |= foregroundGreen
207			}
208			if (n-90)&4 != 0 {
209				attr |= foregroundBlue
210			}
211		case 100 <= n && n <= 107:
212			attr = (attr & foregroundMask)
213			attr |= backgroundIntensity
214			if (n-100)&1 != 0 {
215				attr |= backgroundRed
216			}
217			if (n-100)&2 != 0 {
218				attr |= backgroundGreen
219			}
220			if (n-100)&4 != 0 {
221				attr |= backgroundBlue
222			}
223		}
224	}
225
226	procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(attr))
227}
228