1// Copyright 2014 shiena Authors. All rights reserved.
2// Use of this source code is governed by a MIT-style
3// license that can be found in the LICENSE file.
4
5// +build windows
6
7package ansicolor
8
9import (
10	"bytes"
11	"io"
12	"strings"
13	"syscall"
14	"unsafe"
15)
16
17type csiState int
18
19const (
20	outsideCsiCode csiState = iota
21	firstCsiCode
22	secondCsiCode
23)
24
25type parseResult int
26
27const (
28	noConsole parseResult = iota
29	changedColor
30	unknown
31)
32
33type ansiColorWriter struct {
34	w             io.Writer
35	mode          outputMode
36	state         csiState
37	paramStartBuf bytes.Buffer
38	paramBuf      bytes.Buffer
39}
40
41const (
42	firstCsiChar   byte = '\x1b'
43	secondeCsiChar byte = '['
44	separatorChar  byte = ';'
45	sgrCode        byte = 'm'
46)
47
48const (
49	foregroundBlue      = uint16(0x0001)
50	foregroundGreen     = uint16(0x0002)
51	foregroundRed       = uint16(0x0004)
52	foregroundIntensity = uint16(0x0008)
53	backgroundBlue      = uint16(0x0010)
54	backgroundGreen     = uint16(0x0020)
55	backgroundRed       = uint16(0x0040)
56	backgroundIntensity = uint16(0x0080)
57	underscore          = uint16(0x8000)
58
59	foregroundMask = foregroundBlue | foregroundGreen | foregroundRed | foregroundIntensity
60	backgroundMask = backgroundBlue | backgroundGreen | backgroundRed | backgroundIntensity
61)
62
63const (
64	ansiReset        = "0"
65	ansiIntensityOn  = "1"
66	ansiIntensityOff = "21"
67	ansiUnderlineOn  = "4"
68	ansiUnderlineOff = "24"
69	ansiBlinkOn      = "5"
70	ansiBlinkOff     = "25"
71
72	ansiForegroundBlack   = "30"
73	ansiForegroundRed     = "31"
74	ansiForegroundGreen   = "32"
75	ansiForegroundYellow  = "33"
76	ansiForegroundBlue    = "34"
77	ansiForegroundMagenta = "35"
78	ansiForegroundCyan    = "36"
79	ansiForegroundWhite   = "37"
80	ansiForegroundDefault = "39"
81
82	ansiBackgroundBlack   = "40"
83	ansiBackgroundRed     = "41"
84	ansiBackgroundGreen   = "42"
85	ansiBackgroundYellow  = "43"
86	ansiBackgroundBlue    = "44"
87	ansiBackgroundMagenta = "45"
88	ansiBackgroundCyan    = "46"
89	ansiBackgroundWhite   = "47"
90	ansiBackgroundDefault = "49"
91
92	ansiLightForegroundGray    = "90"
93	ansiLightForegroundRed     = "91"
94	ansiLightForegroundGreen   = "92"
95	ansiLightForegroundYellow  = "93"
96	ansiLightForegroundBlue    = "94"
97	ansiLightForegroundMagenta = "95"
98	ansiLightForegroundCyan    = "96"
99	ansiLightForegroundWhite   = "97"
100
101	ansiLightBackgroundGray    = "100"
102	ansiLightBackgroundRed     = "101"
103	ansiLightBackgroundGreen   = "102"
104	ansiLightBackgroundYellow  = "103"
105	ansiLightBackgroundBlue    = "104"
106	ansiLightBackgroundMagenta = "105"
107	ansiLightBackgroundCyan    = "106"
108	ansiLightBackgroundWhite   = "107"
109)
110
111type drawType int
112
113const (
114	foreground drawType = iota
115	background
116)
117
118type winColor struct {
119	code     uint16
120	drawType drawType
121}
122
123var colorMap = map[string]winColor{
124	ansiForegroundBlack:   {0, foreground},
125	ansiForegroundRed:     {foregroundRed, foreground},
126	ansiForegroundGreen:   {foregroundGreen, foreground},
127	ansiForegroundYellow:  {foregroundRed | foregroundGreen, foreground},
128	ansiForegroundBlue:    {foregroundBlue, foreground},
129	ansiForegroundMagenta: {foregroundRed | foregroundBlue, foreground},
130	ansiForegroundCyan:    {foregroundGreen | foregroundBlue, foreground},
131	ansiForegroundWhite:   {foregroundRed | foregroundGreen | foregroundBlue, foreground},
132	ansiForegroundDefault: {foregroundRed | foregroundGreen | foregroundBlue, foreground},
133
134	ansiBackgroundBlack:   {0, background},
135	ansiBackgroundRed:     {backgroundRed, background},
136	ansiBackgroundGreen:   {backgroundGreen, background},
137	ansiBackgroundYellow:  {backgroundRed | backgroundGreen, background},
138	ansiBackgroundBlue:    {backgroundBlue, background},
139	ansiBackgroundMagenta: {backgroundRed | backgroundBlue, background},
140	ansiBackgroundCyan:    {backgroundGreen | backgroundBlue, background},
141	ansiBackgroundWhite:   {backgroundRed | backgroundGreen | backgroundBlue, background},
142	ansiBackgroundDefault: {0, background},
143
144	ansiLightForegroundGray:    {foregroundIntensity, foreground},
145	ansiLightForegroundRed:     {foregroundIntensity | foregroundRed, foreground},
146	ansiLightForegroundGreen:   {foregroundIntensity | foregroundGreen, foreground},
147	ansiLightForegroundYellow:  {foregroundIntensity | foregroundRed | foregroundGreen, foreground},
148	ansiLightForegroundBlue:    {foregroundIntensity | foregroundBlue, foreground},
149	ansiLightForegroundMagenta: {foregroundIntensity | foregroundRed | foregroundBlue, foreground},
150	ansiLightForegroundCyan:    {foregroundIntensity | foregroundGreen | foregroundBlue, foreground},
151	ansiLightForegroundWhite:   {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},
152
153	ansiLightBackgroundGray:    {backgroundIntensity, background},
154	ansiLightBackgroundRed:     {backgroundIntensity | backgroundRed, background},
155	ansiLightBackgroundGreen:   {backgroundIntensity | backgroundGreen, background},
156	ansiLightBackgroundYellow:  {backgroundIntensity | backgroundRed | backgroundGreen, background},
157	ansiLightBackgroundBlue:    {backgroundIntensity | backgroundBlue, background},
158	ansiLightBackgroundMagenta: {backgroundIntensity | backgroundRed | backgroundBlue, background},
159	ansiLightBackgroundCyan:    {backgroundIntensity | backgroundGreen | backgroundBlue, background},
160	ansiLightBackgroundWhite:   {backgroundIntensity | backgroundRed | backgroundGreen | backgroundBlue, background},
161}
162
163var (
164	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
165	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
166	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
167	defaultAttr                    *textAttributes
168)
169
170func init() {
171	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
172	if screenInfo != nil {
173		colorMap[ansiForegroundDefault] = winColor{
174			screenInfo.WAttributes & (foregroundRed | foregroundGreen | foregroundBlue),
175			foreground,
176		}
177		colorMap[ansiBackgroundDefault] = winColor{
178			screenInfo.WAttributes & (backgroundRed | backgroundGreen | backgroundBlue),
179			background,
180		}
181		defaultAttr = convertTextAttr(screenInfo.WAttributes)
182	}
183}
184
185type coord struct {
186	X, Y int16
187}
188
189type smallRect struct {
190	Left, Top, Right, Bottom int16
191}
192
193type consoleScreenBufferInfo struct {
194	DwSize              coord
195	DwCursorPosition    coord
196	WAttributes         uint16
197	SrWindow            smallRect
198	DwMaximumWindowSize coord
199}
200
201func getConsoleScreenBufferInfo(hConsoleOutput uintptr) *consoleScreenBufferInfo {
202	var csbi consoleScreenBufferInfo
203	ret, _, _ := procGetConsoleScreenBufferInfo.Call(
204		hConsoleOutput,
205		uintptr(unsafe.Pointer(&csbi)))
206	if ret == 0 {
207		return nil
208	}
209	return &csbi
210}
211
212func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
213	ret, _, _ := procSetConsoleTextAttribute.Call(
214		hConsoleOutput,
215		uintptr(wAttributes))
216	return ret != 0
217}
218
219type textAttributes struct {
220	foregroundColor     uint16
221	backgroundColor     uint16
222	foregroundIntensity uint16
223	backgroundIntensity uint16
224	underscore          uint16
225	otherAttributes     uint16
226}
227
228func convertTextAttr(winAttr uint16) *textAttributes {
229	fgColor := winAttr & (foregroundRed | foregroundGreen | foregroundBlue)
230	bgColor := winAttr & (backgroundRed | backgroundGreen | backgroundBlue)
231	fgIntensity := winAttr & foregroundIntensity
232	bgIntensity := winAttr & backgroundIntensity
233	underline := winAttr & underscore
234	otherAttributes := winAttr &^ (foregroundMask | backgroundMask | underscore)
235	return &textAttributes{fgColor, bgColor, fgIntensity, bgIntensity, underline, otherAttributes}
236}
237
238func convertWinAttr(textAttr *textAttributes) uint16 {
239	var winAttr uint16
240	winAttr |= textAttr.foregroundColor
241	winAttr |= textAttr.backgroundColor
242	winAttr |= textAttr.foregroundIntensity
243	winAttr |= textAttr.backgroundIntensity
244	winAttr |= textAttr.underscore
245	winAttr |= textAttr.otherAttributes
246	return winAttr
247}
248
249func changeColor(param []byte) parseResult {
250	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
251	if screenInfo == nil {
252		return noConsole
253	}
254
255	winAttr := convertTextAttr(screenInfo.WAttributes)
256	strParam := string(param)
257	if len(strParam) <= 0 {
258		strParam = "0"
259	}
260	csiParam := strings.Split(strParam, string(separatorChar))
261	for _, p := range csiParam {
262		c, ok := colorMap[p]
263		switch {
264		case !ok:
265			switch p {
266			case ansiReset:
267				winAttr.foregroundColor = defaultAttr.foregroundColor
268				winAttr.backgroundColor = defaultAttr.backgroundColor
269				winAttr.foregroundIntensity = defaultAttr.foregroundIntensity
270				winAttr.backgroundIntensity = defaultAttr.backgroundIntensity
271				winAttr.underscore = 0
272				winAttr.otherAttributes = 0
273			case ansiIntensityOn:
274				winAttr.foregroundIntensity = foregroundIntensity
275			case ansiIntensityOff:
276				winAttr.foregroundIntensity = 0
277			case ansiUnderlineOn:
278				winAttr.underscore = underscore
279			case ansiUnderlineOff:
280				winAttr.underscore = 0
281			case ansiBlinkOn:
282				winAttr.backgroundIntensity = backgroundIntensity
283			case ansiBlinkOff:
284				winAttr.backgroundIntensity = 0
285			default:
286				// unknown code
287			}
288		case c.drawType == foreground:
289			winAttr.foregroundColor = c.code
290		case c.drawType == background:
291			winAttr.backgroundColor = c.code
292		}
293	}
294	winTextAttribute := convertWinAttr(winAttr)
295	setConsoleTextAttribute(uintptr(syscall.Stdout), winTextAttribute)
296
297	return changedColor
298}
299
300func parseEscapeSequence(command byte, param []byte) parseResult {
301	if defaultAttr == nil {
302		return noConsole
303	}
304
305	switch command {
306	case sgrCode:
307		return changeColor(param)
308	default:
309		return unknown
310	}
311}
312
313func (cw *ansiColorWriter) flushBuffer() (int, error) {
314	return cw.flushTo(cw.w)
315}
316
317func (cw *ansiColorWriter) resetBuffer() (int, error) {
318	return cw.flushTo(nil)
319}
320
321func (cw *ansiColorWriter) flushTo(w io.Writer) (int, error) {
322	var n1, n2 int
323	var err error
324
325	startBytes := cw.paramStartBuf.Bytes()
326	cw.paramStartBuf.Reset()
327	if w != nil {
328		n1, err = cw.w.Write(startBytes)
329		if err != nil {
330			return n1, err
331		}
332	} else {
333		n1 = len(startBytes)
334	}
335	paramBytes := cw.paramBuf.Bytes()
336	cw.paramBuf.Reset()
337	if w != nil {
338		n2, err = cw.w.Write(paramBytes)
339		if err != nil {
340			return n1 + n2, err
341		}
342	} else {
343		n2 = len(paramBytes)
344	}
345	return n1 + n2, nil
346}
347
348func isParameterChar(b byte) bool {
349	return ('0' <= b && b <= '9') || b == separatorChar
350}
351
352func (cw *ansiColorWriter) Write(p []byte) (int, error) {
353	r, nw, first, last := 0, 0, 0, 0
354	if cw.mode != DiscardNonColorEscSeq {
355		cw.state = outsideCsiCode
356		cw.resetBuffer()
357	}
358
359	var err error
360	for i, ch := range p {
361		switch cw.state {
362		case outsideCsiCode:
363			if ch == firstCsiChar {
364				cw.paramStartBuf.WriteByte(ch)
365				cw.state = firstCsiCode
366			}
367		case firstCsiCode:
368			switch ch {
369			case firstCsiChar:
370				cw.paramStartBuf.WriteByte(ch)
371				break
372			case secondeCsiChar:
373				cw.paramStartBuf.WriteByte(ch)
374				cw.state = secondCsiCode
375				last = i - 1
376			default:
377				cw.resetBuffer()
378				cw.state = outsideCsiCode
379			}
380		case secondCsiCode:
381			if isParameterChar(ch) {
382				cw.paramBuf.WriteByte(ch)
383			} else {
384				nw, err = cw.w.Write(p[first:last])
385				r += nw
386				if err != nil {
387					return r, err
388				}
389				first = i + 1
390				result := parseEscapeSequence(ch, cw.paramBuf.Bytes())
391				if result == noConsole || (cw.mode == OutputNonColorEscSeq && result == unknown) {
392					cw.paramBuf.WriteByte(ch)
393					nw, err := cw.flushBuffer()
394					if err != nil {
395						return r, err
396					}
397					r += nw
398				} else {
399					n, _ := cw.resetBuffer()
400					// Add one more to the size of the buffer for the last ch
401					r += n + 1
402				}
403
404				cw.state = outsideCsiCode
405			}
406		default:
407			cw.state = outsideCsiCode
408		}
409	}
410
411	if cw.mode != DiscardNonColorEscSeq || cw.state == outsideCsiCode {
412		nw, err = cw.w.Write(p[first:len(p)])
413		r += nw
414	}
415
416	return r, err
417}
418