<commit_before>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor

import (
	"bytes"
	"io"
	"strings"
	"syscall"
	"unsafe"
)

type csiState int

const (
	outsideCsiCode csiState = iota
	firstCsiCode
	secondeCsiCode
)

type ansiColorWriter struct {
	w        io.Writer
	state    csiState
	paramBuf bytes.Buffer
}

const (
	firstCsiChar   byte = '\x1b'
	secondeCsiChar byte = '['
	separatorChar  byte = ';'
	sgrCode        byte = 'm'
)

const (
	foregroundBlue      = uint16(0x0001)
	foregroundGreen     = uint16(0x0002)
	foregroundRed       = uint16(0x0004)
	foregroundIntensity = uint16(0x0008)
	backgroundBlue      = uint16(0x0010)
	backgroundGreen     = uint16(0x0020)
	backgroundRed       = uint16(0x0040)
	backgroundIntensity = uint16(0x0080)
	underscore          = uint16(0x8000)

	foregroundMask = foregroundBlue | foregroundGreen | foregroundRed | foregroundIntensity
	backgroundMask = backgroundBlue | backgroundGreen | backgroundRed | backgroundIntensity
)

const (
	ansiReset        = "0"
	ansiIntensityOn  = "1"
	ansiIntensityOff = "21"
	ansiUnderlineOn  = "4"
	ansiUnderlineOff = "24"
	ansiBlinkOn      = "5"
	ansiBlinkOff     = "25"

	ansiForegroundBlack   = "30"
	ansiForegroundRed     = "31"
	ansiForegroundGreen   = "32"
	ansiForegroundYellow  = "33"
	ansiForegroundBlue    = "34"
	ansiForegroundMagenta = "35"
	ansiForegroundCyan    = "36"
	ansiForegroundWhite   = "37"
	ansiForegroundDefault = "39"

	ansiBackgroundBlack   = "40"
	ansiBackgroundRed     = "41"
	ansiBackgroundGreen   = "42"
	ansiBackgroundYellow  = "43"
	ansiBackgroundBlue    = "44"
	ansiBackgroundMagenta = "45"
	ansiBackgroundCyan    = "46"
	ansiBackgroundWhite   = "47"
	ansiBackgroundDefault = "49"

	ansiLightForegroundGray    = "90"
	ansiLightForegroundRed     = "91"
	ansiLightForegroundGreen   = "92"
	ansiLightForegroundYellow  = "93"
	ansiLightForegroundBlue    = "94"
	ansiLightForegroundMagenta = "95"
	ansiLightForegroundCyan    = "96"
	ansiLightForegroundWhite   = "97"

	ansiLightBackgroundGray    = "100"
	ansiLightBackgroundRed     = "101"
	ansiLightBackgroundGreen   = "102"
	ansiLightBackgroundYellow  = "103"
	ansiLightBackgroundBlue    = "104"
	ansiLightBackgroundMagenta = "105"
	ansiLightBackgroundCyan    = "106"
	ansiLightBackgroundWhite   = "107"
)

type drawType int

const (
	foreground drawType = iota
	background
)

type winColor struct {
	code     uint16
	drawType drawType
}

var colorMap = map[string]winColor{
	ansiForegroundBlack:   {0, foreground},
	ansiForegroundRed:     {foregroundRed, foreground},
	ansiForegroundGreen:   {foregroundGreen, foreground},
	ansiForegroundYellow:  {foregroundRed | foregroundGreen, foreground},
	ansiForegroundBlue:    {foregroundBlue, foreground},
	ansiForegroundMagenta: {foregroundRed | foregroundBlue, foreground},
	ansiForegroundCyan:    {foregroundGreen | foregroundBlue, foreground},
	ansiForegroundWhite:   {foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiForegroundDefault: {foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiBackgroundBlack:   {0, background},
	ansiBackgroundRed:     {backgroundRed, background},
	ansiBackgroundGreen:   {backgroundGreen, background},
	ansiBackgroundYellow:  {backgroundRed | backgroundGreen, background},
	ansiBackgroundBlue:    {backgroundBlue, background},
	ansiBackgroundMagenta: {backgroundRed | backgroundBlue, background},
	ansiBackgroundCyan:    {backgroundGreen | backgroundBlue, background},
	ansiBackgroundWhite:   {backgroundRed | backgroundGreen | backgroundBlue, background},
	ansiBackgroundDefault: {0, background},

	ansiLightForegroundGray:    {foregroundIntensity, foreground},
	ansiLightForegroundRed:     {foregroundIntensity | foregroundRed, foreground},
	ansiLightForegroundGreen:   {foregroundIntensity | foregroundGreen, foreground},
	ansiLightForegroundYellow:  {foregroundIntensity | foregroundRed | foregroundGreen, foreground},
	ansiLightForegroundBlue:    {foregroundIntensity | foregroundBlue, foreground},
	ansiLightForegroundMagenta: {foregroundIntensity | foregroundRed | foregroundBlue, foreground},
	ansiLightForegroundCyan:    {foregroundIntensity | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundWhite:   {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiLightBackgroundGray:    {backgroundIntensity, background},
	ansiLightBackgroundRed:     {backgroundIntensity | backgroundRed, background},
	ansiLightBackgroundGreen:   {backgroundIntensity | backgroundGreen, background},
	ansiLightBackgroundYellow:  {backgroundIntensity | backgroundRed | backgroundGreen, background},
	ansiLightBackgroundBlue:    {backgroundIntensity | backgroundBlue, background},
	ansiLightBackgroundMagenta: {backgroundIntensity | backgroundRed | backgroundBlue, background},
	ansiLightBackgroundCyan:    {backgroundIntensity | backgroundGreen | backgroundBlue, background},
	ansiLightBackgroundWhite:   {backgroundIntensity | backgroundRed | backgroundGreen | backgroundBlue, background},
}

var (
	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
	defaultAttr                    *textAttributes
)

func init() {
	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		colorMap[ansiForegroundDefault] = winColor{
			screenInfo.WAttributes & (foregroundRed | foregroundGreen | foregroundBlue),
			foreground,
		}
		colorMap[ansiBackgroundDefault] = winColor{
			screenInfo.WAttributes & (backgroundRed | backgroundGreen | backgroundBlue),
			background,
		}
		defaultAttr = convertTextAttr(screenInfo.WAttributes)
	}
}

type coord struct {
	X, Y int16
}

type smallRect struct {
	Left, Top, Right, Bottom int16
}

type consoleScreenBufferInfo struct {
	DwSize              coord
	DwCursorPosition    coord
	WAttributes         uint16
	SrWindow            smallRect
	DwMaximumWindowSize coord
}

func getConsoleScreenBufferInfo(hConsoleOutput uintptr) *consoleScreenBufferInfo {
	var csbi consoleScreenBufferInfo
	ret, _, _ := procGetConsoleScreenBufferInfo.Call(
		hConsoleOutput,
		uintptr(unsafe.Pointer(&csbi)))
	if ret == 0 {
		return nil
	}
	return &csbi
}

func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
	ret, _, _ := procSetConsoleTextAttribute.Call(
		hConsoleOutput,
		uintptr(wAttributes))
	return ret != 0
}

type textAttributes struct {
	foregroundColor     uint16
	backgroundColor     uint16
	foregroundIntensity uint16
	backgroundIntensity uint16
	underscore          uint16
	otherAttributes     uint16
}

func convertTextAttr(winAttr uint16) *textAttributes {
	fgColor := winAttr & (foregroundRed | foregroundGreen | foregroundBlue)
	bgColor := winAttr & (backgroundRed | backgroundGreen | backgroundBlue)
	fgIntensity := winAttr & foregroundIntensity
	bgIntensity := winAttr & backgroundIntensity
	underline := winAttr & underscore
	otherAttributes := winAttr &^ (foregroundMask | backgroundMask | underscore)
	return &textAttributes{fgColor, bgColor, fgIntensity, bgIntensity, underline, otherAttributes}
}

func convertWinAttr(textAttr *textAttributes) uint16 {
	var winAttr uint16 = 0
	winAttr |= textAttr.foregroundColor
	winAttr |= textAttr.backgroundColor
	winAttr |= textAttr.foregroundIntensity
	winAttr |= textAttr.backgroundIntensity
	winAttr |= textAttr.underscore
	winAttr |= textAttr.otherAttributes
	return winAttr
}

func changeColor(param []byte) {
	if defaultAttr == nil {
		return
	}

	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		return
	}

	winAttr := convertTextAttr(screenInfo.WAttributes)
	strParam := string(param)
	if len(strParam) <= 0 {
		strParam = "0"
	}
	csiParam := strings.Split(strParam, string(separatorChar))
	for _, p := range csiParam {
		c, ok := colorMap[p]
		switch {
		case !ok:
			switch p {
			case ansiReset:
				winAttr.foregroundColor = defaultAttr.foregroundColor
				winAttr.backgroundColor = defaultAttr.backgroundColor
				winAttr.foregroundIntensity = defaultAttr.foregroundIntensity
				winAttr.backgroundIntensity = defaultAttr.backgroundIntensity
				winAttr.underscore = 0
				winAttr.otherAttributes = 0
			case ansiIntensityOn:
				winAttr.foregroundIntensity = foregroundIntensity
			case ansiIntensityOff:
				winAttr.foregroundIntensity = 0
			case ansiUnderlineOn:
				winAttr.underscore = underscore
			case ansiUnderlineOff:
				winAttr.underscore = 0
			case ansiBlinkOn:
				winAttr.backgroundIntensity = backgroundIntensity
			case ansiBlinkOff:
				winAttr.backgroundIntensity = 0
			default:
				// unknown code
			}
		case c.drawType == foreground:
			winAttr.foregroundColor = c.code
		case c.drawType == background:
			winAttr.backgroundColor = c.code
		}
	}
	winTextAttribute := convertWinAttr(winAttr)
	setConsoleTextAttribute(uintptr(syscall.Stdout), winTextAttribute)
}

func parseEscapeSequence(command byte, param []byte) {
	switch command {
	case sgrCode:
		changeColor(param)
	}
}

func isParameterChar(b byte) bool {
	return ('0' <= b && b <= '9') || b == separatorChar
}

func (cw *ansiColorWriter) Write(p []byte) (int, error) {
	r, nw, nc, first, last := 0, 0, 0, 0, 0
	var err error
	for i, ch := range p {
		switch cw.state {
		case outsideCsiCode:
			if ch == firstCsiChar {
				nc++
				cw.state = firstCsiCode
			}
		case firstCsiCode:
			switch ch {
			case firstCsiChar:
				nc++
				break
			case secondeCsiChar:
				nc++
				cw.state = secondeCsiCode
				last = i - 1
			default:
				cw.state = outsideCsiCode
			}
		case secondeCsiCode:
			nc++
			if isParameterChar(ch) {
				cw.paramBuf.WriteByte(ch)
			} else {
				nw, err = cw.w.Write(p[first:last])
				r += nw
				if err != nil {
					return r, err
				}
				first = i + 1
				param := cw.paramBuf.Bytes()
				cw.paramBuf.Reset()
				parseEscapeSequence(ch, param)
				cw.state = outsideCsiCode
			}
		default:
			cw.state = outsideCsiCode
		}
	}

	if cw.state == outsideCsiCode {
		nw, err = cw.w.Write(p[first:len(p)])
	}

	return r + nw + nc, err
}
<commit_msg>Fix misspelled word<commit_after>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor

import (
	"bytes"
	"io"
	"strings"
	"syscall"
	"unsafe"
)

type csiState int

const (
	outsideCsiCode csiState = iota
	firstCsiCode
	secondCsiCode
)

type ansiColorWriter struct {
	w        io.Writer
	state    csiState
	paramBuf bytes.Buffer
}

const (
	firstCsiChar   byte = '\x1b'
	secondeCsiChar byte = '['
	separatorChar  byte = ';'
	sgrCode        byte = 'm'
)

const (
	foregroundBlue      = uint16(0x0001)
	foregroundGreen     = uint16(0x0002)
	foregroundRed       = uint16(0x0004)
	foregroundIntensity = uint16(0x0008)
	backgroundBlue      = uint16(0x0010)
	backgroundGreen     = uint16(0x0020)
	backgroundRed       = uint16(0x0040)
	backgroundIntensity = uint16(0x0080)
	underscore          = uint16(0x8000)

	foregroundMask = foregroundBlue | foregroundGreen | foregroundRed | foregroundIntensity
	backgroundMask = backgroundBlue | backgroundGreen | backgroundRed | backgroundIntensity
)

const (
	ansiReset        = "0"
	ansiIntensityOn  = "1"
	ansiIntensityOff = "21"
	ansiUnderlineOn  = "4"
	ansiUnderlineOff = "24"
	ansiBlinkOn      = "5"
	ansiBlinkOff     = "25"

	ansiForegroundBlack   = "30"
	ansiForegroundRed     = "31"
	ansiForegroundGreen   = "32"
	ansiForegroundYellow  = "33"
	ansiForegroundBlue    = "34"
	ansiForegroundMagenta = "35"
	ansiForegroundCyan    = "36"
	ansiForegroundWhite   = "37"
	ansiForegroundDefault = "39"

	ansiBackgroundBlack   = "40"
	ansiBackgroundRed     = "41"
	ansiBackgroundGreen   = "42"
	ansiBackgroundYellow  = "43"
	ansiBackgroundBlue    = "44"
	ansiBackgroundMagenta = "45"
	ansiBackgroundCyan    = "46"
	ansiBackgroundWhite   = "47"
	ansiBackgroundDefault = "49"

	ansiLightForegroundGray    = "90"
	ansiLightForegroundRed     = "91"
	ansiLightForegroundGreen   = "92"
	ansiLightForegroundYellow  = "93"
	ansiLightForegroundBlue    = "94"
	ansiLightForegroundMagenta = "95"
	ansiLightForegroundCyan    = "96"
	ansiLightForegroundWhite   = "97"

	ansiLightBackgroundGray    = "100"
	ansiLightBackgroundRed     = "101"
	ansiLightBackgroundGreen   = "102"
	ansiLightBackgroundYellow  = "103"
	ansiLightBackgroundBlue    = "104"
	ansiLightBackgroundMagenta = "105"
	ansiLightBackgroundCyan    = "106"
	ansiLightBackgroundWhite   = "107"
)

type drawType int

const (
	foreground drawType = iota
	background
)

type winColor struct {
	code     uint16
	drawType drawType
}

var colorMap = map[string]winColor{
	ansiForegroundBlack:   {0, foreground},
	ansiForegroundRed:     {foregroundRed, foreground},
	ansiForegroundGreen:   {foregroundGreen, foreground},
	ansiForegroundYellow:  {foregroundRed | foregroundGreen, foreground},
	ansiForegroundBlue:    {foregroundBlue, foreground},
	ansiForegroundMagenta: {foregroundRed | foregroundBlue, foreground},
	ansiForegroundCyan:    {foregroundGreen | foregroundBlue, foreground},
	ansiForegroundWhite:   {foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiForegroundDefault: {foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiBackgroundBlack:   {0, background},
	ansiBackgroundRed:     {backgroundRed, background},
	ansiBackgroundGreen:   {backgroundGreen, background},
	ansiBackgroundYellow:  {backgroundRed | backgroundGreen, background},
	ansiBackgroundBlue:    {backgroundBlue, background},
	ansiBackgroundMagenta: {backgroundRed | backgroundBlue, background},
	ansiBackgroundCyan:    {backgroundGreen | backgroundBlue, background},
	ansiBackgroundWhite:   {backgroundRed | backgroundGreen | backgroundBlue, background},
	ansiBackgroundDefault: {0, background},

	ansiLightForegroundGray:    {foregroundIntensity, foreground},
	ansiLightForegroundRed:     {foregroundIntensity | foregroundRed, foreground},
	ansiLightForegroundGreen:   {foregroundIntensity | foregroundGreen, foreground},
	ansiLightForegroundYellow:  {foregroundIntensity | foregroundRed | foregroundGreen, foreground},
	ansiLightForegroundBlue:    {foregroundIntensity | foregroundBlue, foreground},
	ansiLightForegroundMagenta: {foregroundIntensity | foregroundRed | foregroundBlue, foreground},
	ansiLightForegroundCyan:    {foregroundIntensity | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundWhite:   {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiLightBackgroundGray:    {backgroundIntensity, background},
	ansiLightBackgroundRed:     {backgroundIntensity | backgroundRed, background},
	ansiLightBackgroundGreen:   {backgroundIntensity | backgroundGreen, background},
	ansiLightBackgroundYellow:  {backgroundIntensity | backgroundRed | backgroundGreen, background},
	ansiLightBackgroundBlue:    {backgroundIntensity | backgroundBlue, background},
	ansiLightBackgroundMagenta: {backgroundIntensity | backgroundRed | backgroundBlue, background},
	ansiLightBackgroundCyan:    {backgroundIntensity | backgroundGreen | backgroundBlue, background},
	ansiLightBackgroundWhite:   {backgroundIntensity | backgroundRed | backgroundGreen | backgroundBlue, background},
}

var (
	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
	defaultAttr                    *textAttributes
)

func init() {
	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		colorMap[ansiForegroundDefault] = winColor{
			screenInfo.WAttributes & (foregroundRed | foregroundGreen | foregroundBlue),
			foreground,
		}
		colorMap[ansiBackgroundDefault] = winColor{
			screenInfo.WAttributes & (backgroundRed | backgroundGreen | backgroundBlue),
			background,
		}
		defaultAttr = convertTextAttr(screenInfo.WAttributes)
	}
}

type coord struct {
	X, Y int16
}

type smallRect struct {
	Left, Top, Right, Bottom int16
}

type consoleScreenBufferInfo struct {
	DwSize              coord
	DwCursorPosition    coord
	WAttributes         uint16
	SrWindow            smallRect
	DwMaximumWindowSize coord
}

func getConsoleScreenBufferInfo(hConsoleOutput uintptr) *consoleScreenBufferInfo {
	var csbi consoleScreenBufferInfo
	ret, _, _ := procGetConsoleScreenBufferInfo.Call(
		hConsoleOutput,
		uintptr(unsafe.Pointer(&csbi)))
	if ret == 0 {
		return nil
	}
	return &csbi
}

func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
	ret, _, _ := procSetConsoleTextAttribute.Call(
		hConsoleOutput,
		uintptr(wAttributes))
	return ret != 0
}

type textAttributes struct {
	foregroundColor     uint16
	backgroundColor     uint16
	foregroundIntensity uint16
	backgroundIntensity uint16
	underscore          uint16
	otherAttributes     uint16
}

func convertTextAttr(winAttr uint16) *textAttributes {
	fgColor := winAttr & (foregroundRed | foregroundGreen | foregroundBlue)
	bgColor := winAttr & (backgroundRed | backgroundGreen | backgroundBlue)
	fgIntensity := winAttr & foregroundIntensity
	bgIntensity := winAttr & backgroundIntensity
	underline := winAttr & underscore
	otherAttributes := winAttr &^ (foregroundMask | backgroundMask | underscore)
	return &textAttributes{fgColor, bgColor, fgIntensity, bgIntensity, underline, otherAttributes}
}

func convertWinAttr(textAttr *textAttributes) uint16 {
	var winAttr uint16 = 0
	winAttr |= textAttr.foregroundColor
	winAttr |= textAttr.backgroundColor
	winAttr |= textAttr.foregroundIntensity
	winAttr |= textAttr.backgroundIntensity
	winAttr |= textAttr.underscore
	winAttr |= textAttr.otherAttributes
	return winAttr
}

func changeColor(param []byte) {
	if defaultAttr == nil {
		return
	}

	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		return
	}

	winAttr := convertTextAttr(screenInfo.WAttributes)
	strParam := string(param)
	if len(strParam) <= 0 {
		strParam = "0"
	}
	csiParam := strings.Split(strParam, string(separatorChar))
	for _, p := range csiParam {
		c, ok := colorMap[p]
		switch {
		case !ok:
			switch p {
			case ansiReset:
				winAttr.foregroundColor = defaultAttr.foregroundColor
				winAttr.backgroundColor = defaultAttr.backgroundColor
				winAttr.foregroundIntensity = defaultAttr.foregroundIntensity
				winAttr.backgroundIntensity = defaultAttr.backgroundIntensity
				winAttr.underscore = 0
				winAttr.otherAttributes = 0
			case ansiIntensityOn:
				winAttr.foregroundIntensity = foregroundIntensity
			case ansiIntensityOff:
				winAttr.foregroundIntensity = 0
			case ansiUnderlineOn:
				winAttr.underscore = underscore
			case ansiUnderlineOff:
				winAttr.underscore = 0
			case ansiBlinkOn:
				winAttr.backgroundIntensity = backgroundIntensity
			case ansiBlinkOff:
				winAttr.backgroundIntensity = 0
			default:
				// unknown code
			}
		case c.drawType == foreground:
			winAttr.foregroundColor = c.code
		case c.drawType == background:
			winAttr.backgroundColor = c.code
		}
	}
	winTextAttribute := convertWinAttr(winAttr)
	setConsoleTextAttribute(uintptr(syscall.Stdout), winTextAttribute)
}

func parseEscapeSequence(command byte, param []byte) {
	switch command {
	case sgrCode:
		changeColor(param)
	}
}

func isParameterChar(b byte) bool {
	return ('0' <= b && b <= '9') || b == separatorChar
}

func (cw *ansiColorWriter) Write(p []byte) (int, error) {
	r, nw, nc, first, last := 0, 0, 0, 0, 0
	var err error
	for i, ch := range p {
		switch cw.state {
		case outsideCsiCode:
			if ch == firstCsiChar {
				nc++
				cw.state = firstCsiCode
			}
		case firstCsiCode:
			switch ch {
			case firstCsiChar:
				nc++
				break
			case secondeCsiChar:
				nc++
				cw.state = secondCsiCode
				last = i - 1
			default:
				cw.state = outsideCsiCode
			}
		case secondCsiCode:
			nc++
			if isParameterChar(ch) {
				cw.paramBuf.WriteByte(ch)
			} else {
				nw, err = cw.w.Write(p[first:last])
				r += nw
				if err != nil {
					return r, err
				}
				first = i + 1
				param := cw.paramBuf.Bytes()
				cw.paramBuf.Reset()
				parseEscapeSequence(ch, param)
				cw.state = outsideCsiCode
			}
		default:
			cw.state = outsideCsiCode
		}
	}

	if cw.state == outsideCsiCode {
		nw, err = cw.w.Write(p[first:len(p)])
	}

	return r + nw + nc, err
}
<|endoftext|>


<commit_before>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor

import (
	"bytes"
	"io"
	"strings"
	"syscall"
	"unsafe"
)

type csiState int

const (
	outsideCsiCode csiState = iota
	firstCsiCode
	secondeCsiCode
)

type ansiColorWriter struct {
	w        io.Writer
	state    csiState
	paramBuf bytes.Buffer
}

const (
	firstCsiChar   byte = '\x1b'
	secondeCsiChar byte = '['
	separatorChar  byte = ';'
	sgrCode        byte = 'm'
)

const (
	foregroundBlue      = uint16(0x0001)
	foregroundGreen     = uint16(0x0002)
	foregroundRed       = uint16(0x0004)
	foregroundIntensity = uint16(0x0008)
	backgroundBlue      = uint16(0x0010)
	backgroundGreen     = uint16(0x0020)
	backgroundRed       = uint16(0x0040)
	backgroundIntensity = uint16(0x0080)
	underscore          = uint16(0x8000)

	foregroundMask = foregroundBlue | foregroundGreen | foregroundRed | foregroundIntensity
	backgroundMask = backgroundBlue | backgroundGreen | backgroundRed | backgroundIntensity
)

const (
	ansiReset        = "0"
	ansiIntensityOn  = "1"
	ansiIntensityOff = "21"
	ansiUnderlineOn  = "4"
	ansiUnderlineOff = "24"
	ansiBlinkOn      = "5"
	ansiBlinkOff     = "25"

	ansiForegroundBlack   = "30"
	ansiForegroundRed     = "31"
	ansiForegroundGreen   = "32"
	ansiForegroundYellow  = "33"
	ansiForegroundBlue    = "34"
	ansiForegroundMagenta = "35"
	ansiForegroundCyan    = "36"
	ansiForegroundWhite   = "37"
	ansiForegroundDefault = "39"

	ansiBackgroundBlack   = "40"
	ansiBackgroundRed     = "41"
	ansiBackgroundGreen   = "42"
	ansiBackgroundYellow  = "43"
	ansiBackgroundBlue    = "44"
	ansiBackgroundMagenta = "45"
	ansiBackgroundCyan    = "46"
	ansiBackgroundWhite   = "47"
	ansiBackgroundDefault = "49"

	ansiLightForegroundGray    = "90"
	ansiLightForegroundRed     = "91"
	ansiLightForegroundGreen   = "92"
	ansiLightForegroundYellow  = "93"
	ansiLightForegroundBlue    = "94"
	ansiLightForegroundMagenta = "95"
	ansiLightForegroundCyan    = "96"
	ansiLightForegroundWhite   = "97"
	ansiLightForegroundDefault = "99"
)

type drawType int

const (
	foreground drawType = iota
	background
)

type winColor struct {
	code     uint16
	drawType drawType
}

var colorMap = map[string]winColor{
	ansiForegroundBlack:   {0, foreground},
	ansiForegroundRed:     {foregroundRed, foreground},
	ansiForegroundGreen:   {foregroundGreen, foreground},
	ansiForegroundYellow:  {foregroundRed | foregroundGreen, foreground},
	ansiForegroundBlue:    {foregroundBlue, foreground},
	ansiForegroundMagenta: {foregroundRed | foregroundBlue, foreground},
	ansiForegroundCyan:    {foregroundGreen | foregroundBlue, foreground},
	ansiForegroundWhite:   {foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiForegroundDefault: {foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiBackgroundBlack:   {0, background},
	ansiBackgroundRed:     {backgroundRed, background},
	ansiBackgroundGreen:   {backgroundGreen, background},
	ansiBackgroundYellow:  {backgroundRed | backgroundGreen, background},
	ansiBackgroundBlue:    {backgroundBlue, background},
	ansiBackgroundMagenta: {backgroundRed | backgroundBlue, background},
	ansiBackgroundCyan:    {backgroundGreen | backgroundBlue, background},
	ansiBackgroundWhite:   {backgroundRed | backgroundGreen | backgroundBlue, background},
	ansiBackgroundDefault: {0, background},

	ansiLightForegroundGray:    {foregroundIntensity, foreground},
	ansiLightForegroundRed:     {foregroundIntensity | foregroundRed, foreground},
	ansiLightForegroundGreen:   {foregroundIntensity | foregroundGreen, foreground},
	ansiLightForegroundYellow:  {foregroundIntensity | foregroundRed | foregroundGreen, foreground},
	ansiLightForegroundBlue:    {foregroundIntensity | foregroundBlue, foreground},
	ansiLightForegroundMagenta: {foregroundIntensity | foregroundRed | foregroundBlue, foreground},
	ansiLightForegroundCyan:    {foregroundIntensity | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundWhite:   {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundDefault: {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},
}

var (
	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
	defaultAttr                    *textAttributes
)

func init() {
	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		colorMap[ansiForegroundDefault] = winColor{
			screenInfo.WAttributes & (foregroundRed | foregroundGreen | foregroundBlue),
			foreground,
		}
		colorMap[ansiBackgroundDefault] = winColor{
			screenInfo.WAttributes & (backgroundRed | backgroundGreen | backgroundBlue),
			background,
		}
		defaultAttr = convertTextAttr(screenInfo.WAttributes)
	}
}

type coord struct {
	X, Y int16
}

type smallRect struct {
	Left, Top, Right, Bottom int16
}

type consoleScreenBufferInfo struct {
	DwSize              coord
	DwCursorPosition    coord
	WAttributes         uint16
	SrWindow            smallRect
	DwMaximumWindowSize coord
}

func getConsoleScreenBufferInfo(hConsoleOutput uintptr) *consoleScreenBufferInfo {
	var csbi consoleScreenBufferInfo
	ret, _, _ := procGetConsoleScreenBufferInfo.Call(
		hConsoleOutput,
		uintptr(unsafe.Pointer(&csbi)))
	if ret == 0 {
		return nil
	}
	return &csbi
}

func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
	ret, _, _ := procSetConsoleTextAttribute.Call(
		hConsoleOutput,
		uintptr(wAttributes))
	return ret != 0
}

type textAttributes struct {
	foregroundColor     uint16
	backgroundColor     uint16
	foregroundIntensity uint16
	backgroundIntensity uint16
	underscore          uint16
	otherAttributes     uint16
}

func convertTextAttr(winAttr uint16) *textAttributes {
	fgColor := winAttr & (foregroundRed | foregroundGreen | foregroundBlue)
	bgColor := winAttr & (backgroundRed | backgroundGreen | backgroundBlue)
	fgIntensity := winAttr & foregroundIntensity
	bgIntensity := winAttr & backgroundIntensity
	underline := winAttr & underscore
	otherAttributes := winAttr &^ (foregroundMask | backgroundMask | underscore)
	return &textAttributes{fgColor, bgColor, fgIntensity, bgIntensity, underline, otherAttributes}
}

func convertWinAttr(textAttr *textAttributes) uint16 {
	var winAttr uint16 = 0
	winAttr |= textAttr.foregroundColor
	winAttr |= textAttr.backgroundColor
	winAttr |= textAttr.foregroundIntensity
	winAttr |= textAttr.backgroundIntensity
	winAttr |= textAttr.underscore
	winAttr |= textAttr.otherAttributes
	return winAttr
}

func changeColor(param []byte) {
	if defaultAttr == nil {
		return
	}

	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		return
	}

	winAttr := convertTextAttr(screenInfo.WAttributes)
	strParam := string(param)
	if len(strParam) <= 0 {
		strParam = "0"
	}
	csiParam := strings.Split(strParam, string(separatorChar))
	for _, p := range csiParam {
		c, ok := colorMap[p]
		switch {
		case !ok:
			switch p {
			case ansiReset:
				winAttr.foregroundColor = defaultAttr.foregroundColor
				winAttr.backgroundColor = defaultAttr.backgroundColor
				winAttr.foregroundIntensity = defaultAttr.foregroundIntensity
				winAttr.backgroundIntensity = defaultAttr.backgroundIntensity
				winAttr.underscore = 0
				winAttr.otherAttributes = 0
			case ansiIntensityOn:
				winAttr.foregroundIntensity = foregroundIntensity
			case ansiIntensityOff:
				winAttr.foregroundIntensity = 0
			case ansiUnderlineOn:
				winAttr.underscore = underscore
			case ansiUnderlineOff:
				winAttr.underscore = 0
			case ansiBlinkOn:
				winAttr.backgroundIntensity = backgroundIntensity
			case ansiBlinkOff:
				winAttr.backgroundIntensity = 0
			default:
				// unknown code
			}
		case c.drawType == foreground:
			winAttr.foregroundColor = c.code
		case c.drawType == background:
			winAttr.backgroundColor = c.code
		}
	}
	winTextAttribute := convertWinAttr(winAttr)
	setConsoleTextAttribute(uintptr(syscall.Stdout), winTextAttribute)
}

func parseEscapeSequence(command byte, param []byte) {
	switch command {
	case sgrCode:
		changeColor(param)
	}
}

func isParameterChar(b byte) bool {
	return ('0' <= b && b <= '9') || b == separatorChar
}

func (cw *ansiColorWriter) Write(p []byte) (int, error) {
	r, nw, nc, first, last := 0, 0, 0, 0, 0
	var err error
	for i, ch := range p {
		switch cw.state {
		case outsideCsiCode:
			if ch == firstCsiChar {
				nc++
				cw.state = firstCsiCode
			}
		case firstCsiCode:
			switch ch {
			case firstCsiChar:
				nc++
				break
			case secondeCsiChar:
				nc++
				cw.state = secondeCsiCode
				last = i - 1
			default:
				cw.state = outsideCsiCode
			}
		case secondeCsiCode:
			nc++
			if isParameterChar(ch) {
				cw.paramBuf.WriteByte(ch)
			} else {
				nw, err = cw.w.Write(p[first:last])
				r += nw
				if err != nil {
					return r, err
				}
				first = i + 1
				param := cw.paramBuf.Bytes()
				cw.paramBuf.Reset()
				parseEscapeSequence(ch, param)
				cw.state = outsideCsiCode
			}
		default:
			cw.state = outsideCsiCode
		}
	}

	if cw.state == outsideCsiCode {
		nw, err = cw.w.Write(p[first:len(p)])
	}

	return r + nw + nc, err
}
<commit_msg>Add background  high intensity colors<commit_after>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor

import (
	"bytes"
	"io"
	"strings"
	"syscall"
	"unsafe"
)

type csiState int

const (
	outsideCsiCode csiState = iota
	firstCsiCode
	secondeCsiCode
)

type ansiColorWriter struct {
	w        io.Writer
	state    csiState
	paramBuf bytes.Buffer
}

const (
	firstCsiChar   byte = '\x1b'
	secondeCsiChar byte = '['
	separatorChar  byte = ';'
	sgrCode        byte = 'm'
)

const (
	foregroundBlue      = uint16(0x0001)
	foregroundGreen     = uint16(0x0002)
	foregroundRed       = uint16(0x0004)
	foregroundIntensity = uint16(0x0008)
	backgroundBlue      = uint16(0x0010)
	backgroundGreen     = uint16(0x0020)
	backgroundRed       = uint16(0x0040)
	backgroundIntensity = uint16(0x0080)
	underscore          = uint16(0x8000)

	foregroundMask = foregroundBlue | foregroundGreen | foregroundRed | foregroundIntensity
	backgroundMask = backgroundBlue | backgroundGreen | backgroundRed | backgroundIntensity
)

const (
	ansiReset        = "0"
	ansiIntensityOn  = "1"
	ansiIntensityOff = "21"
	ansiUnderlineOn  = "4"
	ansiUnderlineOff = "24"
	ansiBlinkOn      = "5"
	ansiBlinkOff     = "25"

	ansiForegroundBlack   = "30"
	ansiForegroundRed     = "31"
	ansiForegroundGreen   = "32"
	ansiForegroundYellow  = "33"
	ansiForegroundBlue    = "34"
	ansiForegroundMagenta = "35"
	ansiForegroundCyan    = "36"
	ansiForegroundWhite   = "37"
	ansiForegroundDefault = "39"

	ansiBackgroundBlack   = "40"
	ansiBackgroundRed     = "41"
	ansiBackgroundGreen   = "42"
	ansiBackgroundYellow  = "43"
	ansiBackgroundBlue    = "44"
	ansiBackgroundMagenta = "45"
	ansiBackgroundCyan    = "46"
	ansiBackgroundWhite   = "47"
	ansiBackgroundDefault = "49"

	ansiLightForegroundGray    = "90"
	ansiLightForegroundRed     = "91"
	ansiLightForegroundGreen   = "92"
	ansiLightForegroundYellow  = "93"
	ansiLightForegroundBlue    = "94"
	ansiLightForegroundMagenta = "95"
	ansiLightForegroundCyan    = "96"
	ansiLightForegroundWhite   = "97"

	ansiLightBackgroundGray    = "100"
	ansiLightBackgroundRed     = "101"
	ansiLightBackgroundGreen   = "102"
	ansiLightBackgroundYellow  = "103"
	ansiLightBackgroundBlue    = "104"
	ansiLightBackgroundMagenta = "105"
	ansiLightBackgroundCyan    = "106"
	ansiLightBackgroundWhite   = "107"
)

type drawType int

const (
	foreground drawType = iota
	background
)

type winColor struct {
	code     uint16
	drawType drawType
}

var colorMap = map[string]winColor{
	ansiForegroundBlack:   {0, foreground},
	ansiForegroundRed:     {foregroundRed, foreground},
	ansiForegroundGreen:   {foregroundGreen, foreground},
	ansiForegroundYellow:  {foregroundRed | foregroundGreen, foreground},
	ansiForegroundBlue:    {foregroundBlue, foreground},
	ansiForegroundMagenta: {foregroundRed | foregroundBlue, foreground},
	ansiForegroundCyan:    {foregroundGreen | foregroundBlue, foreground},
	ansiForegroundWhite:   {foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiForegroundDefault: {foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiBackgroundBlack:   {0, background},
	ansiBackgroundRed:     {backgroundRed, background},
	ansiBackgroundGreen:   {backgroundGreen, background},
	ansiBackgroundYellow:  {backgroundRed | backgroundGreen, background},
	ansiBackgroundBlue:    {backgroundBlue, background},
	ansiBackgroundMagenta: {backgroundRed | backgroundBlue, background},
	ansiBackgroundCyan:    {backgroundGreen | backgroundBlue, background},
	ansiBackgroundWhite:   {backgroundRed | backgroundGreen | backgroundBlue, background},
	ansiBackgroundDefault: {0, background},

	ansiLightForegroundGray:    {foregroundIntensity, foreground},
	ansiLightForegroundRed:     {foregroundIntensity | foregroundRed, foreground},
	ansiLightForegroundGreen:   {foregroundIntensity | foregroundGreen, foreground},
	ansiLightForegroundYellow:  {foregroundIntensity | foregroundRed | foregroundGreen, foreground},
	ansiLightForegroundBlue:    {foregroundIntensity | foregroundBlue, foreground},
	ansiLightForegroundMagenta: {foregroundIntensity | foregroundRed | foregroundBlue, foreground},
	ansiLightForegroundCyan:    {foregroundIntensity | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundWhite:   {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiLightBackgroundGray:    {backgroundIntensity, background},
	ansiLightBackgroundRed:     {backgroundIntensity | backgroundRed, background},
	ansiLightBackgroundGreen:   {backgroundIntensity | backgroundGreen, background},
	ansiLightBackgroundYellow:  {backgroundIntensity | backgroundRed | backgroundGreen, background},
	ansiLightBackgroundBlue:    {backgroundIntensity | backgroundBlue, background},
	ansiLightBackgroundMagenta: {backgroundIntensity | backgroundRed | backgroundBlue, background},
	ansiLightBackgroundCyan:    {backgroundIntensity | backgroundGreen | backgroundBlue, background},
	ansiLightBackgroundWhite:   {backgroundIntensity | backgroundRed | backgroundGreen | backgroundBlue, background},
}

var (
	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
	defaultAttr                    *textAttributes
)

func init() {
	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		colorMap[ansiForegroundDefault] = winColor{
			screenInfo.WAttributes & (foregroundRed | foregroundGreen | foregroundBlue),
			foreground,
		}
		colorMap[ansiBackgroundDefault] = winColor{
			screenInfo.WAttributes & (backgroundRed | backgroundGreen | backgroundBlue),
			background,
		}
		defaultAttr = convertTextAttr(screenInfo.WAttributes)
	}
}

type coord struct {
	X, Y int16
}

type smallRect struct {
	Left, Top, Right, Bottom int16
}

type consoleScreenBufferInfo struct {
	DwSize              coord
	DwCursorPosition    coord
	WAttributes         uint16
	SrWindow            smallRect
	DwMaximumWindowSize coord
}

func getConsoleScreenBufferInfo(hConsoleOutput uintptr) *consoleScreenBufferInfo {
	var csbi consoleScreenBufferInfo
	ret, _, _ := procGetConsoleScreenBufferInfo.Call(
		hConsoleOutput,
		uintptr(unsafe.Pointer(&csbi)))
	if ret == 0 {
		return nil
	}
	return &csbi
}

func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
	ret, _, _ := procSetConsoleTextAttribute.Call(
		hConsoleOutput,
		uintptr(wAttributes))
	return ret != 0
}

type textAttributes struct {
	foregroundColor     uint16
	backgroundColor     uint16
	foregroundIntensity uint16
	backgroundIntensity uint16
	underscore          uint16
	otherAttributes     uint16
}

func convertTextAttr(winAttr uint16) *textAttributes {
	fgColor := winAttr & (foregroundRed | foregroundGreen | foregroundBlue)
	bgColor := winAttr & (backgroundRed | backgroundGreen | backgroundBlue)
	fgIntensity := winAttr & foregroundIntensity
	bgIntensity := winAttr & backgroundIntensity
	underline := winAttr & underscore
	otherAttributes := winAttr &^ (foregroundMask | backgroundMask | underscore)
	return &textAttributes{fgColor, bgColor, fgIntensity, bgIntensity, underline, otherAttributes}
}

func convertWinAttr(textAttr *textAttributes) uint16 {
	var winAttr uint16 = 0
	winAttr |= textAttr.foregroundColor
	winAttr |= textAttr.backgroundColor
	winAttr |= textAttr.foregroundIntensity
	winAttr |= textAttr.backgroundIntensity
	winAttr |= textAttr.underscore
	winAttr |= textAttr.otherAttributes
	return winAttr
}

func changeColor(param []byte) {
	if defaultAttr == nil {
		return
	}

	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		return
	}

	winAttr := convertTextAttr(screenInfo.WAttributes)
	strParam := string(param)
	if len(strParam) <= 0 {
		strParam = "0"
	}
	csiParam := strings.Split(strParam, string(separatorChar))
	for _, p := range csiParam {
		c, ok := colorMap[p]
		switch {
		case !ok:
			switch p {
			case ansiReset:
				winAttr.foregroundColor = defaultAttr.foregroundColor
				winAttr.backgroundColor = defaultAttr.backgroundColor
				winAttr.foregroundIntensity = defaultAttr.foregroundIntensity
				winAttr.backgroundIntensity = defaultAttr.backgroundIntensity
				winAttr.underscore = 0
				winAttr.otherAttributes = 0
			case ansiIntensityOn:
				winAttr.foregroundIntensity = foregroundIntensity
			case ansiIntensityOff:
				winAttr.foregroundIntensity = 0
			case ansiUnderlineOn:
				winAttr.underscore = underscore
			case ansiUnderlineOff:
				winAttr.underscore = 0
			case ansiBlinkOn:
				winAttr.backgroundIntensity = backgroundIntensity
			case ansiBlinkOff:
				winAttr.backgroundIntensity = 0
			default:
				// unknown code
			}
		case c.drawType == foreground:
			winAttr.foregroundColor = c.code
		case c.drawType == background:
			winAttr.backgroundColor = c.code
		}
	}
	winTextAttribute := convertWinAttr(winAttr)
	setConsoleTextAttribute(uintptr(syscall.Stdout), winTextAttribute)
}

func parseEscapeSequence(command byte, param []byte) {
	switch command {
	case sgrCode:
		changeColor(param)
	}
}

func isParameterChar(b byte) bool {
	return ('0' <= b && b <= '9') || b == separatorChar
}

func (cw *ansiColorWriter) Write(p []byte) (int, error) {
	r, nw, nc, first, last := 0, 0, 0, 0, 0
	var err error
	for i, ch := range p {
		switch cw.state {
		case outsideCsiCode:
			if ch == firstCsiChar {
				nc++
				cw.state = firstCsiCode
			}
		case firstCsiCode:
			switch ch {
			case firstCsiChar:
				nc++
				break
			case secondeCsiChar:
				nc++
				cw.state = secondeCsiCode
				last = i - 1
			default:
				cw.state = outsideCsiCode
			}
		case secondeCsiCode:
			nc++
			if isParameterChar(ch) {
				cw.paramBuf.WriteByte(ch)
			} else {
				nw, err = cw.w.Write(p[first:last])
				r += nw
				if err != nil {
					return r, err
				}
				first = i + 1
				param := cw.paramBuf.Bytes()
				cw.paramBuf.Reset()
				parseEscapeSequence(ch, param)
				cw.state = outsideCsiCode
			}
		default:
			cw.state = outsideCsiCode
		}
	}

	if cw.state == outsideCsiCode {
		nw, err = cw.w.Write(p[first:len(p)])
	}

	return r + nw + nc, err
}
<|endoftext|>


<commit_before>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor

import (
	"bytes"
	"io"
	"strings"
	"syscall"
	"unsafe"
)

type csiState int

const (
	outsideCsiCode csiState = iota
	firstCsiCode
	secondeCsiCode
)

type ansiColorWriter struct {
	w        io.Writer
	state    csiState
	paramBuf bytes.Buffer
}

const (
	firstCsiChar   byte = '\x1b'
	secondeCsiChar byte = '['
	separatorChar  byte = ';'
	sgrCode        byte = 'm'
)

const (
	foregroundBlue      = uint16(0x0001)
	foregroundGreen     = uint16(0x0002)
	foregroundRed       = uint16(0x0004)
	foregroundIntensity = uint16(0x0008)
	backgroundBlue      = uint16(0x0010)
	backgroundGreen     = uint16(0x0020)
	backgroundRed       = uint16(0x0040)
	backgroundIntensity = uint16(0x0080)
	underscore          = uint16(0x8000)

	foregroundMask = foregroundBlue | foregroundGreen | foregroundRed | foregroundIntensity
	backgroundMask = backgroundBlue | backgroundGreen | backgroundRed | backgroundIntensity
)

const (
	ansiReset        = "0"
	ansiIntensityOn  = "1"
	ansiIntensityOff = "21"
	ansiUnderlineOn  = "4"
	ansiUnderlineOff = "24"
	ansiBlinkOn      = "5"
	ansiBlinkOff     = "25"

	ansiForegroundBlack   = "30"
	ansiForegroundRed     = "31"
	ansiForegroundGreen   = "32"
	ansiForegroundYellow  = "33"
	ansiForegroundBlue    = "34"
	ansiForegroundMagenta = "35"
	ansiForegroundCyan    = "36"
	ansiForegroundWhite   = "37"
	ansiForegroundDefault = "39"

	ansiBackgroundBlack   = "40"
	ansiBackgroundRed     = "41"
	ansiBackgroundGreen   = "42"
	ansiBackgroundYellow  = "43"
	ansiBackgroundBlue    = "44"
	ansiBackgroundMagenta = "45"
	ansiBackgroundCyan    = "46"
	ansiBackgroundWhite   = "47"
	ansiBackgroundDefault = "49"

	ansiLightForegroundGray    = "90"
	ansiLightForegroundRed     = "91"
	ansiLightForegroundGreen   = "92"
	ansiLightForegroundYellow  = "93"
	ansiLightForegroundBlue    = "94"
	ansiLightForegroundMagenta = "95"
	ansiLightForegroundCyan    = "96"
	ansiLightForegroundWhite   = "97"
	ansiLightForegroundDefault = "99"
)

type drawType int

const (
	foreground drawType = iota
	background
)

type winColor struct {
	code     uint16
	drawType drawType
}

var colorMap = map[string]winColor{
	ansiForegroundBlack:   {0, foreground},
	ansiForegroundRed:     {foregroundRed, foreground},
	ansiForegroundGreen:   {foregroundGreen, foreground},
	ansiForegroundYellow:  {foregroundRed | foregroundGreen, foreground},
	ansiForegroundBlue:    {foregroundBlue, foreground},
	ansiForegroundMagenta: {foregroundRed | foregroundBlue, foreground},
	ansiForegroundCyan:    {foregroundGreen | foregroundBlue, foreground},
	ansiForegroundWhite:   {foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiForegroundDefault: {foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiBackgroundBlack:   {0, background},
	ansiBackgroundRed:     {backgroundRed, background},
	ansiBackgroundGreen:   {backgroundGreen, background},
	ansiBackgroundYellow:  {backgroundRed | backgroundGreen, background},
	ansiBackgroundBlue:    {backgroundBlue, background},
	ansiBackgroundMagenta: {backgroundRed | backgroundBlue, background},
	ansiBackgroundCyan:    {backgroundGreen | backgroundBlue, background},
	ansiBackgroundWhite:   {backgroundRed | backgroundGreen | backgroundBlue, background},
	ansiBackgroundDefault: {0, background},

	ansiLightForegroundGray:    {foregroundIntensity, foreground},
	ansiLightForegroundRed:     {foregroundIntensity | foregroundRed, foreground},
	ansiLightForegroundGreen:   {foregroundIntensity | foregroundGreen, foreground},
	ansiLightForegroundYellow:  {foregroundIntensity | foregroundRed | foregroundGreen, foreground},
	ansiLightForegroundBlue:    {foregroundIntensity | foregroundBlue, foreground},
	ansiLightForegroundMagenta: {foregroundIntensity | foregroundRed | foregroundBlue, foreground},
	ansiLightForegroundCyan:    {foregroundIntensity | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundWhite:   {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},
}

var (
	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
	defaultAttr                    *textAttributes
)

func init() {
	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		colorMap[ansiForegroundDefault] = winColor{
			screenInfo.WAttributes & (foregroundRed | foregroundGreen | foregroundBlue),
			foreground,
		}
		colorMap[ansiBackgroundDefault] = winColor{
			screenInfo.WAttributes & (backgroundRed | backgroundGreen | backgroundBlue),
			background,
		}
		defaultAttr = convertTextAttr(screenInfo.WAttributes)
	}
}

type coord struct {
	X, Y int16
}

type smallRect struct {
	Left, Top, Right, Bottom int16
}

type consoleScreenBufferInfo struct {
	DwSize              coord
	DwCursorPosition    coord
	WAttributes         uint16
	SrWindow            smallRect
	DwMaximumWindowSize coord
}

func getConsoleScreenBufferInfo(hConsoleOutput uintptr) *consoleScreenBufferInfo {
	var csbi consoleScreenBufferInfo
	ret, _, _ := procGetConsoleScreenBufferInfo.Call(
		hConsoleOutput,
		uintptr(unsafe.Pointer(&csbi)))
	if ret == 0 {
		return nil
	}
	return &csbi
}

func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
	ret, _, _ := procSetConsoleTextAttribute.Call(
		hConsoleOutput,
		uintptr(wAttributes))
	return ret != 0
}

type textAttributes struct {
	foregroundColor     uint16
	backgroundColor     uint16
	foregroundIntensity uint16
	backgroundIntensity uint16
	underscore          uint16
	otherAttributes     uint16
}

func convertTextAttr(winAttr uint16) *textAttributes {
	fgColor := winAttr & (foregroundRed | foregroundGreen | foregroundBlue)
	bgColor := winAttr & (backgroundRed | backgroundGreen | backgroundBlue)
	fgIntensity := winAttr & foregroundIntensity
	bgIntensity := winAttr & backgroundIntensity
	underline := winAttr & underscore
	otherAttributes := winAttr &^ (foregroundMask | backgroundMask | underscore)
	return &textAttributes{fgColor, bgColor, fgIntensity, bgIntensity, underline, otherAttributes}
}

func convertWinAttr(textAttr *textAttributes) uint16 {
	var winAttr uint16 = 0
	winAttr |= textAttr.foregroundColor
	winAttr |= textAttr.backgroundColor
	winAttr |= textAttr.foregroundIntensity
	winAttr |= textAttr.backgroundIntensity
	winAttr |= textAttr.underscore
	winAttr |= textAttr.otherAttributes
	return winAttr
}

func changeColor(param []byte) {
	if defaultAttr == nil {
		return
	}

	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		return
	}

	winAttr := convertTextAttr(screenInfo.WAttributes)
	strParam := string(param)
	if len(strParam) <= 0 {
		strParam = "0"
	}
	csiParam := strings.Split(strParam, string(separatorChar))
	for _, p := range csiParam {
		c, ok := colorMap[p]
		switch {
		case !ok:
			switch p {
			case ansiReset:
				winAttr.foregroundColor = defaultAttr.foregroundColor
				winAttr.backgroundColor = defaultAttr.backgroundColor
				winAttr.foregroundIntensity = defaultAttr.foregroundIntensity
				winAttr.backgroundIntensity = defaultAttr.backgroundIntensity
				winAttr.underscore = 0
				winAttr.otherAttributes = 0
			case ansiIntensityOn:
				winAttr.foregroundIntensity = foregroundIntensity
			case ansiIntensityOff:
				winAttr.foregroundIntensity = 0
			case ansiUnderlineOn:
				winAttr.underscore = underscore
			case ansiUnderlineOff:
				winAttr.underscore = 0
			case ansiBlinkOn:
				winAttr.backgroundIntensity = backgroundIntensity
			case ansiBlinkOff:
				winAttr.backgroundIntensity = 0
			default:
				// unknown code
			}
		case c.drawType == foreground:
			winAttr.foregroundColor = c.code
		case c.drawType == background:
			winAttr.backgroundColor = c.code
		}
	}
	winTextAttribute := convertWinAttr(winAttr)
	setConsoleTextAttribute(uintptr(syscall.Stdout), winTextAttribute)
}

func parseEscapeSequence(command byte, param []byte) {
	switch command {
	case sgrCode:
		changeColor(param)
	}
}

func isParameterChar(b byte) bool {
	return ('0' <= b && b <= '9') || b == separatorChar
}

func (cw *ansiColorWriter) Write(p []byte) (int, error) {
	r, nw, nc, first, last := 0, 0, 0, 0, 0
	var err error
	for i, ch := range p {
		switch cw.state {
		case outsideCsiCode:
			if ch == firstCsiChar {
				nc++
				cw.state = firstCsiCode
			}
		case firstCsiCode:
			switch ch {
			case firstCsiChar:
				nc++
				break
			case secondeCsiChar:
				nc++
				cw.state = secondeCsiCode
				last = i - 1
			default:
				cw.state = outsideCsiCode
			}
		case secondeCsiCode:
			nc++
			if isParameterChar(ch) {
				cw.paramBuf.WriteByte(ch)
			} else {
				nw, err = cw.w.Write(p[first:last])
				r += nw
				if err != nil {
					return r, err
				}
				first = i + 1
				param := cw.paramBuf.Bytes()
				cw.paramBuf.Reset()
				parseEscapeSequence(ch, param)
				cw.state = outsideCsiCode
			}
		default:
			cw.state = outsideCsiCode
		}
	}

	if cw.state == outsideCsiCode {
		nw, err = cw.w.Write(p[first:len(p)])
	}

	return r + nw + nc, err
}
<commit_msg>Fix miss color<commit_after>// Copyright 2014 shiena Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// +build windows

package ansicolor

import (
	"bytes"
	"io"
	"strings"
	"syscall"
	"unsafe"
)

type csiState int

const (
	outsideCsiCode csiState = iota
	firstCsiCode
	secondeCsiCode
)

type ansiColorWriter struct {
	w        io.Writer
	state    csiState
	paramBuf bytes.Buffer
}

const (
	firstCsiChar   byte = '\x1b'
	secondeCsiChar byte = '['
	separatorChar  byte = ';'
	sgrCode        byte = 'm'
)

const (
	foregroundBlue      = uint16(0x0001)
	foregroundGreen     = uint16(0x0002)
	foregroundRed       = uint16(0x0004)
	foregroundIntensity = uint16(0x0008)
	backgroundBlue      = uint16(0x0010)
	backgroundGreen     = uint16(0x0020)
	backgroundRed       = uint16(0x0040)
	backgroundIntensity = uint16(0x0080)
	underscore          = uint16(0x8000)

	foregroundMask = foregroundBlue | foregroundGreen | foregroundRed | foregroundIntensity
	backgroundMask = backgroundBlue | backgroundGreen | backgroundRed | backgroundIntensity
)

const (
	ansiReset        = "0"
	ansiIntensityOn  = "1"
	ansiIntensityOff = "21"
	ansiUnderlineOn  = "4"
	ansiUnderlineOff = "24"
	ansiBlinkOn      = "5"
	ansiBlinkOff     = "25"

	ansiForegroundBlack   = "30"
	ansiForegroundRed     = "31"
	ansiForegroundGreen   = "32"
	ansiForegroundYellow  = "33"
	ansiForegroundBlue    = "34"
	ansiForegroundMagenta = "35"
	ansiForegroundCyan    = "36"
	ansiForegroundWhite   = "37"
	ansiForegroundDefault = "39"

	ansiBackgroundBlack   = "40"
	ansiBackgroundRed     = "41"
	ansiBackgroundGreen   = "42"
	ansiBackgroundYellow  = "43"
	ansiBackgroundBlue    = "44"
	ansiBackgroundMagenta = "45"
	ansiBackgroundCyan    = "46"
	ansiBackgroundWhite   = "47"
	ansiBackgroundDefault = "49"

	ansiLightForegroundGray    = "90"
	ansiLightForegroundRed     = "91"
	ansiLightForegroundGreen   = "92"
	ansiLightForegroundYellow  = "93"
	ansiLightForegroundBlue    = "94"
	ansiLightForegroundMagenta = "95"
	ansiLightForegroundCyan    = "96"
	ansiLightForegroundWhite   = "97"
	ansiLightForegroundDefault = "99"
)

type drawType int

const (
	foreground drawType = iota
	background
)

type winColor struct {
	code     uint16
	drawType drawType
}

var colorMap = map[string]winColor{
	ansiForegroundBlack:   {0, foreground},
	ansiForegroundRed:     {foregroundRed, foreground},
	ansiForegroundGreen:   {foregroundGreen, foreground},
	ansiForegroundYellow:  {foregroundRed | foregroundGreen, foreground},
	ansiForegroundBlue:    {foregroundBlue, foreground},
	ansiForegroundMagenta: {foregroundRed | foregroundBlue, foreground},
	ansiForegroundCyan:    {foregroundGreen | foregroundBlue, foreground},
	ansiForegroundWhite:   {foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiForegroundDefault: {foregroundRed | foregroundGreen | foregroundBlue, foreground},

	ansiBackgroundBlack:   {0, background},
	ansiBackgroundRed:     {backgroundRed, background},
	ansiBackgroundGreen:   {backgroundGreen, background},
	ansiBackgroundYellow:  {backgroundRed | backgroundGreen, background},
	ansiBackgroundBlue:    {backgroundBlue, background},
	ansiBackgroundMagenta: {backgroundRed | backgroundBlue, background},
	ansiBackgroundCyan:    {backgroundGreen | backgroundBlue, background},
	ansiBackgroundWhite:   {backgroundRed | backgroundGreen | backgroundBlue, background},
	ansiBackgroundDefault: {0, background},

	ansiLightForegroundGray:    {foregroundIntensity, foreground},
	ansiLightForegroundRed:     {foregroundIntensity | foregroundRed, foreground},
	ansiLightForegroundGreen:   {foregroundIntensity | foregroundGreen, foreground},
	ansiLightForegroundYellow:  {foregroundIntensity | foregroundRed | foregroundGreen, foreground},
	ansiLightForegroundBlue:    {foregroundIntensity | foregroundBlue, foreground},
	ansiLightForegroundMagenta: {foregroundIntensity | foregroundRed | foregroundBlue, foreground},
	ansiLightForegroundCyan:    {foregroundIntensity | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundWhite:   {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},
	ansiLightForegroundDefault: {foregroundIntensity | foregroundRed | foregroundGreen | foregroundBlue, foreground},
}

var (
	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
	procSetConsoleTextAttribute    = kernel32.NewProc("SetConsoleTextAttribute")
	procGetConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")
	defaultAttr                    *textAttributes
)

func init() {
	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo != nil {
		colorMap[ansiForegroundDefault] = winColor{
			screenInfo.WAttributes & (foregroundRed | foregroundGreen | foregroundBlue),
			foreground,
		}
		colorMap[ansiBackgroundDefault] = winColor{
			screenInfo.WAttributes & (backgroundRed | backgroundGreen | backgroundBlue),
			background,
		}
		defaultAttr = convertTextAttr(screenInfo.WAttributes)
	}
}

type coord struct {
	X, Y int16
}

type smallRect struct {
	Left, Top, Right, Bottom int16
}

type consoleScreenBufferInfo struct {
	DwSize              coord
	DwCursorPosition    coord
	WAttributes         uint16
	SrWindow            smallRect
	DwMaximumWindowSize coord
}

func getConsoleScreenBufferInfo(hConsoleOutput uintptr) *consoleScreenBufferInfo {
	var csbi consoleScreenBufferInfo
	ret, _, _ := procGetConsoleScreenBufferInfo.Call(
		hConsoleOutput,
		uintptr(unsafe.Pointer(&csbi)))
	if ret == 0 {
		return nil
	}
	return &csbi
}

func setConsoleTextAttribute(hConsoleOutput uintptr, wAttributes uint16) bool {
	ret, _, _ := procSetConsoleTextAttribute.Call(
		hConsoleOutput,
		uintptr(wAttributes))
	return ret != 0
}

type textAttributes struct {
	foregroundColor     uint16
	backgroundColor     uint16
	foregroundIntensity uint16
	backgroundIntensity uint16
	underscore          uint16
	otherAttributes     uint16
}

func convertTextAttr(winAttr uint16) *textAttributes {
	fgColor := winAttr & (foregroundRed | foregroundGreen | foregroundBlue)
	bgColor := winAttr & (backgroundRed | backgroundGreen | backgroundBlue)
	fgIntensity := winAttr & foregroundIntensity
	bgIntensity := winAttr & backgroundIntensity
	underline := winAttr & underscore
	otherAttributes := winAttr &^ (foregroundMask | backgroundMask | underscore)
	return &textAttributes{fgColor, bgColor, fgIntensity, bgIntensity, underline, otherAttributes}
}

func convertWinAttr(textAttr *textAttributes) uint16 {
	var winAttr uint16 = 0
	winAttr |= textAttr.foregroundColor
	winAttr |= textAttr.backgroundColor
	winAttr |= textAttr.foregroundIntensity
	winAttr |= textAttr.backgroundIntensity
	winAttr |= textAttr.underscore
	winAttr |= textAttr.otherAttributes
	return winAttr
}

func changeColor(param []byte) {
	if defaultAttr == nil {
		return
	}

	screenInfo := getConsoleScreenBufferInfo(uintptr(syscall.Stdout))
	if screenInfo == nil {
		return
	}

	winAttr := convertTextAttr(screenInfo.WAttributes)
	strParam := string(param)
	if len(strParam) <= 0 {
		strParam = "0"
	}
	csiParam := strings.Split(strParam, string(separatorChar))
	for _, p := range csiParam {
		c, ok := colorMap[p]
		switch {
		case !ok:
			switch p {
			case ansiReset:
				winAttr.foregroundColor = defaultAttr.foregroundColor
				winAttr.backgroundColor = defaultAttr.backgroundColor
				winAttr.foregroundIntensity = defaultAttr.foregroundIntensity
				winAttr.backgroundIntensity = defaultAttr.backgroundIntensity
				winAttr.underscore = 0
				winAttr.otherAttributes = 0
			case ansiIntensityOn:
				winAttr.foregroundIntensity = foregroundIntensity
			case ansiIntensityOff:
				winAttr.foregroundIntensity = 0
			case ansiUnderlineOn:
				winAttr.underscore = underscore
			case ansiUnderlineOff:
				winAttr.underscore = 0
			case ansiBlinkOn:
				winAttr.backgroundIntensity = backgroundIntensity
			case ansiBlinkOff:
				winAttr.backgroundIntensity = 0
			default:
				// unknown code
			}
		case c.drawType == foreground:
			winAttr.foregroundColor = c.code
		case c.drawType == background:
			winAttr.backgroundColor = c.code
		}
	}
	winTextAttribute := convertWinAttr(winAttr)
	setConsoleTextAttribute(uintptr(syscall.Stdout), winTextAttribute)
}

func parseEscapeSequence(command byte, param []byte) {
	switch command {
	case sgrCode:
		changeColor(param)
	}
}

func isParameterChar(b byte) bool {
	return ('0' <= b && b <= '9') || b == separatorChar
}

func (cw *ansiColorWriter) Write(p []byte) (int, error) {
	r, nw, nc, first, last := 0, 0, 0, 0, 0
	var err error
	for i, ch := range p {
		switch cw.state {
		case outsideCsiCode:
			if ch == firstCsiChar {
				nc++
				cw.state = firstCsiCode
			}
		case firstCsiCode:
			switch ch {
			case firstCsiChar:
				nc++
				break
			case secondeCsiChar:
				nc++
				cw.state = secondeCsiCode
				last = i - 1
			default:
				cw.state = outsideCsiCode
			}
		case secondeCsiCode:
			nc++
			if isParameterChar(ch) {
				cw.paramBuf.WriteByte(ch)
			} else {
				nw, err = cw.w.Write(p[first:last])
				r += nw
				if err != nil {
					return r, err
				}
				first = i + 1
				param := cw.paramBuf.Bytes()
				cw.paramBuf.Reset()
				parseEscapeSequence(ch, param)
				cw.state = outsideCsiCode
			}
		default:
			cw.state = outsideCsiCode
		}
	}

	if cw.state == outsideCsiCode {
		nw, err = cw.w.Write(p[first:len(p)])
	}

	return r + nw + nc, err
}
<|endoftext|>


<commit_before>package bruxism

import (
	"errors"
	"strings"
	"sync"
	"time"

	"google.golang.org/api/youtube/v3"
)

// YTLiveChannel is a monitor that will send new live videos to a provided channel.
type YTLiveChannel struct {
	sync.RWMutex
	service *youtube.Service
	// map channelID -> chan
	liveVideoChans map[string][]chan *youtube.Video
	channelNames   map[string]string
}

func NewYTLiveChannel(service *youtube.Service) *YTLiveChannel {
	return &YTLiveChannel{service: service}
}

func (y *YTLiveChannel) Monitor(channel string, liveVideoChan chan *youtube.Video) error {
	y.Lock()
	defer y.Unlock()

	if y.channelNames[channel] == "" {
		clr, err := y.service.Channels.List("snippet").Id(channel).Do()
		if err != nil {
			return errors.New("Error loading channel.")
		}
		if len(clr.Items) != 1 {
			return errors.New("No channel found.")
		}
		if y.channelNames == nil {
			y.channelNames = map[string]string{}
		}
		y.channelNames[channel] = clr.Items[0].Snippet.Title
	}

	if y.liveVideoChans == nil {
		y.liveVideoChans = map[string][]chan *youtube.Video{}
	}
	created := len(y.liveVideoChans[channel]) == 0
	y.liveVideoChans[channel] = append(y.liveVideoChans[channel], liveVideoChan)
	if created {
		go y.poll(channel)
	}
	return nil
}

func (y *YTLiveChannel) ChannelName(channel string) string {
	y.RLock()
	defer y.RUnlock()

	return y.channelNames[channel]
}

func (y *YTLiveChannel) poll(channel string) {
	var lastAnnounce time.Time
	seen := map[string]time.Time{}
	now := time.Now()
	first := true
	for {
		videos, _ := y.getLiveVideos(channel)
		for _, v := range videos {
			if now.After(seen[v.Id].Add(6 * time.Hour)) {
				seen[v.Id] = now
				// Don't announce the videos that are already live.
				if first {
					continue
				}
				// Don't allow more than 1 announcement per hour.
				if !now.After(lastAnnounce.Add(1 * time.Hour)) {
					continue
				}
				lastAnnounce = now
				y.RLock()
				for _, c := range y.liveVideoChans[channel] {
					c <- v
				}
				y.RUnlock()
			}
		}
		first = false
		<-time.After(5 * time.Minute)
	}

}

func (y *YTLiveChannel) getLiveVideos(channel string) (map[string]*youtube.Video, error) {
	if y.service == nil {
		return nil, errors.New("Service not available.")
	}

	search, err := y.service.Search.List("id").ChannelId(channel).EventType("live").Type("video").Do()
	if err != nil {
		return nil, err
	}

	ids := []string{}
	for _, searchResult := range search.Items {
		ids = append(ids, searchResult.Id.VideoId)
	}

	m := map[string]*youtube.Video{}

	i := 0
	for i < len(ids) {
		next := i + 50
		if next >= len(ids) {
			next = len(ids)
		}
		videoList, err := y.service.Videos.List("id,snippet,liveStreamingDetails").MaxResults(50).Id(strings.Join(ids[i:next], ",")).Do()
		if err != nil {
			return nil, err
		}

		for _, v := range videoList.Items {
			if v.LiveStreamingDetails.ActualEndTime == "" {
				m[v.Id] = v
			}
		}

		i = next
	}

	return m, nil
}
<commit_msg>Maybe fix live plugin.<commit_after>package bruxism

import (
	"errors"
	"strings"
	"sync"
	"time"

	"google.golang.org/api/youtube/v3"
)

// YTLiveChannel is a monitor that will send new live videos to a provided channel.
type YTLiveChannel struct {
	sync.RWMutex
	service *youtube.Service
	// map channelID -> chan
	liveVideoChans map[string][]chan *youtube.Video
	channelNames   map[string]string
}

func NewYTLiveChannel(service *youtube.Service) *YTLiveChannel {
	return &YTLiveChannel{service: service}
}

func (y *YTLiveChannel) Monitor(channel string, liveVideoChan chan *youtube.Video) error {
	y.Lock()
	defer y.Unlock()

	if y.channelNames[channel] == "" {
		clr, err := y.service.Channels.List("snippet").Id(channel).Do()
		if err != nil {
			return errors.New("Error loading channel.")
		}
		if len(clr.Items) != 1 {
			return errors.New("No channel found.")
		}
		if y.channelNames == nil {
			y.channelNames = map[string]string{}
		}
		y.channelNames[channel] = clr.Items[0].Snippet.Title
	}

	if y.liveVideoChans == nil {
		y.liveVideoChans = map[string][]chan *youtube.Video{}
	}
	created := len(y.liveVideoChans[channel]) == 0
	y.liveVideoChans[channel] = append(y.liveVideoChans[channel], liveVideoChan)
	if created {
		go y.poll(channel)
	}
	return nil
}

func (y *YTLiveChannel) ChannelName(channel string) string {
	y.RLock()
	defer y.RUnlock()

	return y.channelNames[channel]
}

func (y *YTLiveChannel) poll(channel string) {
	var lastAnnounce time.Time
	seen := map[string]bool{}
	now := time.Now()
	first := true
	for {
		videos, _ := y.getLiveVideos(channel)
		for k, v := range videos {
			if !seen[k] {
				seen[k] = true
				// Don't announce the videos that are already live.
				if first {
					continue
				}
				// Don't allow more than 1 announcement per hour.
				if !now.After(lastAnnounce.Add(1 * time.Hour)) {
					continue
				}
				lastAnnounce = now
				y.RLock()
				for _, c := range y.liveVideoChans[channel] {
					c <- v
				}
				y.RUnlock()
			}
		}
		first = false
		<-time.After(5 * time.Minute)
	}

}

func (y *YTLiveChannel) getLiveVideos(channel string) (map[string]*youtube.Video, error) {
	if y.service == nil {
		return nil, errors.New("Service not available.")
	}

	search, err := y.service.Search.List("id").ChannelId(channel).EventType("live").Type("video").Do()
	if err != nil {
		return nil, err
	}

	ids := []string{}
	for _, searchResult := range search.Items {
		ids = append(ids, searchResult.Id.VideoId)
	}

	m := map[string]*youtube.Video{}

	i := 0
	for i < len(ids) {
		next := i + 50
		if next >= len(ids) {
			next = len(ids)
		}
		videoList, err := y.service.Videos.List("id,snippet,liveStreamingDetails").MaxResults(50).Id(strings.Join(ids[i:next], ",")).Do()
		if err != nil {
			return nil, err
		}

		for _, v := range videoList.Items {
			if v.LiveStreamingDetails.ActualEndTime == "" {
				m[v.Id] = v
			}
		}

		i = next
	}

	return m, nil
}
<|endoftext|>


<commit_before>/*
	cryptokit is used to authenticate and secure messages using
	public-key cryptography. It provides an interface similar to NaCL,
	but uses ECIES using ephemeral ECDH for shared keys, and secret
	box for securing messages.

	Messages should be secured using the Seal function, and recovered
	using the Open function. A box (or authenticated and encrypted
	message) will be Overhead bytes longer than the message it
	came from; this package will not obscure the length of the
	message. Keys, if they are not generated using the GenerateKey
	function, should be KeySize bytes long. The KeyIsSuitable function
	may be used to test a key is the proper length.

	This package also provides signed boxes: these digitally sign the
	message before sealing them, and the signature can be checked
	on opening. These must be opened with the OpenSigned function,
	and use ECDSA for signatures.

	The boxes used in this package are suitable for 20-year security.
*/
package cryptokit

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"github.com/gokyle/cryptokit/secretbox"
	"math/big"
)

type PublicKey []byte
type PrivateKey []byte

const (
	publicKeySize  = 65
	privateKeySize = 32
	sigSize        = 64
)

const (
	SharedKeySize  = 48
	ecdhSharedSize = 32
)

// Overhead is the number of bytes of overhead when boxing a message.
var Overhead = publicKeySize + secretbox.Overhead

// SignedOverhead is the number of bytes of overhead when signing and
// boxing a message.
var Overhead = publicKeySize + secretbox.Overhead + sigSize

// The default source for random data is the crypto/rand package's Reader.
var PRNG = rand.Reader

var curve = elliptic.P256()

// ecdh performs the ECDH key agreement method to generate a shared key
// between a pair of keys.
func ecdh(key PrivateKey, peer PublicKey) ([]byte, bool) {
	x, y := elliptic.Unmarshal(curve, peer)
	if x == nil {
		fmt.Println("failed to unmarshal key")
		return nil, false
	}
	x, _ = curve.ScalarMult(x, y, key)
	if x == nil {
		return nil, false
	}
	xb := zeroPad(x.Bytes(), ecdhSharedSize)

	skey := xb[:16]
	mkey := xb[16:]
	h := sha256.New()
	h.Write(mkey)
	mkey = h.Sum(nil)

	return append(skey, mkey...), true
}

//
func GenerateKey() (PrivateKey, PublicKey, bool) {
	key, x, y, err := elliptic.GenerateKey(curve, PRNG)
	if err != nil {
		return nil, nil, false
	}
	peer := elliptic.Marshal(curve, x, y)
	if peer == nil {
	}
	if len(key) != privateKeySize || len(peer) != publicKeySize {
		return nil, nil, false
	}
	return key, peer, true
}

// Seal returns an authenticated and encrypted message, and a boolean
// indicating whether the sealing operation was successful. If it returns
// true, the message was successfully sealed. The box will be Overhead
// bytes longer than the message. These boxes are not dependent on having
// a private key.
func Seal(message []byte, peer PublicKey) (box []byte, ok bool) {
	if !KeyIsSuitable(nil, peer) {
		return
	}

	eph_key, eph_peer, ok := GenerateKey()
	if !ok {
		fmt.Println("failed to generate ephem key")
		return
	}

	skey, ok := ecdh(eph_key, peer)
	if !ok {
		fmt.Println("failed to generate shared key")
		return
	}

	sbox, ok := secretbox.Seal(message, skey)
	if !ok {
		fmt.Println("failed to seal SecretBox")
		return
	}

	box = make([]byte, publicKeySize+len(sbox))
	copy(box, eph_peer)
	copy(box[publicKeySize:], sbox)
	return box, true
}

// Open authenticates and decrypts a sealed message, also returning
// whether the message was successfully opened. If this is false, the
// message must be discarded. The returned message will be Overhead
// bytes shorter than the box.
func Open(box []byte, key PrivateKey) (message []byte, ok bool) {
	if !KeyIsSuitable(key, nil) {
		return
	}

	if len(box) < publicKeySize+secretbox.Overhead {
		fmt.Println("box size is invalid")
		return
	}

	eph_peer := box[:publicKeySize]
	shared, ok := ecdh(key, eph_peer)
	if !ok {
		fmt.Println("couldn't generate shared key")
		return
	}

	message, ok = secretbox.Open(box[publicKeySize:], shared)
	return
}

func ecdsa_private(key PrivateKey, pub PublicKey) (skey *ecdsa.PrivateKey, ok bool) {
	x, y := elliptic.Unmarshal(curve, pub)
	if x == nil {
		return
	}

	skey = new(ecdsa.PrivateKey)
	skey.D = new(big.Int).SetBytes(key)
	skey.PublicKey.Curve = curve
	skey.X = x
	skey.Y = y
	ok = true
	return
}

func ecdsa_public(peer PublicKey) (pkey *ecdsa.PublicKey, ok bool) {
	x, y := elliptic.Unmarshal(curve, peer)
	if x == nil {
		return
	}
	pkey = &ecdsa.PublicKey{
		Curve: curve,
		X:     x,
		Y:     y,
	}
	return pkey, true
}

func sign(message []byte, key PrivateKey, pub PublicKey) (smessage []byte, ok bool) {
	h := sha256.New()
	h.Write(message)
	hash := h.Sum(nil)

	skey, ok := ecdsa_private(key, pub)
	if !ok {
		fmt.Println("failed to generate ecdsa key")
		return
	}
	r, s, err := ecdsa.Sign(PRNG, skey, hash)
	if err == nil {
		smessage = make([]byte, len(message)+64)
		copy(smessage, message)
		sig := marshalECDSASignature(r, s)
		copy(smessage[len(message):], sig)
		ok = true
	}
	return
}

func verify(smessage []byte, peer PublicKey) bool {
	if len(smessage) <= sigSize {
		return false
	}
	sigPos := len(smessage) - sigSize
	message := smessage[:sigPos]
	sig := smessage[sigPos:]
	h := sha256.New()
	h.Write(message)

	pub, ok := ecdsa_public(peer)
	if !ok {
		return false
	}
	r, s := unmarshalECDSASignature(sig)
	if r == nil {
		return false
	}
	return ecdsa.Verify(pub, h.Sum(nil), r, s)
}

func marshalECDSASignature(r, s *big.Int) []byte {
	if r == nil || s == nil {
		return make([]byte, sigSize)
	}
	sig := make([]byte, sigSize)
	rb := r.Bytes()
	rb = zeroPad(rb, 32)
	sb := s.Bytes()
	sb = zeroPad(sb, 32)
	copy(sig, rb)
	copy(sig[32:], sb)
	return sig
}

func unmarshalECDSASignature(sig []byte) (r, s *big.Int) {
	if len(sig) != sigSize {
		return
	}
	r = new(big.Int).SetBytes(sig[:32])
	s = new(big.Int).SetBytes(sig[32:])
	return
}

// SignAndSeal adds a digital signature to the message before sealing it.
func SignAndSeal(message []byte, key PrivateKey, public PublicKey, peer PublicKey) (box []byte, ok bool) {
	smessage, ok := sign(message, key, public)
	if !ok {
		return
	}
	box, ok = Seal(smessage, peer)
	return
}

// OpenSigned opens a signed box, and verifies the signature. If the box
// couldn't be opened or the signature is invalid, OpenSigned returns false,
// and the message value must be discarded.
func OpenSigned(box []byte, key PrivateKey, peer PublicKey) (message []byte, ok bool) {
	smessage, ok := Open(box, key)
	if !ok {
		return
	}

	ok = verify(smessage, peer)
	message = smessage[:len(smessage)-sigSize]
	return
}

// zeroPad returns a new slice of length size. The contents of input are right
// aligned in the new slice.
func zeroPad(in []byte, outlen int) (out []byte) {
	var inLen int
	if inLen = len(in); inLen > outlen {
		inLen = outlen
	} else if inLen == outlen {
		return in
	}
	start := outlen - inLen - 1
	out = make([]byte, outlen)
	copy(out[start:], in)
	return
}

// IsKeySuitable takes a private and/or public key, and returns true if
// all keys passed in are valid. If no key is passed in, or any key passed
// in is invalid, it will return false.
func KeyIsSuitable(key PrivateKey, pub PublicKey) bool {
	if key == nil && pub == nil {
		return false
	} else if key != nil && len(key) != privateKeySize {
		return false
	} else if pub != nil && len(pub) != publicKeySize {
		return false
	}
	return true
}
<commit_msg>Fix typo'd variable.<commit_after>/*
	cryptokit is used to authenticate and secure messages using
	public-key cryptography. It provides an interface similar to NaCL,
	but uses ECIES using ephemeral ECDH for shared keys, and secret
	box for securing messages.

	Messages should be secured using the Seal function, and recovered
	using the Open function. A box (or authenticated and encrypted
	message) will be Overhead bytes longer than the message it
	came from; this package will not obscure the length of the
	message. Keys, if they are not generated using the GenerateKey
	function, should be KeySize bytes long. The KeyIsSuitable function
	may be used to test a key is the proper length.

	This package also provides signed boxes: these digitally sign the
	message before sealing them, and the signature can be checked
	on opening. These must be opened with the OpenSigned function,
	and use ECDSA for signatures.

	The boxes used in this package are suitable for 20-year security.
*/
package cryptokit

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"github.com/gokyle/cryptokit/secretbox"
	"math/big"
)

type PublicKey []byte
type PrivateKey []byte

const (
	publicKeySize  = 65
	privateKeySize = 32
	sigSize        = 64
)

const (
	SharedKeySize  = 48
	ecdhSharedSize = 32
)

// Overhead is the number of bytes of overhead when boxing a message.
var Overhead = publicKeySize + secretbox.Overhead

// SignedOverhead is the number of bytes of overhead when signing and
// boxing a message.
var SignedOverhead = publicKeySize + secretbox.Overhead + sigSize

// The default source for random data is the crypto/rand package's Reader.
var PRNG = rand.Reader

var curve = elliptic.P256()

// ecdh performs the ECDH key agreement method to generate a shared key
// between a pair of keys.
func ecdh(key PrivateKey, peer PublicKey) ([]byte, bool) {
	x, y := elliptic.Unmarshal(curve, peer)
	if x == nil {
		fmt.Println("failed to unmarshal key")
		return nil, false
	}
	x, _ = curve.ScalarMult(x, y, key)
	if x == nil {
		return nil, false
	}
	xb := zeroPad(x.Bytes(), ecdhSharedSize)

	skey := xb[:16]
	mkey := xb[16:]
	h := sha256.New()
	h.Write(mkey)
	mkey = h.Sum(nil)

	return append(skey, mkey...), true
}

//
func GenerateKey() (PrivateKey, PublicKey, bool) {
	key, x, y, err := elliptic.GenerateKey(curve, PRNG)
	if err != nil {
		return nil, nil, false
	}
	peer := elliptic.Marshal(curve, x, y)
	if peer == nil {
	}
	if len(key) != privateKeySize || len(peer) != publicKeySize {
		return nil, nil, false
	}
	return key, peer, true
}

// Seal returns an authenticated and encrypted message, and a boolean
// indicating whether the sealing operation was successful. If it returns
// true, the message was successfully sealed. The box will be Overhead
// bytes longer than the message. These boxes are not dependent on having
// a private key.
func Seal(message []byte, peer PublicKey) (box []byte, ok bool) {
	if !KeyIsSuitable(nil, peer) {
		return
	}

	eph_key, eph_peer, ok := GenerateKey()
	if !ok {
		fmt.Println("failed to generate ephem key")
		return
	}

	skey, ok := ecdh(eph_key, peer)
	if !ok {
		fmt.Println("failed to generate shared key")
		return
	}

	sbox, ok := secretbox.Seal(message, skey)
	if !ok {
		fmt.Println("failed to seal SecretBox")
		return
	}

	box = make([]byte, publicKeySize+len(sbox))
	copy(box, eph_peer)
	copy(box[publicKeySize:], sbox)
	return box, true
}

// Open authenticates and decrypts a sealed message, also returning
// whether the message was successfully opened. If this is false, the
// message must be discarded. The returned message will be Overhead
// bytes shorter than the box.
func Open(box []byte, key PrivateKey) (message []byte, ok bool) {
	if !KeyIsSuitable(key, nil) {
		return
	}

	if len(box) < publicKeySize+secretbox.Overhead {
		fmt.Println("box size is invalid")
		return
	}

	eph_peer := box[:publicKeySize]
	shared, ok := ecdh(key, eph_peer)
	if !ok {
		fmt.Println("couldn't generate shared key")
		return
	}

	message, ok = secretbox.Open(box[publicKeySize:], shared)
	return
}

func ecdsa_private(key PrivateKey, pub PublicKey) (skey *ecdsa.PrivateKey, ok bool) {
	x, y := elliptic.Unmarshal(curve, pub)
	if x == nil {
		return
	}

	skey = new(ecdsa.PrivateKey)
	skey.D = new(big.Int).SetBytes(key)
	skey.PublicKey.Curve = curve
	skey.X = x
	skey.Y = y
	ok = true
	return
}

func ecdsa_public(peer PublicKey) (pkey *ecdsa.PublicKey, ok bool) {
	x, y := elliptic.Unmarshal(curve, peer)
	if x == nil {
		return
	}
	pkey = &ecdsa.PublicKey{
		Curve: curve,
		X:     x,
		Y:     y,
	}
	return pkey, true
}

func sign(message []byte, key PrivateKey, pub PublicKey) (smessage []byte, ok bool) {
	h := sha256.New()
	h.Write(message)
	hash := h.Sum(nil)

	skey, ok := ecdsa_private(key, pub)
	if !ok {
		fmt.Println("failed to generate ecdsa key")
		return
	}
	r, s, err := ecdsa.Sign(PRNG, skey, hash)
	if err == nil {
		smessage = make([]byte, len(message)+64)
		copy(smessage, message)
		sig := marshalECDSASignature(r, s)
		copy(smessage[len(message):], sig)
		ok = true
	}
	return
}

func verify(smessage []byte, peer PublicKey) bool {
	if len(smessage) <= sigSize {
		return false
	}
	sigPos := len(smessage) - sigSize
	message := smessage[:sigPos]
	sig := smessage[sigPos:]
	h := sha256.New()
	h.Write(message)

	pub, ok := ecdsa_public(peer)
	if !ok {
		return false
	}
	r, s := unmarshalECDSASignature(sig)
	if r == nil {
		return false
	}
	return ecdsa.Verify(pub, h.Sum(nil), r, s)
}

func marshalECDSASignature(r, s *big.Int) []byte {
	if r == nil || s == nil {
		return make([]byte, sigSize)
	}
	sig := make([]byte, sigSize)
	rb := r.Bytes()
	rb = zeroPad(rb, 32)
	sb := s.Bytes()
	sb = zeroPad(sb, 32)
	copy(sig, rb)
	copy(sig[32:], sb)
	return sig
}

func unmarshalECDSASignature(sig []byte) (r, s *big.Int) {
	if len(sig) != sigSize {
		return
	}
	r = new(big.Int).SetBytes(sig[:32])
	s = new(big.Int).SetBytes(sig[32:])
	return
}

// SignAndSeal adds a digital signature to the message before sealing it.
func SignAndSeal(message []byte, key PrivateKey, public PublicKey, peer PublicKey) (box []byte, ok bool) {
	smessage, ok := sign(message, key, public)
	if !ok {
		return
	}
	box, ok = Seal(smessage, peer)
	return
}

// OpenSigned opens a signed box, and verifies the signature. If the box
// couldn't be opened or the signature is invalid, OpenSigned returns false,
// and the message value must be discarded.
func OpenSigned(box []byte, key PrivateKey, peer PublicKey) (message []byte, ok bool) {
	smessage, ok := Open(box, key)
	if !ok {
		return
	}

	ok = verify(smessage, peer)
	message = smessage[:len(smessage)-sigSize]
	return
}

// zeroPad returns a new slice of length size. The contents of input are right
// aligned in the new slice.
func zeroPad(in []byte, outlen int) (out []byte) {
	var inLen int
	if inLen = len(in); inLen > outlen {
		inLen = outlen
	} else if inLen == outlen {
		return in
	}
	start := outlen - inLen - 1
	out = make([]byte, outlen)
	copy(out[start:], in)
	return
}

// IsKeySuitable takes a private and/or public key, and returns true if
// all keys passed in are valid. If no key is passed in, or any key passed
// in is invalid, it will return false.
func KeyIsSuitable(key PrivateKey, pub PublicKey) bool {
	if key == nil && pub == nil {
		return false
	} else if key != nil && len(key) != privateKeySize {
		return false
	} else if pub != nil && len(pub) != publicKeySize {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package discovery

import (
	"github.com/gkarlik/quark-go/service"
	lb "github.com/gkarlik/quark-go/service/loadbalancer"
)

// Option represents function which is used to apply service discovery options.
type Option func(*Options)

// Options represents service discovery options.
type Options struct {
	Info     service.Info             // service info
	Strategy lb.LoadBalancingStrategy // load balancing strategy
}

// ByInfo allows to discover service by its info metadata.
func ByInfo(i service.Info) Option {
	return func(opts *Options) {
		opts.Info = i
	}
}

// WithInfo allows to register service by its info metadata.
func WithInfo(i service.Info) Option {
	return ByInfo(i)
}

// ByName allows to discover service by its name.
func ByName(name string) Option {
	return func(opts *Options) {
		opts.Info.Name = name
	}
}

// ByTag allows to discover service by its tag(s).
func ByTag(tag string) Option {
	return func(opts *Options) {
		opts.Info.Tags = []string{tag}
	}
}

// UsingLBStrategy allows to discover service using specified load balancing strategy.
func UsingLBStrategy(s lb.LoadBalancingStrategy) Option {
	return func(opts *Options) {
		opts.Strategy = s
	}
}
<commit_msg>Discovery options: fix + extension.<commit_after>package discovery

import (
	"github.com/gkarlik/quark-go/service"
	lb "github.com/gkarlik/quark-go/service/loadbalancer"
)

// Option represents function which is used to apply service discovery options.
type Option func(*Options)

// Options represents service discovery options.
type Options struct {
	Info     service.Info             // service info
	Strategy lb.LoadBalancingStrategy // load balancing strategy
}

// ByInfo allows to discover service by its info metadata.
func ByInfo(i service.Info) Option {
	return func(opts *Options) {
		opts.Info = i
	}
}

// WithInfo allows to register service by its info metadata.
func WithInfo(i service.Info) Option {
	return ByInfo(i)
}

// ByName allows to discover service by its name.
func ByName(name string) Option {
	return func(opts *Options) {
		opts.Info.Name = name
	}
}

// ByVersion allows to discover service by its version.
func ByVersion(version string) Option {
	return func(opts *Options) {
		opts.Info.Version = version
	}
}

// ByTag allows to discover service by its tag(s).
func ByTag(tag string) Option {
	return func(opts *Options) {
		opts.Info.Tags = append(opts.Info.Tags, tag)
	}
}

// UsingLBStrategy allows to discover service using specified load balancing strategy.
func UsingLBStrategy(s lb.LoadBalancingStrategy) Option {
	return func(opts *Options) {
		opts.Strategy = s
	}
}
<|endoftext|>


<commit_before>package logs

import (
	"bufio"
	"io"
	"net/url"
	"strconv"

	log "github.com/Sirupsen/logrus"

	"github.com/rancher/websocket-proxy/backend"
	"github.com/rancher/websocket-proxy/common"

	"bytes"
	"github.com/docker/distribution/context"
	"github.com/docker/docker/api/types"
	"github.com/rancher/agent/service/hostapi/auth"
	"github.com/rancher/agent/service/hostapi/events"
)

var (
	stdoutHead = []byte{1, 0, 0, 0}
	stderrHead = []byte{2, 0, 0, 0}
)

type Handler struct {
}

func (l *Handler) Handle(key string, initialMessage string, incomingMessages <-chan string, response chan<- common.Message) {
	defer backend.SignalHandlerClosed(key, response)

	requestURL, err := url.Parse(initialMessage)
	if err != nil {
		log.WithFields(log.Fields{"error": err, "url": initialMessage}).Error("Couldn't parse url.")
		return
	}
	tokenString := requestURL.Query().Get("token")
	token, valid := auth.GetAndCheckToken(tokenString)
	if !valid {
		return
	}

	logs := token.Claims["logs"].(map[string]interface{})
	container := logs["Container"].(string)
	follow, found := logs["Follow"].(bool)

	if !found {
		follow = true
	}

	tailTemp, found := logs["Lines"].(int)
	var tail string
	if found {
		tail = strconv.Itoa(int(tailTemp))
	} else {
		tail = "100"
	}

	client, err := events.NewDockerClient()
	if err != nil {
		log.WithFields(log.Fields{"error": err}).Error("Couldn't get docker client.")
		return
	}

	bothPrefix := "00 "
	stdoutPrefix := "01 "
	stderrPrefix := "02 "
	logOpts := types.ContainerLogsOptions{
		Follow:     follow,
		ShowStdout: true,
		ShowStderr: true,
		Timestamps: true,
		Tail:       tail,
	}

	stdoutReader, err := client.ContainerLogs(context.Background(), container, logOpts)
	if err != nil {
		log.Error(err)
		return
	}

	go func() {
		for {
			_, ok := <-incomingMessages
			if !ok {
				return
			}
		}
	}()

	go func(stdout io.ReadCloser) {
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			body := ""
			data := scanner.Bytes()
			if bytes.Contains(data, stdoutHead) {
				if len(data) > 8 {
					body = stdoutPrefix + string(data[8:])
				}
			} else if bytes.Contains(data, stderrHead) {
				if len(data) > 8 {
					body = stderrPrefix + string(data[8:])
				}
			} else {
				body = bothPrefix + string(data)
			}
			message := common.Message{
				Key:  key,
				Type: common.Body,
				Body: body,
			}
			response <- message
		}
		if err := scanner.Err(); err != nil {
			log.WithFields(log.Fields{"error": err}).Error("Error with the container log scanner.")
		}
		stdout.Close()
	}(stdoutReader)

	select {}
}
<commit_msg>Close logs when user cancels request<commit_after>package logs

import (
	"bufio"
	"bytes"
	"net/url"
	"strconv"
	"strings"

	log "github.com/Sirupsen/logrus"

	"github.com/rancher/websocket-proxy/backend"
	"github.com/rancher/websocket-proxy/common"

	"github.com/docker/docker/api/types"
	"github.com/rancher/agent/service/hostapi/auth"
	"github.com/rancher/agent/service/hostapi/events"
	"golang.org/x/net/context"
)

var (
	stdoutHead = []byte{1, 0, 0, 0}
	stderrHead = []byte{2, 0, 0, 0}
)

type Handler struct {
}

func (l *Handler) Handle(key string, initialMessage string, incomingMessages <-chan string, response chan<- common.Message) {
	defer backend.SignalHandlerClosed(key, response)

	requestURL, err := url.Parse(initialMessage)
	if err != nil {
		log.WithFields(log.Fields{"error": err, "url": initialMessage}).Error("Couldn't parse url.")
		return
	}
	tokenString := requestURL.Query().Get("token")
	token, valid := auth.GetAndCheckToken(tokenString)
	if !valid {
		return
	}

	logs := token.Claims["logs"].(map[string]interface{})
	container := logs["Container"].(string)
	follow, found := logs["Follow"].(bool)

	if !found {
		follow = true
	}

	tailTemp, found := logs["Lines"].(int)
	var tail string
	if found {
		tail = strconv.Itoa(int(tailTemp))
	} else {
		tail = "100"
	}

	client, err := events.NewDockerClient()
	if err != nil {
		log.WithFields(log.Fields{"error": err}).Error("Couldn't get docker client.")
		return
	}

	bothPrefix := "00 "
	stdoutPrefix := "01 "
	stderrPrefix := "02 "
	logOpts := types.ContainerLogsOptions{
		Follow:     follow,
		ShowStdout: true,
		ShowStderr: true,
		Timestamps: true,
		Tail:       tail,
	}

	ctx, cancelFnc := context.WithCancel(context.Background())
	stdout, err := client.ContainerLogs(ctx, container, logOpts)
	if err != nil {
		log.Error(err)
		return
	}
	defer stdout.Close()

	go func() {
		for {
			_, ok := <-incomingMessages
			if !ok {
				cancelFnc()
				return
			}
		}
	}()

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		body := ""
		data := scanner.Bytes()
		if bytes.Contains(data, stdoutHead) {
			if len(data) > 8 {
				body = stdoutPrefix + string(data[8:])
			}
		} else if bytes.Contains(data, stderrHead) {
			if len(data) > 8 {
				body = stderrPrefix + string(data[8:])
			}
		} else {
			body = bothPrefix + string(data)
		}
		message := common.Message{
			Key:  key,
			Type: common.Body,
			Body: body,
		}
		response <- message
	}
	if err := scanner.Err(); err != nil {
		// hacky, but can't do a type assertion on the cancellation error, which is the "normal" error received
		// when the logs are closed properly
		if !strings.Contains(err.Error(), "request canceled") {
			log.WithFields(log.Fields{"error": err}).Error("Error with the container log scanner.")
		}
	}
}
<|endoftext|>