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