github.com/Secbyte/godog@v0.7.14-0.20200116175429-d8f0aeeb70cf/colors/ansi_windows.go (about)

     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  
     7  package colors
     8  
     9  import (
    10  	"bytes"
    11  	"io"
    12  	"strings"
    13  	"syscall"
    14  	"unsafe"
    15  )
    16  
    17  type csiState int
    18  
    19  const (
    20  	outsideCsiCode csiState = iota
    21  	firstCsiCode
    22  	secondCsiCode
    23  )
    24  
    25  type parseResult int
    26  
    27  const (
    28  	noConsole parseResult = iota
    29  	changedColor
    30  	unknown
    31  )
    32  
    33  type ansiColorWriter struct {
    34  	w             io.Writer
    35  	mode          outputMode
    36  	state         csiState
    37  	paramStartBuf bytes.Buffer
    38  	paramBuf      bytes.Buffer
    39  }
    40  
    41  const (
    42  	firstCsiChar   byte = '\x1b'
    43  	secondeCsiChar byte = '['
    44  	separatorChar  byte = ';'
    45  	sgrCode        byte = 'm'
    46  )
    47  
    48  const (
    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  
    63  const (
    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  
   111  type drawType int
   112  
   113  const (
   114  	foreground drawType = iota
   115  	background
   116  )
   117  
   118  type winColor struct {
   119  	code     uint16
   120  	drawType drawType
   121  }
   122  
   123  var 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  
   163  var (
   164  	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
   165  	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
   166  	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
   167  	defaultAttr                    *textAttributes
   168  )
   169  
   170  func 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  
   185  type coord struct {
   186  	X, Y int16
   187  }
   188  
   189  type smallRect struct {
   190  	Left, Top, Right, Bottom int16
   191  }
   192  
   193  type consoleScreenBufferInfo struct {
   194  	DwSize              coord
   195  	DwCursorPosition    coord
   196  	WAttributes         uint16
   197  	SrWindow            smallRect
   198  	DwMaximumWindowSize coord
   199  }
   200  
   201  func 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  
   212  func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
   213  	ret, _, _ := procSetConsoleTextAttribute.Call(
   214  		hConsoleOutput,
   215  		uintptr(wAttributes))
   216  	return ret != 0
   217  }
   218  
   219  type textAttributes struct {
   220  	foregroundColor     uint16
   221  	backgroundColor     uint16
   222  	foregroundIntensity uint16
   223  	backgroundIntensity uint16
   224  	underscore          uint16
   225  	otherAttributes     uint16
   226  }
   227  
   228  func 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  
   238  func 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  
   249  func 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  
   300  func 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  
   313  func (cw *ansiColorWriter) flushBuffer() (int, error) {
   314  	return cw.flushTo(cw.w)
   315  }
   316  
   317  func (cw *ansiColorWriter) resetBuffer() (int, error) {
   318  	return cw.flushTo(nil)
   319  }
   320  
   321  func (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  
   348  func isParameterChar(b byte) bool {
   349  	return ('0' <= b && b <= '9') || b == separatorChar
   350  }
   351  
   352  func (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  }