1// The code in this file is adapted from github.com/mattn/go-colorable.
2
3package term
4
5import (
6	"bytes"
7	"fmt"
8	"io"
9	"strconv"
10	"strings"
11	"syscall"
12	"unsafe"
13)
14
15type colorWriter struct {
16	out     io.Writer
17	handle  syscall.Handle
18	lastbuf bytes.Buffer
19	oldattr word
20}
21
22// NewColorWriter returns an io.Writer that writes to w and provides cross
23// platform support for ANSI color codes. If w is not a terminal it is
24// returned unmodified.
25func NewColorWriter(w io.Writer) io.Writer {
26	if !IsConsole(w) {
27		return w
28	}
29
30	var csbi consoleScreenBufferInfo
31	handle := syscall.Handle(w.(fder).Fd())
32	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
33
34	return &colorWriter{
35		out:     w,
36		handle:  handle,
37		oldattr: csbi.attributes,
38	}
39}
40
41func (w *colorWriter) Write(data []byte) (n int, err error) {
42	var csbi consoleScreenBufferInfo
43	procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
44
45	er := bytes.NewBuffer(data)
46loop:
47	for {
48		r1, _, err := procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
49		if r1 == 0 {
50			break loop
51		}
52
53		c1, _, err := er.ReadRune()
54		if err != nil {
55			break loop
56		}
57		if c1 != 0x1b {
58			fmt.Fprint(w.out, string(c1))
59			continue
60		}
61		c2, _, err := er.ReadRune()
62		if err != nil {
63			w.lastbuf.WriteRune(c1)
64			break loop
65		}
66		if c2 != 0x5b {
67			w.lastbuf.WriteRune(c1)
68			w.lastbuf.WriteRune(c2)
69			continue
70		}
71
72		var buf bytes.Buffer
73		var m rune
74		for {
75			c, _, err := er.ReadRune()
76			if err != nil {
77				w.lastbuf.WriteRune(c1)
78				w.lastbuf.WriteRune(c2)
79				w.lastbuf.Write(buf.Bytes())
80				break loop
81			}
82			if ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z') || c == '@' {
83				m = c
84				break
85			}
86			buf.Write([]byte(string(c)))
87		}
88
89		switch m {
90		case 'm':
91			attr := csbi.attributes
92			cs := buf.String()
93			if cs == "" {
94				procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(w.oldattr))
95				continue
96			}
97			token := strings.Split(cs, ";")
98			intensityMode := word(0)
99			for _, ns := range token {
100				if n, err = strconv.Atoi(ns); err == nil {
101					switch {
102					case n == 0:
103						attr = w.oldattr
104					case n == 1:
105						attr |= intensityMode
106					case 30 <= n && n <= 37:
107						attr = (attr & backgroundMask)
108						if (n-30)&1 != 0 {
109							attr |= foregroundRed
110						}
111						if (n-30)&2 != 0 {
112							attr |= foregroundGreen
113						}
114						if (n-30)&4 != 0 {
115							attr |= foregroundBlue
116						}
117						intensityMode = foregroundIntensity
118					case n == 39: // reset foreground color
119						attr &= backgroundMask
120						attr |= w.oldattr & foregroundMask
121					case 40 <= n && n <= 47:
122						attr = (attr & foregroundMask)
123						if (n-40)&1 != 0 {
124							attr |= backgroundRed
125						}
126						if (n-40)&2 != 0 {
127							attr |= backgroundGreen
128						}
129						if (n-40)&4 != 0 {
130							attr |= backgroundBlue
131						}
132						intensityMode = backgroundIntensity
133					case n == 49: // reset background color
134						attr &= foregroundMask
135						attr |= w.oldattr & backgroundMask
136					}
137					procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(attr))
138				}
139			}
140		}
141	}
142	return len(data) - w.lastbuf.Len(), nil
143}
144
145var (
146	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
147	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
148)
149
150const (
151	foregroundBlue      = 0x1
152	foregroundGreen     = 0x2
153	foregroundRed       = 0x4
154	foregroundIntensity = 0x8
155	foregroundMask      = (foregroundRed | foregroundBlue | foregroundGreen | foregroundIntensity)
156	backgroundBlue      = 0x10
157	backgroundGreen     = 0x20
158	backgroundRed       = 0x40
159	backgroundIntensity = 0x80
160	backgroundMask      = (backgroundRed | backgroundBlue | backgroundGreen | backgroundIntensity)
161)
162
163type (
164	wchar uint16
165	short int16
166	dword uint32
167	word  uint16
168)
169
170type coord struct {
171	x short
172	y short
173}
174
175type smallRect struct {
176	left   short
177	top    short
178	right  short
179	bottom short
180}
181
182type consoleScreenBufferInfo struct {
183	size              coord
184	cursorPosition    coord
185	attributes        word
186	window            smallRect
187	maximumWindowSize coord
188}
189