<commit_before>// Copyright 2012 Kevin Gillette. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"flag"
	"fmt" // flag uses fmt anyway
	"io"
	"os"
)

func main() {
	var (
		input        string
		n, rep, size int64
		raw, b64     bool
		code         bool
		cols         uint
		str, esc     bool
	)

	flag.StringVar(&input, "i", "", "input file or stdin if '-' (default is urandom)")
	flag.Int64Var(&n, "n", 16, "number of (unencoded) bytes to output")
	flag.Int64Var(&rep, "r", 1, "number of repititions (output lines)")
	flag.BoolVar(&raw, "raw", false, "raw output")
	flag.BoolVar(&b64, "64", false, "base64 output")
	flag.BoolVar(&code, "code", false, "comma-separated hex literals (typical array syntax)")
	flag.UintVar(&cols, "cols", 12, "number of columns to use with -code")
	flag.BoolVar(&str, "str", false, "double-quoted string output")
	flag.BoolVar(&esc, "esc", false, "hex-escape all characters (C compat) when using -str")
	flag.Parse()

	var (
		w, sink io.WriteCloser
		r       io.Reader
	)

	switch input {
	case "":
		r = rand.Reader
	case "-":
		r = os.Stdin
	default:
		if f, err := os.Open(input); err != nil {
			fmt.Fprintln(os.Stderr, "Error:", err)
			os.Exit(1)
		} else {
			if n <= 0 || rep <= 0 {
				if info, err := f.Stat(); err == nil {
					size = info.Size()
					switch {
					case n <= 0 && rep <= 0:
						n = size
						rep = 1
					case n <= 0:
						n = size / rep
						if size%rep > 0 {
							n++
						}
					case rep <= 0:
						rep = size / n
						if size%rep > 0 {
							rep++
						}
					}
				}
			}
			r = f
		}
	}

	sink = NopWriteCloser(os.Stdout)

	for rep != 0 {
		switch {
		case raw:
			w = sink
		case code:
			w = NopWriteCloser(NewCodeWriter(sink, cols))
		case str:
			w = NewQuotedWriter(sink, esc)
		case b64:
			w = base64.NewEncoder(base64.StdEncoding, sink)
		default:
			w = NopWriteCloser(NewHexWriter(sink))
		}

		_, err := io.CopyN(w, r, n)
		w.Close() // ignoring error
		if !raw {
			fmt.Println()
		}
		if err != nil {
			fmt.Fprintln(os.Stderr, "Error:", err)
			os.Exit(1)
		}
		rep--
		if size > 0 {
			size -= n
			if size == 0 {
				break
			} else if size < n {
				n = size
				rep = 1
			}
		}
	}
}

var ErrPartialChunkWrite = errors.New("Partial chunk write")

func NopWriteCloser(w io.Writer) io.WriteCloser {
	return writecloser{w}
}

type writecloser struct {
	io.Writer
}

func (w writecloser) Close() error {
	return nil
}
<commit_msg>Added support for multiple output files. -r and -n handling is a bit more consistent.<commit_after>// Copyright 2012 Kevin Gillette. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"flag"
	"fmt" // flag uses fmt anyway
	"io"
	"os"
)

func main() {
	var (
		input, output string
		n, rep, size  int64
		raw, b64      bool
		code          bool
		cols          uint
		str, esc      bool
	)

	flag.StringVar(&input, "i", "", "input file or stdin if '-' (default is urandom)")
	flag.StringVar(&output, "o", "", "output filename printf-style pattern (default is stdout)")
	flag.Int64Var(&n, "n", 0, "number of (unencoded) bytes to output. Default 16")
	flag.Int64Var(&rep, "r", 0, "number of repititions (output lines). Default 1")
	flag.BoolVar(&raw, "raw", false, "raw output")
	flag.BoolVar(&b64, "64", false, "base64 output")
	flag.BoolVar(&code, "code", false, "comma-separated hex literals (typical array syntax)")
	flag.UintVar(&cols, "cols", 12, "number of columns to use with -code")
	flag.BoolVar(&str, "str", false, "double-quoted string output")
	flag.BoolVar(&esc, "esc", false, "hex-escape all characters (C compat) when using -str")
	flag.Parse()

	var (
		w, sink io.WriteCloser
		r       io.Reader
	)

	switch input {
	case "":
		r = rand.Reader
	case "-":
		r = os.Stdin
	default:
		if f, err := os.Open(input); err != nil {
			fmt.Fprintln(os.Stderr, "Error:", err)
			os.Exit(1)
		} else {
			if n <= 0 || rep <= 0 {
				if info, err := f.Stat(); err == nil {
					size = info.Size()
					switch {
					case n <= 0 && rep <= 0:
						n = size
						rep = 1
					case n <= 0:
						n = size / rep
						if size%rep > 0 {
							n++
						}
					case rep <= 0:
						rep = size / n
						if size%rep > 0 {
							rep++
						}
					}
				}
			}
			r = f
		}
	}

	if n == 0 {
		n = 16
	}
	if rep == 0 {
		rep = 1
	}

	var multioutput bool

	if output == "" {
		sink = NopWriteCloser(os.Stdout)
	} else {
		i := 0
		for _, r := range output {
			if r == '%' {
				i++
				break
			}
		}
		switch i {
		case 0:
			if f, err := os.Create(output); err != nil {
				fmt.Fprintln(os.Stderr, "Output file error:", err)
				os.Exit(1)
			} else {
				sink = NopWriteCloser(f)
			}
		case 1:
			multioutput = true
		case 2:
			// ignored for now
			fmt.Fprintf(os.Stderr, "Invalid filename pattern: %q\n", output)
			os.Exit(2)
		}
	}

	for i := int64(0); i != rep; i++ {
		if multioutput {
			if f, err := os.Create(fmt.Sprintf(output, i)); err != nil {
				fmt.Fprintln(os.Stderr, "Output file error:", err)
				os.Exit(1)
			} else {
				sink = f
			}
		}

		switch {
		case raw:
			w = sink
		case code:
			w = NopWriteCloser(NewCodeWriter(sink, cols))
		case str:
			w = NewQuotedWriter(sink, esc)
		case b64:
			w = base64.NewEncoder(base64.StdEncoding, sink)
		default:
			w = NopWriteCloser(NewHexWriter(sink))
		}

		_, err := io.CopyN(w, r, n)
		w.Close() // ignoring error
		if !raw && !multioutput {
			fmt.Fprintln(sink)
		}
		if err != nil {
			fmt.Fprintln(os.Stderr, "Error:", err)
			os.Exit(1)
		}
		if size > 0 {
			size -= n
			if size == 0 {
				break
			} else if size < n {
				n = size
				// one more iteration
				i = rep - 2
			}
		}
	}
}

var ErrPartialChunkWrite = errors.New("Partial chunk write")

func NopWriteCloser(w io.Writer) io.WriteCloser {
	return writecloser{w}
}

type writecloser struct {
	io.Writer
}

func (w writecloser) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"unicode/utf16"
)

type File struct {
	stringPool  *ResStringPool
	resourceMap []uint32
	Namespace   *ResXMLTreeNamespaceExt
	XMLBuffer   bytes.Buffer
}

const (
	RES_NULL_TYPE        = 0x0000
	RES_STRING_POOL_TYPE = 0x0001
	RES_TABLE_TYPE       = 0x0002
	RES_XML_TYPE         = 0x0003

	// Chunk types in RES_XML_TYPE
	RES_XML_FIRST_CHUNK_TYPE     = 0x0100
	RES_XML_START_NAMESPACE_TYPE = 0x0100
	RES_XML_END_NAMESPACE_TYPE   = 0x0101
	RES_XML_START_ELEMENT_TYPE   = 0x0102
	RES_XML_END_ELEMENT_TYPE     = 0x0103
	RES_XML_CDATA_TYPE           = 0x0104
	RES_XML_LAST_CHUNK_TYPE      = 0x017f

	// This contains a uint32_t array mapping strings in the string
	// pool back to resource identifiers.  It is optional.
	RES_XML_RESOURCE_MAP_TYPE = 0x0180

	// Chunk types in RES_TABLE_TYPE
	RES_TABLE_PACKAGE_TYPE   = 0x0200
	RES_TABLE_TYPE_TYPE      = 0x0201
	RES_TABLE_TYPE_SPEC_TYPE = 0x0202
)

type ResChunkHeader struct {
	Type       uint16
	HeaderSize uint16
	Size       uint32
}

const SORTED_FLAG = 1 << 0
const UTF8_FLAG = 1 << 8

type ResStringPoolRef uint32

const NilResStringPoolRef = ResStringPoolRef(0xFFFFFFFF)

type ResStringPoolHeader struct {
	Header      ResChunkHeader
	StringCount uint32
	StyleCount  uint32
	Flags       uint32
	StringStart uint32
	StylesStart uint32
}

type ResStringPool struct {
	Header  ResStringPoolHeader
	Strings []string
	Styles  []string
}

type ResXMLTreeNode struct {
	Header     ResChunkHeader
	LineNumber uint32
	Comment    ResStringPoolRef
}

type ResXMLTreeNamespaceExt struct {
	Prefix ResStringPoolRef
	Uri    ResStringPoolRef
}

type ResXMLTreeAttrExt struct {
	NS             ResStringPoolRef
	Name           ResStringPoolRef
	AttributeStart uint16
	AttributeSize  uint16
	AttributeCount uint16
	IdIndex        uint16
	ClassIndex     uint16
	StyleIndex     uint16
}

type ResXMLTreeAttribute struct {
	NS         ResStringPoolRef
	Name       ResStringPoolRef
	RawValue   ResStringPoolRef
	TypedValue ResValue
}

const (
	TYPE_NULL            = 0x00
	TYPE_REFERENCE       = 0x01
	TYPE_ATTRIBUTE       = 0x02
	TYPE_STRING          = 0x03
	TYPE_FLOAT           = 0x04
	TYPE_DIMENSION       = 0x05
	TYPE_FRACTION        = 0x06
	TYPE_FIRST_INT       = 0x10
	TYPE_INT_DEC         = 0x10
	TYPE_INT_HEX         = 0x11
	TYPE_INT_BOOLEAN     = 0x12
	TYPE_FIRST_COLOR_INT = 0x1c
	TYPE_INT_COLOR_ARGB8 = 0x1c
	TYPE_INT_COLOR_RGB8  = 0x1d
	TYPE_INT_COLOR_ARGB4 = 0x1e
	TYPE_INT_COLOR_RGB4  = 0x1f
	TYPE_LAST_COLOR_INT  = 0x1f
	TYPE_LAST_INT        = 0x1f
)

type ResValue struct {
	Size     uint16
	Res0     uint8
	DataType uint8
	Data     uint32
}

type ResXMLTreeEndElementExt struct {
	NS   ResStringPoolRef
	Name ResStringPoolRef
}

func NewFile(r io.ReaderAt) (*File, error) {
	f := new(File)
	sr := io.NewSectionReader(r, 0, 1<<63-1)

	header := new(ResChunkHeader)
	binary.Read(sr, binary.LittleEndian, header)
	offset := uint32(header.HeaderSize)

	for offset < header.Size {
		sr.Seek(int64(offset), os.SEEK_SET)
		chunkHeader := &ResChunkHeader{}
		binary.Read(sr, binary.LittleEndian, chunkHeader)

		var err error
		chunkReader := io.NewSectionReader(r, int64(offset), int64(chunkHeader.Size))
		switch chunkHeader.Type {
		case RES_STRING_POOL_TYPE:
			f.stringPool, err = ReadStringPool(chunkReader)
		case RES_XML_RESOURCE_MAP_TYPE:
			f.resourceMap, err = ReadResourceMap(chunkReader)
		case RES_XML_START_NAMESPACE_TYPE:
			err = f.ReadStartNamespace(chunkReader)
		case RES_XML_START_ELEMENT_TYPE:
			err = f.ReadStartElement(chunkReader)
		case RES_XML_END_ELEMENT_TYPE:
			err = f.ReadEndElement(chunkReader)
		}
		if err != nil {
			return nil, err
		}

		offset += chunkHeader.Size
	}
	return f, nil
}

func (f *File) GetString(ref ResStringPoolRef) string {
	if ref == NilResStringPoolRef {
		return ""
	}
	return f.stringPool.Strings[int(ref)]
}

func ReadStringPool(sr *io.SectionReader) (*ResStringPool, error) {
	sp := new(ResStringPool)
	binary.Read(sr, binary.LittleEndian, &sp.Header)

	stringStarts := make([]uint32, sp.Header.StringCount)
	binary.Read(sr, binary.LittleEndian, stringStarts)
	styleStarts := make([]uint32, sp.Header.StyleCount)
	binary.Read(sr, binary.LittleEndian, styleStarts)

	sp.Strings = make([]string, sp.Header.StringCount)
	for i, start := range stringStarts {
		var str string
		var err error
		if (sp.Header.Flags & UTF8_FLAG) == 0 {
			str, err = ReadUTF16(sr, int64(sp.Header.StringStart+start))
		} else {
			str, err = ReadUTF8(sr, int64(sp.Header.StringStart+start))
		}
		if err != nil {
			return nil, err
		}
		sp.Strings[i] = str
	}

	sp.Styles = make([]string, sp.Header.StyleCount)
	for i, start := range styleStarts {
		var str string
		var err error
		if (sp.Header.Flags & UTF8_FLAG) == 0 {
			str, err = ReadUTF16(sr, int64(sp.Header.StylesStart+start))
		} else {
			str, err = ReadUTF8(sr, int64(sp.Header.StylesStart+start))
		}
		if err != nil {
			return nil, err
		}
		sp.Styles[i] = str
	}

	return sp, nil
}

func ReadUTF16(sr *io.SectionReader, offset int64) (string, error) {
	var size uint16
	sr.Seek(offset, os.SEEK_SET)
	if err := binary.Read(sr, binary.LittleEndian, &size); err != nil {
		return "", err
	}
	buf := make([]uint16, size)
	if err := binary.Read(sr, binary.LittleEndian, buf); err != nil {
		return "", err
	}
	return string(utf16.Decode(buf)), nil
}

func ReadUTF8(sr *io.SectionReader, offset int64) (string, error) {
	var size uint16
	sr.Seek(offset, os.SEEK_SET)
	if err := binary.Read(sr, binary.LittleEndian, &size); err != nil {
		return "", err
	}
	buf := make([]uint8, size)
	if err := binary.Read(sr, binary.LittleEndian, buf); err != nil {
		return "", err
	}
	return string(buf), nil
}

func ReadResourceMap(sr *io.SectionReader) ([]uint32, error) {
	header := new(ResChunkHeader)
	binary.Read(sr, binary.LittleEndian, header)
	count := (header.Size - uint32(header.HeaderSize)) / 4
	resourceMap := make([]uint32, count)
	if err := binary.Read(sr, binary.LittleEndian, resourceMap); err != nil {
		return nil, err
	}
	return resourceMap, nil
}

func (f *File) ReadStartNamespace(sr *io.SectionReader) error {
	header := new(ResXMLTreeNode)
	binary.Read(sr, binary.LittleEndian, header)
	sr.Seek(int64(header.Header.HeaderSize), os.SEEK_SET)
	namespace := new(ResXMLTreeNamespaceExt)
	binary.Read(sr, binary.LittleEndian, namespace)
	f.Namespace = namespace
	return nil
}

func (f *File) ReadStartElement(sr *io.SectionReader) error {
	header := new(ResXMLTreeNode)
	binary.Read(sr, binary.LittleEndian, header)
	sr.Seek(int64(header.Header.HeaderSize), os.SEEK_SET)
	ext := new(ResXMLTreeAttrExt)
	binary.Read(sr, binary.LittleEndian, ext)

	fmt.Fprintf(&f.XMLBuffer, "<%s", f.AddNamespace(ext.NS, ext.Name))

	offset := int64(ext.AttributeStart + header.Header.HeaderSize)
	for i := 0; i < int(ext.AttributeCount); i++ {
		sr.Seek(offset, os.SEEK_SET)
		attr := new(ResXMLTreeAttribute)
		binary.Read(sr, binary.LittleEndian, attr)

		var value string
		if attr.RawValue != NilResStringPoolRef {
			value = f.GetString(attr.RawValue)
		} else {
			data := attr.TypedValue.Data
			switch attr.TypedValue.DataType {
			case TYPE_NULL:
				value = ""
			case TYPE_REFERENCE:
				value = fmt.Sprintf("@0x%08X", data)
			case TYPE_INT_DEC:
				value = fmt.Sprintf("%d", data)
			case TYPE_INT_HEX:
				value = fmt.Sprintf("0x%08X", data)
			case TYPE_INT_BOOLEAN:
				if data != 0 {
					value = "true"
				} else {
					value = "false"
				}
			default:
				value = fmt.Sprintf("@0x%08X", data)
			}
		}

		fmt.Fprintf(&f.XMLBuffer, " %s=\"", f.AddNamespace(attr.NS, attr.Name))
		xml.Escape(&f.XMLBuffer, []byte(value))
		fmt.Fprint(&f.XMLBuffer, "\"")
		offset += int64(ext.AttributeSize)
	}
	fmt.Fprint(&f.XMLBuffer, ">")
	return nil
}

func (f *File) ReadEndElement(sr *io.SectionReader) error {
	header := new(ResXMLTreeNode)
	if err := binary.Read(sr, binary.LittleEndian, header); err != nil {
		return err
	}
	sr.Seek(int64(header.Header.HeaderSize), os.SEEK_SET)
	ext := new(ResXMLTreeEndElementExt)
	if err := binary.Read(sr, binary.LittleEndian, ext); err != nil {
		return err
	}
	fmt.Fprintf(&f.XMLBuffer, "</%s>", f.AddNamespace(ext.NS, ext.Name))
	return nil
}

func (f *File) AddNamespace(ns, name ResStringPoolRef) string {
	if ns != NilResStringPoolRef {
		return fmt.Sprintf("%s:%s", f.GetString(f.Namespace.Prefix), f.GetString(name))
	} else {
		return f.GetString(name)
	}
}

func main() {
	f, _ := os.Open("AndroidManifest.xml")
	xml, _ := NewFile(f)
	fmt.Println(xml.XMLBuffer.String())
}
<commit_msg>XML宣言を追加<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"unicode/utf16"
)

type File struct {
	stringPool  *ResStringPool
	resourceMap []uint32
	Namespace   *ResXMLTreeNamespaceExt
	XMLBuffer   bytes.Buffer
}

const (
	RES_NULL_TYPE        = 0x0000
	RES_STRING_POOL_TYPE = 0x0001
	RES_TABLE_TYPE       = 0x0002
	RES_XML_TYPE         = 0x0003

	// Chunk types in RES_XML_TYPE
	RES_XML_FIRST_CHUNK_TYPE     = 0x0100
	RES_XML_START_NAMESPACE_TYPE = 0x0100
	RES_XML_END_NAMESPACE_TYPE   = 0x0101
	RES_XML_START_ELEMENT_TYPE   = 0x0102
	RES_XML_END_ELEMENT_TYPE     = 0x0103
	RES_XML_CDATA_TYPE           = 0x0104
	RES_XML_LAST_CHUNK_TYPE      = 0x017f

	// This contains a uint32_t array mapping strings in the string
	// pool back to resource identifiers.  It is optional.
	RES_XML_RESOURCE_MAP_TYPE = 0x0180

	// Chunk types in RES_TABLE_TYPE
	RES_TABLE_PACKAGE_TYPE   = 0x0200
	RES_TABLE_TYPE_TYPE      = 0x0201
	RES_TABLE_TYPE_SPEC_TYPE = 0x0202
)

type ResChunkHeader struct {
	Type       uint16
	HeaderSize uint16
	Size       uint32
}

const SORTED_FLAG = 1 << 0
const UTF8_FLAG = 1 << 8

type ResStringPoolRef uint32

const NilResStringPoolRef = ResStringPoolRef(0xFFFFFFFF)

type ResStringPoolHeader struct {
	Header      ResChunkHeader
	StringCount uint32
	StyleCount  uint32
	Flags       uint32
	StringStart uint32
	StylesStart uint32
}

type ResStringPool struct {
	Header  ResStringPoolHeader
	Strings []string
	Styles  []string
}

type ResXMLTreeNode struct {
	Header     ResChunkHeader
	LineNumber uint32
	Comment    ResStringPoolRef
}

type ResXMLTreeNamespaceExt struct {
	Prefix ResStringPoolRef
	Uri    ResStringPoolRef
}

type ResXMLTreeAttrExt struct {
	NS             ResStringPoolRef
	Name           ResStringPoolRef
	AttributeStart uint16
	AttributeSize  uint16
	AttributeCount uint16
	IdIndex        uint16
	ClassIndex     uint16
	StyleIndex     uint16
}

type ResXMLTreeAttribute struct {
	NS         ResStringPoolRef
	Name       ResStringPoolRef
	RawValue   ResStringPoolRef
	TypedValue ResValue
}

const (
	TYPE_NULL            = 0x00
	TYPE_REFERENCE       = 0x01
	TYPE_ATTRIBUTE       = 0x02
	TYPE_STRING          = 0x03
	TYPE_FLOAT           = 0x04
	TYPE_DIMENSION       = 0x05
	TYPE_FRACTION        = 0x06
	TYPE_FIRST_INT       = 0x10
	TYPE_INT_DEC         = 0x10
	TYPE_INT_HEX         = 0x11
	TYPE_INT_BOOLEAN     = 0x12
	TYPE_FIRST_COLOR_INT = 0x1c
	TYPE_INT_COLOR_ARGB8 = 0x1c
	TYPE_INT_COLOR_RGB8  = 0x1d
	TYPE_INT_COLOR_ARGB4 = 0x1e
	TYPE_INT_COLOR_RGB4  = 0x1f
	TYPE_LAST_COLOR_INT  = 0x1f
	TYPE_LAST_INT        = 0x1f
)

type ResValue struct {
	Size     uint16
	Res0     uint8
	DataType uint8
	Data     uint32
}

type ResXMLTreeEndElementExt struct {
	NS   ResStringPoolRef
	Name ResStringPoolRef
}

func NewFile(r io.ReaderAt) (*File, error) {
	f := new(File)
	sr := io.NewSectionReader(r, 0, 1<<63-1)

	header := new(ResChunkHeader)
	binary.Read(sr, binary.LittleEndian, header)
	offset := uint32(header.HeaderSize)

	fmt.Fprintf(&f.XMLBuffer, "<?xml version=\"1.0\" encoding=\"utf-8\"?>")

	for offset < header.Size {
		sr.Seek(int64(offset), os.SEEK_SET)
		chunkHeader := &ResChunkHeader{}
		binary.Read(sr, binary.LittleEndian, chunkHeader)

		var err error
		chunkReader := io.NewSectionReader(r, int64(offset), int64(chunkHeader.Size))
		switch chunkHeader.Type {
		case RES_STRING_POOL_TYPE:
			f.stringPool, err = ReadStringPool(chunkReader)
		case RES_XML_RESOURCE_MAP_TYPE:
			f.resourceMap, err = ReadResourceMap(chunkReader)
		case RES_XML_START_NAMESPACE_TYPE:
			err = f.ReadStartNamespace(chunkReader)
		case RES_XML_START_ELEMENT_TYPE:
			err = f.ReadStartElement(chunkReader)
		case RES_XML_END_ELEMENT_TYPE:
			err = f.ReadEndElement(chunkReader)
		}
		if err != nil {
			return nil, err
		}

		offset += chunkHeader.Size
	}
	return f, nil
}

func (f *File) GetString(ref ResStringPoolRef) string {
	if ref == NilResStringPoolRef {
		return ""
	}
	return f.stringPool.Strings[int(ref)]
}

func ReadStringPool(sr *io.SectionReader) (*ResStringPool, error) {
	sp := new(ResStringPool)
	binary.Read(sr, binary.LittleEndian, &sp.Header)

	stringStarts := make([]uint32, sp.Header.StringCount)
	binary.Read(sr, binary.LittleEndian, stringStarts)
	styleStarts := make([]uint32, sp.Header.StyleCount)
	binary.Read(sr, binary.LittleEndian, styleStarts)

	sp.Strings = make([]string, sp.Header.StringCount)
	for i, start := range stringStarts {
		var str string
		var err error
		if (sp.Header.Flags & UTF8_FLAG) == 0 {
			str, err = ReadUTF16(sr, int64(sp.Header.StringStart+start))
		} else {
			str, err = ReadUTF8(sr, int64(sp.Header.StringStart+start))
		}
		if err != nil {
			return nil, err
		}
		sp.Strings[i] = str
	}

	sp.Styles = make([]string, sp.Header.StyleCount)
	for i, start := range styleStarts {
		var str string
		var err error
		if (sp.Header.Flags & UTF8_FLAG) == 0 {
			str, err = ReadUTF16(sr, int64(sp.Header.StylesStart+start))
		} else {
			str, err = ReadUTF8(sr, int64(sp.Header.StylesStart+start))
		}
		if err != nil {
			return nil, err
		}
		sp.Styles[i] = str
	}

	return sp, nil
}

func ReadUTF16(sr *io.SectionReader, offset int64) (string, error) {
	var size uint16
	sr.Seek(offset, os.SEEK_SET)
	if err := binary.Read(sr, binary.LittleEndian, &size); err != nil {
		return "", err
	}
	buf := make([]uint16, size)
	if err := binary.Read(sr, binary.LittleEndian, buf); err != nil {
		return "", err
	}
	return string(utf16.Decode(buf)), nil
}

func ReadUTF8(sr *io.SectionReader, offset int64) (string, error) {
	var size uint16
	sr.Seek(offset, os.SEEK_SET)
	if err := binary.Read(sr, binary.LittleEndian, &size); err != nil {
		return "", err
	}
	buf := make([]uint8, size)
	if err := binary.Read(sr, binary.LittleEndian, buf); err != nil {
		return "", err
	}
	return string(buf), nil
}

func ReadResourceMap(sr *io.SectionReader) ([]uint32, error) {
	header := new(ResChunkHeader)
	binary.Read(sr, binary.LittleEndian, header)
	count := (header.Size - uint32(header.HeaderSize)) / 4
	resourceMap := make([]uint32, count)
	if err := binary.Read(sr, binary.LittleEndian, resourceMap); err != nil {
		return nil, err
	}
	return resourceMap, nil
}

func (f *File) ReadStartNamespace(sr *io.SectionReader) error {
	header := new(ResXMLTreeNode)
	binary.Read(sr, binary.LittleEndian, header)
	sr.Seek(int64(header.Header.HeaderSize), os.SEEK_SET)
	namespace := new(ResXMLTreeNamespaceExt)
	binary.Read(sr, binary.LittleEndian, namespace)
	f.Namespace = namespace
	return nil
}

func (f *File) ReadStartElement(sr *io.SectionReader) error {
	header := new(ResXMLTreeNode)
	binary.Read(sr, binary.LittleEndian, header)
	sr.Seek(int64(header.Header.HeaderSize), os.SEEK_SET)
	ext := new(ResXMLTreeAttrExt)
	binary.Read(sr, binary.LittleEndian, ext)

	fmt.Fprintf(&f.XMLBuffer, "<%s", f.AddNamespace(ext.NS, ext.Name))

	offset := int64(ext.AttributeStart + header.Header.HeaderSize)
	for i := 0; i < int(ext.AttributeCount); i++ {
		sr.Seek(offset, os.SEEK_SET)
		attr := new(ResXMLTreeAttribute)
		binary.Read(sr, binary.LittleEndian, attr)

		var value string
		if attr.RawValue != NilResStringPoolRef {
			value = f.GetString(attr.RawValue)
		} else {
			data := attr.TypedValue.Data
			switch attr.TypedValue.DataType {
			case TYPE_NULL:
				value = ""
			case TYPE_REFERENCE:
				value = fmt.Sprintf("@0x%08X", data)
			case TYPE_INT_DEC:
				value = fmt.Sprintf("%d", data)
			case TYPE_INT_HEX:
				value = fmt.Sprintf("0x%08X", data)
			case TYPE_INT_BOOLEAN:
				if data != 0 {
					value = "true"
				} else {
					value = "false"
				}
			default:
				value = fmt.Sprintf("@0x%08X", data)
			}
		}

		fmt.Fprintf(&f.XMLBuffer, " %s=\"", f.AddNamespace(attr.NS, attr.Name))
		xml.Escape(&f.XMLBuffer, []byte(value))
		fmt.Fprint(&f.XMLBuffer, "\"")
		offset += int64(ext.AttributeSize)
	}
	fmt.Fprint(&f.XMLBuffer, ">")
	return nil
}

func (f *File) ReadEndElement(sr *io.SectionReader) error {
	header := new(ResXMLTreeNode)
	if err := binary.Read(sr, binary.LittleEndian, header); err != nil {
		return err
	}
	sr.Seek(int64(header.Header.HeaderSize), os.SEEK_SET)
	ext := new(ResXMLTreeEndElementExt)
	if err := binary.Read(sr, binary.LittleEndian, ext); err != nil {
		return err
	}
	fmt.Fprintf(&f.XMLBuffer, "</%s>", f.AddNamespace(ext.NS, ext.Name))
	return nil
}

func (f *File) AddNamespace(ns, name ResStringPoolRef) string {
	if ns != NilResStringPoolRef {
		return fmt.Sprintf("%s:%s", f.GetString(f.Namespace.Prefix), f.GetString(name))
	} else {
		return f.GetString(name)
	}
}

func main() {
	f, _ := os.Open("AndroidManifest.xml")
	xml, _ := NewFile(f)
	fmt.Println(xml.XMLBuffer.String())
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	cmp "github.com/kshvmdn/fsql/compare"
	"github.com/kshvmdn/fsql/query"
)

const (
	uBYTE     = 1.0
	uKILOBYTE = 1024 * uBYTE
	uMEGABYTE = 1024 * uKILOBYTE
	uGIGABYTE = 1024 * uMEGABYTE
)

func compare(condition query.Condition, file os.FileInfo) bool {
	var retval bool

	switch condition.Attribute {
	case "name":
		retval = cmp.Alpha(condition.Comparator, file.Name(), condition.Value)

	case "size":
		mult := uBYTE

		if len(condition.Value) > 2 {
			unit := strings.ToLower(condition.Value[len(condition.Value)-2:])
			switch unit {
			case "kb":
				mult = uKILOBYTE
			case "mb":
				mult = uMEGABYTE
			case "gb":
				mult = uGIGABYTE
			}

			if mult > 1 {
				condition.Value = condition.Value[:len(condition.Value)-2]
			}
		}

		size, err := strconv.ParseFloat(condition.Value, 64)
		if err != nil {
			return false
		}
		retval = cmp.Numeric(condition.Comparator, file.Size(), int64(size*mult))

	case "time":
		t, err := time.Parse("Jan 02 2006 15 04", condition.Value)
		if err != nil {
			return false
		}
		retval = cmp.Time(condition.Comparator, file.ModTime(), t)

	case "file":
		retval = cmp.File(condition.Comparator, file, condition.Value)
	}

	if condition.Negate {
		return !retval
	}

	return retval
}

func containsAny(exclusions []string, path string) bool {
	for _, exclusion := range exclusions {
		if strings.Contains(path, exclusion) {
			return true
		}
	}

	return false
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Expected query.")
		os.Exit(1)
	}
	input := os.Args[1]

	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}

	q, err := query.Run(input)
	if err != nil {
		if err == io.ErrUnexpectedEOF {
			log.Fatal("Unexpected end of line.")
		}
		log.Fatal(err)
	}

	err = q.ReduceInclusions()
	if err != nil {
		log.Fatal(err)
	}

	for _, src := range q.Sources["include"] {
		if strings.Contains(src, "~") {
			src = filepath.Join(usr.HomeDir, src[1:])
		}

		filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
			if path == "." || path == ".." || containsAny(q.Sources["exclude"], path) {
				return nil
			}

			if q.ConditionTree.Evaluate(info, compare) {
				if q.HasAttribute("mode") {
					fmt.Printf("%s\t", info.Mode())
				}

				if q.HasAttribute("size") {
					fmt.Printf("%d\t", info.Size())
				}

				if q.HasAttribute("time") {
					fmt.Printf("%s\t", info.ModTime().Format(time.Stamp))
				}

				if q.HasAttribute("name") {
					if strings.Contains(path, usr.HomeDir) {
						path = filepath.Join("~", path[len(usr.HomeDir):])
					}
					fmt.Printf("%s", path)
				}

				fmt.Printf("\n")
			}

			return nil
		})
	}
}
<commit_msg>Expand tildes in source dirs prior to filepath.Rel call<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	cmp "github.com/kshvmdn/fsql/compare"
	"github.com/kshvmdn/fsql/query"
)

const (
	uBYTE     = 1.0
	uKILOBYTE = 1024 * uBYTE
	uMEGABYTE = 1024 * uKILOBYTE
	uGIGABYTE = 1024 * uMEGABYTE
)

func compare(condition query.Condition, file os.FileInfo) bool {
	var retval bool

	switch condition.Attribute {
	case "name":
		retval = cmp.Alpha(condition.Comparator, file.Name(), condition.Value)

	case "size":
		mult := uBYTE

		if len(condition.Value) > 2 {
			unit := strings.ToLower(condition.Value[len(condition.Value)-2:])
			switch unit {
			case "kb":
				mult = uKILOBYTE
			case "mb":
				mult = uMEGABYTE
			case "gb":
				mult = uGIGABYTE
			}

			if mult > 1 {
				condition.Value = condition.Value[:len(condition.Value)-2]
			}
		}

		size, err := strconv.ParseFloat(condition.Value, 64)
		if err != nil {
			return false
		}
		retval = cmp.Numeric(condition.Comparator, file.Size(), int64(size*mult))

	case "time":
		t, err := time.Parse("Jan 02 2006 15 04", condition.Value)
		if err != nil {
			return false
		}
		retval = cmp.Time(condition.Comparator, file.ModTime(), t)

	case "file":
		retval = cmp.File(condition.Comparator, file, condition.Value)
	}

	if condition.Negate {
		return !retval
	}

	return retval
}

func containsAny(exclusions []string, path string) bool {
	for _, exclusion := range exclusions {
		if strings.Contains(path, exclusion) {
			return true
		}
	}

	return false
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Expected query.")
		os.Exit(1)
	}
	input := os.Args[1]

	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}

	q, err := query.Run(input)
	if err != nil {
		if err == io.ErrUnexpectedEOF {
			log.Fatal("Unexpected end of line.")
		}
		log.Fatal(err)
	}

	// Replace ~ with the home directory in each source directory.
	for _, sourceType := range []string{"include", "exclude"} {
		for i, src := range q.Sources[sourceType] {
			if strings.Contains(src, "~") {
				q.Sources[sourceType][i] = filepath.Join(usr.HomeDir, src[1:])
			}
		}
	}

	err = q.ReduceInclusions()
	if err != nil {
		log.Fatal(err)
	}

	for _, src := range q.Sources["include"] {
		filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
			if path == "." || path == ".." || containsAny(q.Sources["exclude"], path) {
				return nil
			}

			if q.ConditionTree.Evaluate(info, compare) {
				if q.HasAttribute("mode") {
					fmt.Printf("%s\t", info.Mode())
				}

				if q.HasAttribute("size") {
					fmt.Printf("%d\t", info.Size())
				}

				if q.HasAttribute("time") {
					fmt.Printf("%s\t", info.ModTime().Format(time.Stamp))
				}

				if q.HasAttribute("name") {
					if strings.Contains(path, usr.HomeDir) {
						path = filepath.Join("~", path[len(usr.HomeDir):])
					}
					fmt.Printf("%s", path)
				}

				fmt.Printf("\n")
			}

			return nil
		})
	}
}
<|endoftext|>


<commit_before>/*-
 * Copyright 2015 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"crypto/tls"
	"fmt"
	"log"
	"log/syslog"
	"os"
	"regexp"
	"runtime"
	"sync"
	"time"

	"github.com/kavu/go_reuseport"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	listenAddress  = kingpin.Flag("listen", "Address and port to listen on.").Required().TCP()
	forwardAddress = kingpin.Flag("target", "Address to foward connections to.").Required().TCP()
	unsafeTarget   = kingpin.Flag("unsafe-target", "If set, does not limit target to localhost, 127.0.0.1 or ::1").Bool()
	keystorePath   = kingpin.Flag("keystore", "Path to certificate and keystore (PKCS12).").PlaceHolder("PATH").Required().String()
	keystorePass   = kingpin.Flag("storepass", "Password for certificate and keystore.").PlaceHolder("PASS").Required().String()
	caBundlePath   = kingpin.Flag("cacert", "Path to certificate authority bundle file (PEM/X509).").Required().String()
	autoReload     = kingpin.Flag("auto-reload", "Watch keystores file with inotify/fswatch and reload on changes.").Bool()
	timedReload    = kingpin.Flag("timed-reload", "Reload keystores every N minutes, refresh listener on changes.").PlaceHolder("N").Int()
	allowAll       = kingpin.Flag("allow-all", "Allow all clients, do not check client cert subject.").Bool()
	allowedCNs     = kingpin.Flag("allow-cn", "Allow clients with given common name (can be repeated).").PlaceHolder("CN").Strings()
	allowedOUs     = kingpin.Flag("allow-ou", "Allow clients with organizational unit name (can be repeated).").PlaceHolder("OU").Strings()
	useSyslog      = kingpin.Flag("syslog", "Send logs to syslog instead of stderr.").Bool()
)

// Global logger instance
var logger = log.New(os.Stderr, "", log.LstdFlags|log.Lmicroseconds)

func initLogger() {
	if *useSyslog {
		var err error
		logger, err = syslog.NewLogger(syslog.LOG_NOTICE|syslog.LOG_DAEMON, log.LstdFlags|log.Lmicroseconds)
		panicOnError(err)
	}

	// Set log prefix to process ID to distinguish parent/child
	logger.SetPrefix(fmt.Sprintf("[%5d] ", os.Getpid()))
}

// panicOnError panics if err is not nil
func panicOnError(err error) {
	if err != nil {
		panic(err)
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	kingpin.Parse()

	// Validate flags
	if !(*allowAll) && len(*allowedCNs) == 0 && len(*allowedOUs) == 0 {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: at least one of --allow-all, --allow-cn or --allow-ou is required")
		os.Exit(1)
	}
	if *allowAll && len(*allowedCNs) != 0 {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --allow-all and --allow-cn are mutually exclusive")
		os.Exit(1)
	}
	if *allowAll && len(*allowedOUs) != 0 {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --allow-all and --allow-ou are mutually exclusive")
		os.Exit(1)
	}
	if !validateTarget((*forwardAddress).String()) {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --target must be localhost:port, 127.0.0.1:port or [::1]:port")
		os.Exit(1)
	}
	if *autoReload && (*timedReload > 0) {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --auto-reload and --timed-reload are mutually exclusive")
		os.Exit(1)
	}

	initLogger()

	listeners := &sync.WaitGroup{}
	listeners.Add(1)

	// Set up file watchers (if requested)
	watcher := make(chan bool, 1)
	if *autoReload {
		go watchAuto([]string{*keystorePath, *caBundlePath}, watcher)
	} else if *timedReload > 0 {
		go watchTimed([]string{*keystorePath, *caBundlePath}, time.Duration(*timedReload)*time.Second, watcher)
	}

	// A channel to notify us that the listener is running.
	started := make(chan bool, 1)
	go listen(started, listeners, watcher)

	up := <-started
	if !up {
		logger.Printf("failed to start initial listener")
		os.Exit(1)
	}

	logger.Printf("initial startup completed, waiting for connections")
	listeners.Wait()
	logger.Printf("all listeners closed, shutting down")
}

func validateTarget(addr string) bool {
	if *unsafeTarget {
		return true
	}
	if matched, _ := regexp.MatchString("^127[.]0[.]0[.]1:\\d+$", addr); matched {
		return true
	}
	if matched, _ := regexp.MatchString("^\\[::1\\]:\\d+$", addr); matched {
		return true
	}
	if matched, _ := regexp.MatchString("^localhost:\\d+$", addr); matched {
		return true
	}
	return false
}

// Open listening socket. Take note that we create a "reusable port
// listener", meaning we pass SO_REUSEPORT to the kernel. This allows
// us to have multiple sockets listening on the same port and accept
// connections. This is useful for the purposes of replacing certificates
// in-place without having to take downtime, e.g. if a certificate is
// expiring.
func listen(started chan bool, listeners *sync.WaitGroup, watcher chan bool) {
	// Open raw listening socket
	network, address := decodeAddress(*listenAddress)
	rawListener, err := reuseport.NewReusablePortListener(network, address)
	if err != nil {
		logger.Printf("error opening socket: %s", err)
		started <- false
		return
	}

	// Wrap listening socket with TLS listener.
	tlsConfig, err := buildConfig(*keystorePath, *keystorePass, *caBundlePath)
	if err != nil {
		logger.Printf("error setting up TLS: %s", err)
		started <- false
		return
	}

	leaf := tlsConfig.Certificates[0].Leaf

	listener := tls.NewListener(rawListener, tlsConfig)
	logger.Printf("listening on %s", *listenAddress)
	defer listener.Close()

	handlers := &sync.WaitGroup{}
	handlers.Add(1)

	// A channel to allow signal handlers to notify our accept loop that it
	// should shut down.
	stopper := make(chan bool, 1)

	go accept(listener, handlers, stopper, leaf)
	go signalHandler(listener, stopper, listeners, watcher)

	started <- true

	logger.Printf("listening with cert serial no. %d (expiring %s)", leaf.SerialNumber, leaf.NotAfter.String())
	handlers.Wait()

	listeners.Done()
}
<commit_msg>Fix flag documentation<commit_after>/*-
 * Copyright 2015 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"crypto/tls"
	"fmt"
	"log"
	"log/syslog"
	"os"
	"regexp"
	"runtime"
	"sync"
	"time"

	"github.com/kavu/go_reuseport"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	listenAddress  = kingpin.Flag("listen", "Address and port to listen on.").Required().TCP()
	forwardAddress = kingpin.Flag("target", "Address to foward connections to.").Required().TCP()
	unsafeTarget   = kingpin.Flag("unsafe-target", "If set, does not limit target to localhost, 127.0.0.1 or ::1").Bool()
	keystorePath   = kingpin.Flag("keystore", "Path to certificate and keystore (PKCS12).").PlaceHolder("PATH").Required().String()
	keystorePass   = kingpin.Flag("storepass", "Password for certificate and keystore.").PlaceHolder("PASS").Required().String()
	caBundlePath   = kingpin.Flag("cacert", "Path to certificate authority bundle file (PEM/X509).").Required().String()
	autoReload     = kingpin.Flag("auto-reload", "Watch keystores file with inotify/fswatch and reload on changes.").Bool()
	timedReload    = kingpin.Flag("timed-reload", "Reload keystores every N seconds, refresh listener on changes.").PlaceHolder("N").Int()
	allowAll       = kingpin.Flag("allow-all", "Allow all clients, do not check client cert subject.").Bool()
	allowedCNs     = kingpin.Flag("allow-cn", "Allow clients with given common name (can be repeated).").PlaceHolder("CN").Strings()
	allowedOUs     = kingpin.Flag("allow-ou", "Allow clients with organizational unit name (can be repeated).").PlaceHolder("OU").Strings()
	useSyslog      = kingpin.Flag("syslog", "Send logs to syslog instead of stderr.").Bool()
)

// Global logger instance
var logger = log.New(os.Stderr, "", log.LstdFlags|log.Lmicroseconds)

func initLogger() {
	if *useSyslog {
		var err error
		logger, err = syslog.NewLogger(syslog.LOG_NOTICE|syslog.LOG_DAEMON, log.LstdFlags|log.Lmicroseconds)
		panicOnError(err)
	}

	// Set log prefix to process ID to distinguish parent/child
	logger.SetPrefix(fmt.Sprintf("[%5d] ", os.Getpid()))
}

// panicOnError panics if err is not nil
func panicOnError(err error) {
	if err != nil {
		panic(err)
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	kingpin.Parse()

	// Validate flags
	if !(*allowAll) && len(*allowedCNs) == 0 && len(*allowedOUs) == 0 {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: at least one of --allow-all, --allow-cn or --allow-ou is required")
		os.Exit(1)
	}
	if *allowAll && len(*allowedCNs) != 0 {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --allow-all and --allow-cn are mutually exclusive")
		os.Exit(1)
	}
	if *allowAll && len(*allowedOUs) != 0 {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --allow-all and --allow-ou are mutually exclusive")
		os.Exit(1)
	}
	if !validateTarget((*forwardAddress).String()) {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --target must be localhost:port, 127.0.0.1:port or [::1]:port")
		os.Exit(1)
	}
	if *autoReload && (*timedReload > 0) {
		fmt.Fprintf(os.Stderr, "ghostunnel: error: --auto-reload and --timed-reload are mutually exclusive")
		os.Exit(1)
	}

	initLogger()

	listeners := &sync.WaitGroup{}
	listeners.Add(1)

	// Set up file watchers (if requested)
	watcher := make(chan bool, 1)
	if *autoReload {
		go watchAuto([]string{*keystorePath, *caBundlePath}, watcher)
	} else if *timedReload > 0 {
		go watchTimed([]string{*keystorePath, *caBundlePath}, time.Duration(*timedReload)*time.Second, watcher)
	}

	// A channel to notify us that the listener is running.
	started := make(chan bool, 1)
	go listen(started, listeners, watcher)

	up := <-started
	if !up {
		logger.Printf("failed to start initial listener")
		os.Exit(1)
	}

	logger.Printf("initial startup completed, waiting for connections")
	listeners.Wait()
	logger.Printf("all listeners closed, shutting down")
}

func validateTarget(addr string) bool {
	if *unsafeTarget {
		return true
	}
	if matched, _ := regexp.MatchString("^127[.]0[.]0[.]1:\\d+$", addr); matched {
		return true
	}
	if matched, _ := regexp.MatchString("^\\[::1\\]:\\d+$", addr); matched {
		return true
	}
	if matched, _ := regexp.MatchString("^localhost:\\d+$", addr); matched {
		return true
	}
	return false
}

// Open listening socket. Take note that we create a "reusable port
// listener", meaning we pass SO_REUSEPORT to the kernel. This allows
// us to have multiple sockets listening on the same port and accept
// connections. This is useful for the purposes of replacing certificates
// in-place without having to take downtime, e.g. if a certificate is
// expiring.
func listen(started chan bool, listeners *sync.WaitGroup, watcher chan bool) {
	// Open raw listening socket
	network, address := decodeAddress(*listenAddress)
	rawListener, err := reuseport.NewReusablePortListener(network, address)
	if err != nil {
		logger.Printf("error opening socket: %s", err)
		started <- false
		return
	}

	// Wrap listening socket with TLS listener.
	tlsConfig, err := buildConfig(*keystorePath, *keystorePass, *caBundlePath)
	if err != nil {
		logger.Printf("error setting up TLS: %s", err)
		started <- false
		return
	}

	leaf := tlsConfig.Certificates[0].Leaf

	listener := tls.NewListener(rawListener, tlsConfig)
	logger.Printf("listening on %s", *listenAddress)
	defer listener.Close()

	handlers := &sync.WaitGroup{}
	handlers.Add(1)

	// A channel to allow signal handlers to notify our accept loop that it
	// should shut down.
	stopper := make(chan bool, 1)

	go accept(listener, handlers, stopper, leaf)
	go signalHandler(listener, stopper, listeners, watcher)

	started <- true

	logger.Printf("listening with cert serial no. %d (expiring %s)", leaf.SerialNumber, leaf.NotAfter.String())
	handlers.Wait()

	listeners.Done()
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2014 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package donna

import(`github.com/michaeldv/donna/expect`; `testing`)

// Doubled pawns.
func TestEvaluatePawns100(t *testing.T) {
	p := NewGame(`Ke1,h2,h3`, `Ke8,a7,a6`).Start(White)
	score := p.Evaluate()
	expect.Eq(t, score, rightToMove.endgame) // Right to move only.
}

func TestEvaluatePawns110(t *testing.T) {
	game := NewGame(`Ke1,h2,h3`, `Ke8,a7,h7`)
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -19)
}

func TestEvaluatePawns120(t *testing.T) {
	game := NewGame(`Ke1,f4,f5`, `Ke8,f7,h7`)
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -23)
}

// Passed pawns.
func TestEvaluatePawns200(t *testing.T) {
	game := NewGame(`Ke1,h4`, `Ke8,h5`) // Blocked.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns210(t *testing.T) {
	game := NewGame(`Ke1,h4`, `Ke8,g7`) // Can't pass.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 9)
}

func TestEvaluatePawns220(t *testing.T) {
	game := NewGame(`Ke1,e4`, `Ke8,d6`) // Can't pass.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns230(t *testing.T) {
	game := NewGame(`Ke1,e5`, `Ke8,e4`) // Both passing.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns240(t *testing.T) {
	game := NewGame(`Kd1,e5`, `Ke8,d5`) // Both passing but white is closer.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 34)
}

func TestEvaluatePawns250(t *testing.T) {
	game := NewGame(`Ke1,a5`, `Kd8,h7`) // Both passing but white is much closer.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 97)
}

// Isolated pawns.
func TestEvaluatePawns300(t *testing.T) {
	game := NewGame(`Ke1,a5,c5`, `Kd8,f4,h4`) // All pawns are isolated.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns310(t *testing.T) {
	game := NewGame(`Ke1,a2,c2,e2`, `Ke8,a7,b7,c7`) // White pawns are isolated.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -52)
}

// Rooks.
func TestEvaluatePawns400(t *testing.T) {
	game := NewGame(`Ke1,Ra7`, `Ke8,Rh3`) // White on 7th.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 15)
}

func TestEvaluatePawns410(t *testing.T) {
	game := NewGame(`Ke1,Rb1,Ng2,a2`, `Ke8,Rh8,Nb7,h7`) // White on open file.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 112)
}

func TestEvaluatePawns420(t *testing.T) {
	game := NewGame(`Ke1,Rb1,a2,g2`, `Ke8,Rh8,h7,b7`) // White on semi-open file.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 119)
}

// King shield.
func TestEvaluatePawns500(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,f7,g7,h7,Qa6,Na5`) // h2,g2,h2 == f7,g7,h7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 8)
}
func TestEvaluatePawns505(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,f7,g6,h7,Qa6,Na5`) // h2,g2,h2 vs f7,G6,h7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 14)
}

func TestEvaluatePawns510(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,f5,g6,h7,Qa6,Na5`) // h2,g2,h2 vs F5,G6,h7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 22)
}

func TestEvaluatePawns520(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,a7,f7,g7,Qa6,Na5`) // h2,g2,h2 vs A7,f7,g7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 45)
}

func TestEvaluatePawns530(t *testing.T) {
	game := NewGame(`Kb1,a3,b2,c2,Qh3,Nh4`, `Kb8,a7,b7,c7,Qh6,Nh5`) // A3,b2,c2 vs a7,b7,c7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 2)
}

func TestEvaluatePawns540(t *testing.T) {
	game := NewGame(`Kb1,a3,b4,c2,Qh3,Nh4`, `Kb8,a7,b7,c7,Qh6,Nh5`) // A3,B4,c2 vs a7,b7,c7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -5)
}

func TestEvaluatePawns550(t *testing.T) {
	game := NewGame(`Kb1,b2,c2,h2,Qh3,Nh4`, `Kb8,a7,b7,c7,Qh6,Nh5`) // b2,c2,H2 vs a7,b7,c7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -29)
}

func TestEvaluatePawns560(t *testing.T) {
	game := NewGame(`Ka1,a3,b2,Qc1,Nd2`, `Kh8,g7,h6,Qf8,Ne7`) // a3,b2 == g7,h6
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 8)
}

func TestEvaluatePawns570(t *testing.T) {
	game := NewGame(`Kb1,a2,c2,f2,g2,h2`, `Kg8,a7,c7,f7,g7,h7`) // B2 hole but not enough power to bother.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}
<commit_msg>Updated pawn evaluation tests<commit_after>// Copyright (c) 2013-2014 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package donna

import(`github.com/michaeldv/donna/expect`; `testing`)

// Doubled pawns.
func TestEvaluatePawns100(t *testing.T) {
	p := NewGame(`Ke1,h2,h3`, `Ke8,a7,a6`).Start(White)
	score := p.Evaluate()
	expect.Eq(t, score, rightToMove.endgame) // Right to move only.
}

func TestEvaluatePawns110(t *testing.T) {
	game := NewGame(`Ke1,h2,h3`, `Ke8,a7,h7`)
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -19)
}

func TestEvaluatePawns120(t *testing.T) {
	game := NewGame(`Ke1,f4,f5`, `Ke8,f7,h7`)
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -23)
}

// Passed pawns.
func TestEvaluatePawns200(t *testing.T) {
	game := NewGame(`Ke1,h4`, `Ke8,h5`) // Blocked.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns210(t *testing.T) {
	game := NewGame(`Ke1,h4`, `Ke8,g7`) // Can't pass.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 9)
}

func TestEvaluatePawns220(t *testing.T) {
	game := NewGame(`Ke1,e4`, `Ke8,d6`) // Can't pass.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns230(t *testing.T) {
	game := NewGame(`Ke1,e5`, `Ke8,e4`) // Both passing.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns240(t *testing.T) {
	game := NewGame(`Kd1,e5`, `Ke8,d5`) // Both passing but white is closer.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 34)
}

func TestEvaluatePawns250(t *testing.T) {
	game := NewGame(`Ke1,a5`, `Kd8,h7`) // Both passing but white is much closer.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 97)
}

// Isolated pawns.
func TestEvaluatePawns300(t *testing.T) {
	game := NewGame(`Ke1,a5,c5`, `Kd8,f4,h4`) // All pawns are isolated.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}

func TestEvaluatePawns310(t *testing.T) {
	game := NewGame(`Ke1,a2,c2,e2`, `Ke8,a7,b7,c7`) // White pawns are isolated.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -52)
}

// Rooks.
func TestEvaluatePawns400(t *testing.T) {
	game := NewGame(`Ke1,Ra7`, `Ke8,Rh3`) // White on 7th.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 15)
}

func TestEvaluatePawns410(t *testing.T) {
	game := NewGame(`Ke1,Rb1,Ng2,a2`, `Ke8,Rh8,Nb7,h7`) // White on open file.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 112)
}

func TestEvaluatePawns420(t *testing.T) {
	game := NewGame(`Ke1,Rb1,a2,g2`, `Ke8,Rh8,h7,b7`) // White on semi-open file.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 119)
}

// King shield.
func TestEvaluatePawns500(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,f7,g7,h7,Qa6,Na5`) // h2,g2,h2 == f7,g7,h7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 8)
}
func TestEvaluatePawns505(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,f7,g6,h7,Qa6,Na5`) // h2,g2,h2 vs f7,G6,h7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 13)
}

func TestEvaluatePawns510(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,f5,g6,h7,Qa6,Na5`) // h2,g2,h2 vs F5,G6,h7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 26)
}

func TestEvaluatePawns520(t *testing.T) {
	game := NewGame(`Kg1,f2,g2,h2,Qa3,Na4`, `Kg8,a7,f7,g7,Qa6,Na5`) // h2,g2,h2 vs A7,f7,g7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 47)
}

func TestEvaluatePawns530(t *testing.T) {
	game := NewGame(`Kb1,a3,b2,c2,Qh3,Nh4`, `Kb8,a7,b7,c7,Qh6,Nh5`) // A3,b2,c2 vs a7,b7,c7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 2)
}

func TestEvaluatePawns540(t *testing.T) {
	game := NewGame(`Kb1,a3,b4,c2,Qh3,Nh4`, `Kb8,a7,b7,c7,Qh6,Nh5`) // A3,B4,c2 vs a7,b7,c7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -9)
}

func TestEvaluatePawns550(t *testing.T) {
	game := NewGame(`Kb1,b2,c2,h2,Qh3,Nh4`, `Kb8,a7,b7,c7,Qh6,Nh5`) // b2,c2,H2 vs a7,b7,c7
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, -31)
}

func TestEvaluatePawns560(t *testing.T) {
	game := NewGame(`Ka1,a3,b2,Qc1,Nd2`, `Kh8,g7,h6,Qf8,Ne7`) // a3,b2 == g7,h6
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 8)
}

func TestEvaluatePawns570(t *testing.T) {
	game := NewGame(`Kb1,a2,c2,f2,g2,h2`, `Kg8,a7,c7,f7,g7,h7`) // B2 hole but not enough power to bother.
	score := game.Start(White).Evaluate()

	expect.Eq(t, score, 5)
}
<|endoftext|>


<commit_before>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	_ "net/http/pprof" // Comment this line to disable pprof endpoint.
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/prometheus/log"

	clientmodel "github.com/prometheus/client_golang/model"
	registry "github.com/prometheus/client_golang/prometheus"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/notification"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/rules"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/local"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/prometheus/prometheus/storage/remote/influxdb"
	"github.com/prometheus/prometheus/storage/remote/opentsdb"
	"github.com/prometheus/prometheus/web"
	"github.com/prometheus/prometheus/web/api/legacy"
	"github.com/prometheus/prometheus/web/api/v1"
)

const deletionBatchSize = 100

// Commandline flags.
var (
	configFile = flag.String("config.file", "prometheus.yml", "Prometheus configuration file name.")

	alertmanagerURL           = flag.String("alertmanager.url", "", "The URL of the alert manager to send notifications to.")
	notificationQueueCapacity = flag.Int("alertmanager.notification-queue-capacity", 100, "The capacity of the queue for pending alert manager notifications.")

	persistenceStoragePath = flag.String("storage.local.path", "/tmp/metrics", "Base path for metrics storage.")

	opentsdbURL          = flag.String("storage.remote.opentsdb-url", "", "The URL of the remote OpenTSDB server to send samples to. None, if empty.")
	influxdbURL          = flag.String("storage.remote.influxdb-url", "", "The URL of the remote InfluxDB server to send samples to. None, if empty.")
	remoteStorageTimeout = flag.Duration("storage.remote.timeout", 30*time.Second, "The timeout to use when sending samples to the remote storage.")

	numMemoryChunks = flag.Int("storage.local.memory-chunks", 1024*1024, "How many chunks to keep in memory. While the size of a chunk is 1kiB, the total memory usage will be significantly higher than this value * 1kiB. Furthermore, for various reasons, more chunks might have to be kept in memory temporarily.")

	persistenceRetentionPeriod = flag.Duration("storage.local.retention", 15*24*time.Hour, "How long to retain samples in the local storage.")
	maxChunksToPersist         = flag.Int("storage.local.max-chunks-to-persist", 1024*1024, "How many chunks can be waiting for persistence before sample ingestion will stop. Many chunks waiting to be persisted will increase the checkpoint size.")

	checkpointInterval         = flag.Duration("storage.local.checkpoint-interval", 5*time.Minute, "The period at which the in-memory metrics and the chunks not yet persisted to series files are checkpointed.")
	checkpointDirtySeriesLimit = flag.Int("storage.local.checkpoint-dirty-series-limit", 5000, "If approx. that many time series are in a state that would require a recovery operation after a crash, a checkpoint is triggered, even if the checkpoint interval hasn't passed yet. A recovery operation requires a disk seek. The default limit intends to keep the recovery time below 1min even on spinning disks. With SSD, recovery is much faster, so you might want to increase this value in that case to avoid overly frequent checkpoints.")
	seriesSyncStrategy         = flag.String("storage.local.series-sync-strategy", "adaptive", "When to sync series files after modification. Possible values: 'never', 'always', 'adaptive'. Sync'ing slows down storage performance but reduces the risk of data loss in case of an OS crash. With the 'adaptive' strategy, series files are sync'd for as long as the storage is not too much behind on chunk persistence.")

	storageDirty          = flag.Bool("storage.local.dirty", false, "If set, the local storage layer will perform crash recovery even if the last shutdown appears to be clean.")
	storagePedanticChecks = flag.Bool("storage.local.pedantic-checks", false, "If set, a crash recovery will perform checks on each series file. This might take a very long time.")

	pathPrefix = flag.String("web.path-prefix", "", "Prefix for all web paths.")

	printVersion = flag.Bool("version", false, "Print version information.")
)

type prometheus struct {
	queryEngine         *promql.Engine
	ruleManager         *rules.Manager
	targetManager       *retrieval.TargetManager
	notificationHandler *notification.NotificationHandler
	statusHandler       *web.PrometheusStatusHandler
	storage             local.Storage
	remoteStorageQueues []*remote.StorageQueueManager

	webService *web.WebService

	closeOnce sync.Once
}

// NewPrometheus creates a new prometheus object based on flag values.
// Call Serve() to start serving and Close() for clean shutdown.
func NewPrometheus() *prometheus {
	notificationHandler := notification.NewNotificationHandler(*alertmanagerURL, *notificationQueueCapacity)

	var syncStrategy local.SyncStrategy
	switch *seriesSyncStrategy {
	case "never":
		syncStrategy = local.Never
	case "always":
		syncStrategy = local.Always
	case "adaptive":
		syncStrategy = local.Adaptive
	default:
		log.Errorf("Invalid flag value for 'storage.local.series-sync-strategy': %s\n", *seriesSyncStrategy)
		os.Exit(2)
	}

	o := &local.MemorySeriesStorageOptions{
		MemoryChunks:               *numMemoryChunks,
		MaxChunksToPersist:         *maxChunksToPersist,
		PersistenceStoragePath:     *persistenceStoragePath,
		PersistenceRetentionPeriod: *persistenceRetentionPeriod,
		CheckpointInterval:         *checkpointInterval,
		CheckpointDirtySeriesLimit: *checkpointDirtySeriesLimit,
		Dirty:          *storageDirty,
		PedanticChecks: *storagePedanticChecks,
		SyncStrategy:   syncStrategy,
	}
	memStorage := local.NewMemorySeriesStorage(o)

	var sampleAppender storage.SampleAppender
	var remoteStorageQueues []*remote.StorageQueueManager
	if *opentsdbURL == "" && *influxdbURL == "" {
		log.Warnf("No remote storage URLs provided; not sending any samples to long-term storage")
		sampleAppender = memStorage
	} else {
		fanout := storage.Fanout{memStorage}

		addRemoteStorage := func(c remote.StorageClient) {
			qm := remote.NewStorageQueueManager(c, 100*1024)
			fanout = append(fanout, qm)
			remoteStorageQueues = append(remoteStorageQueues, qm)
		}

		if *opentsdbURL != "" {
			addRemoteStorage(opentsdb.NewClient(*opentsdbURL, *remoteStorageTimeout))
		}
		if *influxdbURL != "" {
			addRemoteStorage(influxdb.NewClient(*influxdbURL, *remoteStorageTimeout))
		}

		sampleAppender = fanout
	}

	targetManager := retrieval.NewTargetManager(sampleAppender)

	queryEngine := promql.NewEngine(memStorage)

	ruleManager := rules.NewManager(&rules.ManagerOptions{
		SampleAppender:      sampleAppender,
		NotificationHandler: notificationHandler,
		QueryEngine:         queryEngine,
		PrometheusURL:       web.MustBuildServerURL(*pathPrefix),
		PathPrefix:          *pathPrefix,
	})

	flags := map[string]string{}
	flag.VisitAll(func(f *flag.Flag) {
		flags[f.Name] = f.Value.String()
	})
	prometheusStatus := &web.PrometheusStatusHandler{
		BuildInfo:   BuildInfo,
		RuleManager: ruleManager,
		TargetPools: targetManager.Pools,
		Flags:       flags,
		Birth:       time.Now(),
		PathPrefix:  *pathPrefix,
	}

	alertsHandler := &web.AlertsHandler{
		RuleManager: ruleManager,
		PathPrefix:  *pathPrefix,
	}

	consolesHandler := &web.ConsolesHandler{
		QueryEngine: queryEngine,
		PathPrefix:  *pathPrefix,
	}

	graphsHandler := &web.GraphsHandler{
		PathPrefix: *pathPrefix,
	}

	apiLegacy := &legacy.API{
		Now:         clientmodel.Now,
		Storage:     memStorage,
		QueryEngine: queryEngine,
	}

	apiv1 := &v1.API{
		Storage:     memStorage,
		QueryEngine: queryEngine,
	}

	webService := web.NewWebService(&web.WebServiceOptions{
		PathPrefix:      *pathPrefix,
		StatusHandler:   prometheusStatus,
		APILegacy:       apiLegacy,
		APIv1:           apiv1,
		ConsolesHandler: consolesHandler,
		AlertsHandler:   alertsHandler,
		GraphsHandler:   graphsHandler,
	})

	p := &prometheus{
		queryEngine:         queryEngine,
		ruleManager:         ruleManager,
		targetManager:       targetManager,
		notificationHandler: notificationHandler,
		statusHandler:       prometheusStatus,
		storage:             memStorage,
		remoteStorageQueues: remoteStorageQueues,

		webService: webService,
	}

	if !p.reloadConfig() {
		os.Exit(1)
	}

	return p
}

func (p *prometheus) reloadConfig() bool {
	log.Infof("Loading configuration file %s", *configFile)

	conf, err := config.LoadFromFile(*configFile)
	if err != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, err)
		log.Errorf("Note: The configuration format has changed with version 0.14. Please see the documentation (http://prometheus.io/docs/operating/configuration/) and the provided configuration migration tool (https://github.com/prometheus/migrate).")
		return false
	}
	success := true

	success = success && p.statusHandler.ApplyConfig(conf)
	success = success && p.targetManager.ApplyConfig(conf)
	success = success && p.ruleManager.ApplyConfig(conf)

	return success
}

// Serve starts the Prometheus server. It returns after the server has been shut
// down. The method installs an interrupt handler, allowing to trigger a
// shutdown by sending SIGTERM to the process.
func (p *prometheus) Serve() {
	// Start all components.
	if err := p.storage.Start(); err != nil {
		log.Errorln("Error opening memory series storage:", err)
		os.Exit(1)
	}
	defer func() {
		if err := p.storage.Stop(); err != nil {
			log.Errorln("Error stopping storage:", err)
		}
	}()

	// The storage has to be fully initialized before registering Prometheus.
	registry.MustRegister(p)

	for _, q := range p.remoteStorageQueues {
		go q.Run()
		defer q.Stop()
	}

	go p.ruleManager.Run()
	defer p.ruleManager.Stop()

	go p.notificationHandler.Run()
	defer p.notificationHandler.Stop()

	go p.targetManager.Run()
	defer p.targetManager.Stop()

	defer p.queryEngine.Stop()

	go p.webService.Run()

	// Wait for reload or termination signals.
	hup := make(chan os.Signal)
	signal.Notify(hup, syscall.SIGHUP)
	go func() {
		for range hup {
			p.reloadConfig()
		}
	}()

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	select {
	case <-term:
		log.Warn("Received SIGTERM, exiting gracefully...")
	case <-p.webService.QuitChan:
		log.Warn("Received termination request via web service, exiting gracefully...")
	}

	close(hup)

	log.Info("See you next time!")
}

// Describe implements registry.Collector.
func (p *prometheus) Describe(ch chan<- *registry.Desc) {
	p.notificationHandler.Describe(ch)
	p.storage.Describe(ch)
	for _, q := range p.remoteStorageQueues {
		q.Describe(ch)
	}
}

// Collect implements registry.Collector.
func (p *prometheus) Collect(ch chan<- registry.Metric) {
	p.notificationHandler.Collect(ch)
	p.storage.Collect(ch)
	for _, q := range p.remoteStorageQueues {
		q.Collect(ch)
	}
}

func usage() {
	groups := make(map[string][]*flag.Flag)
	// Set a default group for ungrouped flags.
	groups["."] = make([]*flag.Flag, 0)

	// Bucket flags into groups based on the first of their dot-separated levels.
	flag.VisitAll(func(fl *flag.Flag) {
		parts := strings.SplitN(fl.Name, ".", 2)
		if len(parts) == 1 {
			groups["."] = append(groups["."], fl)
		} else {
			name := parts[0]
			groups[name] = append(groups[name], fl)
		}
	})

	groupsOrdered := make(sort.StringSlice, 0, len(groups))
	for groupName := range groups {
		groupsOrdered = append(groupsOrdered, groupName)
	}
	sort.Sort(groupsOrdered)

	fmt.Fprintf(os.Stderr, "Usage: %s [options ...]:\n\n", os.Args[0])

	const (
		maxLineLength = 80
		lineSep       = "\n      "
	)
	for _, groupName := range groupsOrdered {
		if groupName != "." {
			fmt.Fprintf(os.Stderr, "\n%s:\n", strings.Title(groupName))
		}

		for _, fl := range groups[groupName] {
			format := "  -%s=%s"
			if strings.Contains(fl.DefValue, " ") || fl.DefValue == "" {
				format = "  -%s=%q"
			}
			flagUsage := fmt.Sprintf(format+lineSep, fl.Name, fl.DefValue)

			// Format the usage text to not exceed maxLineLength characters per line.
			words := strings.SplitAfter(fl.Usage, " ")
			lineLength := len(lineSep) - 1
			for _, w := range words {
				if lineLength+len(w) > maxLineLength {
					flagUsage += lineSep
					lineLength = len(lineSep) - 1
				}
				flagUsage += w
				lineLength += len(w)
			}
			fmt.Fprintf(os.Stderr, "%s\n", flagUsage)
		}
	}
}

func main() {
	flag.CommandLine.Init(os.Args[0], flag.ContinueOnError)
	flag.CommandLine.Usage = usage

	if err := flag.CommandLine.Parse(os.Args[1:]); err != nil {
		if err != flag.ErrHelp {
			log.Errorf("Invalid command line arguments. Help: %s -h", os.Args[0])
		}
		os.Exit(2)
	}

	*pathPrefix = strings.TrimRight(*pathPrefix, "/")
	if *pathPrefix != "" && !strings.HasPrefix(*pathPrefix, "/") {
		*pathPrefix = "/" + *pathPrefix
	}

	versionInfoTmpl.Execute(os.Stdout, BuildInfo)

	if *printVersion {
		os.Exit(0)
	}

	p := NewPrometheus()
	p.Serve()
}
<commit_msg>Place storage under working directory by default.<commit_after>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	_ "net/http/pprof" // Comment this line to disable pprof endpoint.
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/prometheus/log"

	clientmodel "github.com/prometheus/client_golang/model"
	registry "github.com/prometheus/client_golang/prometheus"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/notification"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/rules"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/local"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/prometheus/prometheus/storage/remote/influxdb"
	"github.com/prometheus/prometheus/storage/remote/opentsdb"
	"github.com/prometheus/prometheus/web"
	"github.com/prometheus/prometheus/web/api/legacy"
	"github.com/prometheus/prometheus/web/api/v1"
)

const deletionBatchSize = 100

// Commandline flags.
var (
	configFile = flag.String("config.file", "prometheus.yml", "Prometheus configuration file name.")

	alertmanagerURL           = flag.String("alertmanager.url", "", "The URL of the alert manager to send notifications to.")
	notificationQueueCapacity = flag.Int("alertmanager.notification-queue-capacity", 100, "The capacity of the queue for pending alert manager notifications.")

	persistenceStoragePath = flag.String("storage.local.path", "data", "Base path for metrics storage.")

	opentsdbURL          = flag.String("storage.remote.opentsdb-url", "", "The URL of the remote OpenTSDB server to send samples to. None, if empty.")
	influxdbURL          = flag.String("storage.remote.influxdb-url", "", "The URL of the remote InfluxDB server to send samples to. None, if empty.")
	remoteStorageTimeout = flag.Duration("storage.remote.timeout", 30*time.Second, "The timeout to use when sending samples to the remote storage.")

	numMemoryChunks = flag.Int("storage.local.memory-chunks", 1024*1024, "How many chunks to keep in memory. While the size of a chunk is 1kiB, the total memory usage will be significantly higher than this value * 1kiB. Furthermore, for various reasons, more chunks might have to be kept in memory temporarily.")

	persistenceRetentionPeriod = flag.Duration("storage.local.retention", 15*24*time.Hour, "How long to retain samples in the local storage.")
	maxChunksToPersist         = flag.Int("storage.local.max-chunks-to-persist", 1024*1024, "How many chunks can be waiting for persistence before sample ingestion will stop. Many chunks waiting to be persisted will increase the checkpoint size.")

	checkpointInterval         = flag.Duration("storage.local.checkpoint-interval", 5*time.Minute, "The period at which the in-memory metrics and the chunks not yet persisted to series files are checkpointed.")
	checkpointDirtySeriesLimit = flag.Int("storage.local.checkpoint-dirty-series-limit", 5000, "If approx. that many time series are in a state that would require a recovery operation after a crash, a checkpoint is triggered, even if the checkpoint interval hasn't passed yet. A recovery operation requires a disk seek. The default limit intends to keep the recovery time below 1min even on spinning disks. With SSD, recovery is much faster, so you might want to increase this value in that case to avoid overly frequent checkpoints.")
	seriesSyncStrategy         = flag.String("storage.local.series-sync-strategy", "adaptive", "When to sync series files after modification. Possible values: 'never', 'always', 'adaptive'. Sync'ing slows down storage performance but reduces the risk of data loss in case of an OS crash. With the 'adaptive' strategy, series files are sync'd for as long as the storage is not too much behind on chunk persistence.")

	storageDirty          = flag.Bool("storage.local.dirty", false, "If set, the local storage layer will perform crash recovery even if the last shutdown appears to be clean.")
	storagePedanticChecks = flag.Bool("storage.local.pedantic-checks", false, "If set, a crash recovery will perform checks on each series file. This might take a very long time.")

	pathPrefix = flag.String("web.path-prefix", "", "Prefix for all web paths.")

	printVersion = flag.Bool("version", false, "Print version information.")
)

type prometheus struct {
	queryEngine         *promql.Engine
	ruleManager         *rules.Manager
	targetManager       *retrieval.TargetManager
	notificationHandler *notification.NotificationHandler
	statusHandler       *web.PrometheusStatusHandler
	storage             local.Storage
	remoteStorageQueues []*remote.StorageQueueManager

	webService *web.WebService

	closeOnce sync.Once
}

// NewPrometheus creates a new prometheus object based on flag values.
// Call Serve() to start serving and Close() for clean shutdown.
func NewPrometheus() *prometheus {
	notificationHandler := notification.NewNotificationHandler(*alertmanagerURL, *notificationQueueCapacity)

	var syncStrategy local.SyncStrategy
	switch *seriesSyncStrategy {
	case "never":
		syncStrategy = local.Never
	case "always":
		syncStrategy = local.Always
	case "adaptive":
		syncStrategy = local.Adaptive
	default:
		log.Errorf("Invalid flag value for 'storage.local.series-sync-strategy': %s\n", *seriesSyncStrategy)
		os.Exit(2)
	}

	o := &local.MemorySeriesStorageOptions{
		MemoryChunks:               *numMemoryChunks,
		MaxChunksToPersist:         *maxChunksToPersist,
		PersistenceStoragePath:     *persistenceStoragePath,
		PersistenceRetentionPeriod: *persistenceRetentionPeriod,
		CheckpointInterval:         *checkpointInterval,
		CheckpointDirtySeriesLimit: *checkpointDirtySeriesLimit,
		Dirty:          *storageDirty,
		PedanticChecks: *storagePedanticChecks,
		SyncStrategy:   syncStrategy,
	}
	memStorage := local.NewMemorySeriesStorage(o)

	var sampleAppender storage.SampleAppender
	var remoteStorageQueues []*remote.StorageQueueManager
	if *opentsdbURL == "" && *influxdbURL == "" {
		log.Warnf("No remote storage URLs provided; not sending any samples to long-term storage")
		sampleAppender = memStorage
	} else {
		fanout := storage.Fanout{memStorage}

		addRemoteStorage := func(c remote.StorageClient) {
			qm := remote.NewStorageQueueManager(c, 100*1024)
			fanout = append(fanout, qm)
			remoteStorageQueues = append(remoteStorageQueues, qm)
		}

		if *opentsdbURL != "" {
			addRemoteStorage(opentsdb.NewClient(*opentsdbURL, *remoteStorageTimeout))
		}
		if *influxdbURL != "" {
			addRemoteStorage(influxdb.NewClient(*influxdbURL, *remoteStorageTimeout))
		}

		sampleAppender = fanout
	}

	targetManager := retrieval.NewTargetManager(sampleAppender)

	queryEngine := promql.NewEngine(memStorage)

	ruleManager := rules.NewManager(&rules.ManagerOptions{
		SampleAppender:      sampleAppender,
		NotificationHandler: notificationHandler,
		QueryEngine:         queryEngine,
		PrometheusURL:       web.MustBuildServerURL(*pathPrefix),
		PathPrefix:          *pathPrefix,
	})

	flags := map[string]string{}
	flag.VisitAll(func(f *flag.Flag) {
		flags[f.Name] = f.Value.String()
	})
	prometheusStatus := &web.PrometheusStatusHandler{
		BuildInfo:   BuildInfo,
		RuleManager: ruleManager,
		TargetPools: targetManager.Pools,
		Flags:       flags,
		Birth:       time.Now(),
		PathPrefix:  *pathPrefix,
	}

	alertsHandler := &web.AlertsHandler{
		RuleManager: ruleManager,
		PathPrefix:  *pathPrefix,
	}

	consolesHandler := &web.ConsolesHandler{
		QueryEngine: queryEngine,
		PathPrefix:  *pathPrefix,
	}

	graphsHandler := &web.GraphsHandler{
		PathPrefix: *pathPrefix,
	}

	apiLegacy := &legacy.API{
		Now:         clientmodel.Now,
		Storage:     memStorage,
		QueryEngine: queryEngine,
	}

	apiv1 := &v1.API{
		Storage:     memStorage,
		QueryEngine: queryEngine,
	}

	webService := web.NewWebService(&web.WebServiceOptions{
		PathPrefix:      *pathPrefix,
		StatusHandler:   prometheusStatus,
		APILegacy:       apiLegacy,
		APIv1:           apiv1,
		ConsolesHandler: consolesHandler,
		AlertsHandler:   alertsHandler,
		GraphsHandler:   graphsHandler,
	})

	p := &prometheus{
		queryEngine:         queryEngine,
		ruleManager:         ruleManager,
		targetManager:       targetManager,
		notificationHandler: notificationHandler,
		statusHandler:       prometheusStatus,
		storage:             memStorage,
		remoteStorageQueues: remoteStorageQueues,

		webService: webService,
	}

	if !p.reloadConfig() {
		os.Exit(1)
	}

	return p
}

func (p *prometheus) reloadConfig() bool {
	log.Infof("Loading configuration file %s", *configFile)

	conf, err := config.LoadFromFile(*configFile)
	if err != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, err)
		log.Errorf("Note: The configuration format has changed with version 0.14. Please see the documentation (http://prometheus.io/docs/operating/configuration/) and the provided configuration migration tool (https://github.com/prometheus/migrate).")
		return false
	}
	success := true

	success = success && p.statusHandler.ApplyConfig(conf)
	success = success && p.targetManager.ApplyConfig(conf)
	success = success && p.ruleManager.ApplyConfig(conf)

	return success
}

// Serve starts the Prometheus server. It returns after the server has been shut
// down. The method installs an interrupt handler, allowing to trigger a
// shutdown by sending SIGTERM to the process.
func (p *prometheus) Serve() {
	// Start all components.
	if err := p.storage.Start(); err != nil {
		log.Errorln("Error opening memory series storage:", err)
		os.Exit(1)
	}
	defer func() {
		if err := p.storage.Stop(); err != nil {
			log.Errorln("Error stopping storage:", err)
		}
	}()

	// The storage has to be fully initialized before registering Prometheus.
	registry.MustRegister(p)

	for _, q := range p.remoteStorageQueues {
		go q.Run()
		defer q.Stop()
	}

	go p.ruleManager.Run()
	defer p.ruleManager.Stop()

	go p.notificationHandler.Run()
	defer p.notificationHandler.Stop()

	go p.targetManager.Run()
	defer p.targetManager.Stop()

	defer p.queryEngine.Stop()

	go p.webService.Run()

	// Wait for reload or termination signals.
	hup := make(chan os.Signal)
	signal.Notify(hup, syscall.SIGHUP)
	go func() {
		for range hup {
			p.reloadConfig()
		}
	}()

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	select {
	case <-term:
		log.Warn("Received SIGTERM, exiting gracefully...")
	case <-p.webService.QuitChan:
		log.Warn("Received termination request via web service, exiting gracefully...")
	}

	close(hup)

	log.Info("See you next time!")
}

// Describe implements registry.Collector.
func (p *prometheus) Describe(ch chan<- *registry.Desc) {
	p.notificationHandler.Describe(ch)
	p.storage.Describe(ch)
	for _, q := range p.remoteStorageQueues {
		q.Describe(ch)
	}
}

// Collect implements registry.Collector.
func (p *prometheus) Collect(ch chan<- registry.Metric) {
	p.notificationHandler.Collect(ch)
	p.storage.Collect(ch)
	for _, q := range p.remoteStorageQueues {
		q.Collect(ch)
	}
}

func usage() {
	groups := make(map[string][]*flag.Flag)
	// Set a default group for ungrouped flags.
	groups["."] = make([]*flag.Flag, 0)

	// Bucket flags into groups based on the first of their dot-separated levels.
	flag.VisitAll(func(fl *flag.Flag) {
		parts := strings.SplitN(fl.Name, ".", 2)
		if len(parts) == 1 {
			groups["."] = append(groups["."], fl)
		} else {
			name := parts[0]
			groups[name] = append(groups[name], fl)
		}
	})

	groupsOrdered := make(sort.StringSlice, 0, len(groups))
	for groupName := range groups {
		groupsOrdered = append(groupsOrdered, groupName)
	}
	sort.Sort(groupsOrdered)

	fmt.Fprintf(os.Stderr, "Usage: %s [options ...]:\n\n", os.Args[0])

	const (
		maxLineLength = 80
		lineSep       = "\n      "
	)
	for _, groupName := range groupsOrdered {
		if groupName != "." {
			fmt.Fprintf(os.Stderr, "\n%s:\n", strings.Title(groupName))
		}

		for _, fl := range groups[groupName] {
			format := "  -%s=%s"
			if strings.Contains(fl.DefValue, " ") || fl.DefValue == "" {
				format = "  -%s=%q"
			}
			flagUsage := fmt.Sprintf(format+lineSep, fl.Name, fl.DefValue)

			// Format the usage text to not exceed maxLineLength characters per line.
			words := strings.SplitAfter(fl.Usage, " ")
			lineLength := len(lineSep) - 1
			for _, w := range words {
				if lineLength+len(w) > maxLineLength {
					flagUsage += lineSep
					lineLength = len(lineSep) - 1
				}
				flagUsage += w
				lineLength += len(w)
			}
			fmt.Fprintf(os.Stderr, "%s\n", flagUsage)
		}
	}
}

func main() {
	flag.CommandLine.Init(os.Args[0], flag.ContinueOnError)
	flag.CommandLine.Usage = usage

	if err := flag.CommandLine.Parse(os.Args[1:]); err != nil {
		if err != flag.ErrHelp {
			log.Errorf("Invalid command line arguments. Help: %s -h", os.Args[0])
		}
		os.Exit(2)
	}

	*pathPrefix = strings.TrimRight(*pathPrefix, "/")
	if *pathPrefix != "" && !strings.HasPrefix(*pathPrefix, "/") {
		*pathPrefix = "/" + *pathPrefix
	}

	versionInfoTmpl.Execute(os.Stdout, BuildInfo)

	if *printVersion {
		os.Exit(0)
	}

	p := NewPrometheus()
	p.Serve()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
)

var (
	flagVerbose  = flag.Bool("v", false, "be verbose")
	flagHTTPPort = flag.Int("hport", 80, "HTTP server port")
	flagDNSPort  = flag.Int("dport", 53, "DNS server port")
	cntMsgs      = expvar.NewInt("statsQuestions")
	cntRelayed   = expvar.NewInt("statsRelayed")
	cntBlocked   = expvar.NewInt("statsBlocked")
	cntServed    = expvar.NewInt("statsServed")
	cntErrors    = expvar.NewInt("statsErrors")
	pixel        = "\x47\x49\x46\x38\x39\x61\x01\x00\x01\x00\x80\x00\x00\xff\xff" +
		"\xff\x00\x00\x00\x21\xf9\x04\x01\x00\x00\x00\x00\x2c\x00\x00" +
		"\x00\x00\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b"
	exts = map[string]bool{
		"jpg":  true,
		"jpeg": true,
		"png":  true,
		"gif":  true,
	}
	proxy    *net.UDPConn
	upstream *net.UDPConn
	queries  map[int]*net.UDPAddr
	blocked  map[string]bool
	answer   []byte
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] upstream proxy list.txt\n\n"+
			"upstream - real upstream DNS address, e.g. 8.8.8.8\n"+
			"proxy    - your address, e.g. 127.0.0.1\n"+
			"list.txt - text file with addresses to kill\n\n",
			os.Args[0],
		)
		flag.PrintDefaults()
		return
	}
	flag.Parse()

	if len(flag.Args()) < 3 {
		flag.Usage()
		os.Exit(1)
	}

	upIP := net.ParseIP(flag.Arg(0))
	if upIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse upstream IP '%s'\n", flag.Arg(0))
		os.Exit(2)
	}

	proxyIP := net.ParseIP(flag.Arg(1))
	if proxyIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse proxy IP '%s'\n", flag.Arg(1))
		os.Exit(2)
	}

	answer = []byte("\x00\x01\x00\x01\xff\xff\xff\xff\x00\x04")
	answerIP := proxyIP.To4()
	if answerIP == nil {
		fmt.Fprintln(os.Stderr, "ERROR: IPv6 is not supported, sorry")
		os.Exit(3)
	}
	answer = append(answer, answerIP...)

	parseList(flag.Arg(2))

	var err error
	upAddr := &net.UDPAddr{IP: upIP, Port: 53}
	upstream, err = net.DialUDP("udp", nil, upAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer upstream.Close()

	proxyAddr := &net.UDPAddr{IP: proxyIP, Port: *flagDNSPort}
	proxy, err = net.ListenUDP("udp", proxyAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer proxy.Close()

	queries = make(map[int]*net.UDPAddr, 4096)

	go runServerHTTP(flag.Arg(1))
	go runServerLocalDNS()
	go runServerUpstreamDNS()

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)

forever:
	for {
		select {
		case <-sig:
			log.Println("Signal received, stopping")
			break forever
		}
	}
}

func parseList(path string) {
	file, err := os.Open(path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR:", err)
		os.Exit(2)
	}
	defer file.Close()

	blocked = make(map[string]bool, 4096)
	counter := 0
	scn := bufio.NewScanner(file)
	for scn.Scan() {
		counter++
		blocked[scn.Text()+"."] = true
	}

	log.Printf("DNS: Parsed %d entries from list\n", counter)
	return
}

func runServerLocalDNS() {
	log.Println("DNS: Started local server at", proxy.LocalAddr())

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, addr, err := proxy.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		msg := make([]byte, n)
		copy(msg, buf[:n])
		cntMsgs.Add(1)
		go handleDNS(msg, addr)
	}

	panic("not reachable")
}

func runServerUpstreamDNS() {
	log.Println("DNS: Started upstream server")

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, _, err := upstream.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		id := int(uint16(buf[0])<<8 + uint16(buf[1]))
		if to, ok := queries[id]; ok {
			delete(queries, id)
			_, err := proxy.WriteTo(buf[:n], to)
			if err != nil {
				log.Println("DNS ERROR:", err)
				cntErrors.Add(1)
				continue
			}
			cntRelayed.Add(1)
			if *flagVerbose {
				log.Println("DNS: Relayed answer to query", id)
			}
		}
	}

	panic("not reachable")
}

func popPart(host string) string {
	parts := strings.Split(host, ".")
	if len(parts) < 3 {
		return ""
	}

	return strings.Join(parts[1:], ".")
}

func handleDNS(msg []byte, from *net.UDPAddr) {
	var domain bytes.Buffer
	var host string
	var block bool
	var try int

	id := int(uint16(msg[0])<<8 + uint16(msg[1]))
	if *flagVerbose {
		log.Printf("DNS: Query id %d from %s\n", id, from)
	}

	// peak query
	count := uint16(msg[5]) // question counter
	offset := uint16(12)    // point to first domain name

	if count != 1 {
		log.Printf("DNS WARN: Query id %d from %s has %d questions\n", id, from, count)
		return
	}

outer:
	for count > 0 {
	inner:
		for {
			length := int8(msg[offset])
			if length == 0 {
				break inner
			}
			offset++
			domain.WriteString(string(msg[offset : offset+uint16(length)]))
			domain.WriteString(".")
			offset += uint16(length)
		}
		host = domain.String()
		testHost := host
		try = 1

	test:
		for {
			if _, ok := blocked[testHost]; ok {
				block = true
				break outer
			}
			testHost = popPart(testHost)
			if testHost == "" {
				break test
			}
			try++
		}
		domain.Reset()

		offset += 4
		count--
	}
	// end peak query

	if block {
		// fake answer
		cntBlocked.Add(1)
		if *flagVerbose {
			log.Printf("DNS: Blocking (%d) %s\n", try, host)
		}

		msg[2] = uint8(129) // flags upper byte
		msg[3] = uint8(128) // flags lower byte
		msg[7] = uint8(1)   // answer counter

		res := append(msg, msg[12:12+1+len(host)]...)
		res = append(res, answer...)
		_, err := proxy.WriteTo(res, from)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			return
		}
		if *flagVerbose {
			log.Println("DNS: Sent fake answer")
		}
		return
		// end fake answer
	} else {
		if *flagVerbose {
			log.Println("DNS: Asking upstream")
		}
		queries[id] = from
		_, err := upstream.Write(msg)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			goto clean
		}
		return
	}

clean:
	delete(queries, id)

	return
}

func handleHTTP(w http.ResponseWriter, req *http.Request) {
	if *flagVerbose {
		log.Printf("HTTP: Request %s %s %s\n", req.Method, req.Host, req.RequestURI)
	}
	cntServed.Add(1)
	w.Header()["Content-type"] = []string{"image/gif"}
	io.WriteString(w, pixel)

	return
}

func runServerHTTP(host string) {
	addr := fmt.Sprintf("%s:%d", host, *flagHTTPPort)
	http.HandleFunc("/", handleHTTP)
	log.Println("HTTP: Started at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))

	panic("not reachable")
}

// vim: ts=4 sw=4 sts=4
<commit_msg>List reloading<commit_after>package main

import (
	"bufio"
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
)

var (
	flagVerbose  = flag.Bool("v", false, "be verbose")
	flagHTTPPort = flag.Int("hport", 80, "HTTP server port")
	flagDNSPort  = flag.Int("dport", 53, "DNS server port")
	cntMsgs      = expvar.NewInt("statsQuestions")
	cntRelayed   = expvar.NewInt("statsRelayed")
	cntBlocked   = expvar.NewInt("statsBlocked")
	cntServed    = expvar.NewInt("statsServed")
	cntErrors    = expvar.NewInt("statsErrors")
	cntRules     = expvar.NewInt("statsRules")
	pixel        = "\x47\x49\x46\x38\x39\x61\x01\x00\x01\x00\x80\x00\x00\xff\xff" +
		"\xff\x00\x00\x00\x21\xf9\x04\x01\x00\x00\x00\x00\x2c\x00\x00" +
		"\x00\x00\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b"
	exts = map[string]bool{
		"jpg":  true,
		"jpeg": true,
		"png":  true,
		"gif":  true,
	}
	proxy    *net.UDPConn
	upstream *net.UDPConn
	queries  map[int]*net.UDPAddr
	blocked  map[string]bool
	answer   []byte
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] upstream proxy list.txt\n\n"+
			"upstream - real upstream DNS address, e.g. 8.8.8.8\n"+
			"proxy    - your address, e.g. 127.0.0.1\n"+
			"list.txt - text file with addresses to kill\n\n",
			os.Args[0],
		)
		flag.PrintDefaults()
		return
	}
	flag.Parse()

	if len(flag.Args()) < 3 {
		flag.Usage()
		os.Exit(1)
	}

	upIP := net.ParseIP(flag.Arg(0))
	if upIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse upstream IP '%s'\n", flag.Arg(0))
		os.Exit(2)
	}

	proxyIP := net.ParseIP(flag.Arg(1))
	if proxyIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse proxy IP '%s'\n", flag.Arg(1))
		os.Exit(2)
	}

	answer = []byte("\x00\x01\x00\x01\xff\xff\xff\xff\x00\x04")
	answerIP := proxyIP.To4()
	if answerIP == nil {
		fmt.Fprintln(os.Stderr, "ERROR: IPv6 is not supported, sorry")
		os.Exit(3)
	}
	answer = append(answer, answerIP...)

	parseList(flag.Arg(2))

	var err error
	upAddr := &net.UDPAddr{IP: upIP, Port: 53}
	upstream, err = net.DialUDP("udp", nil, upAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer upstream.Close()

	proxyAddr := &net.UDPAddr{IP: proxyIP, Port: *flagDNSPort}
	proxy, err = net.ListenUDP("udp", proxyAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer proxy.Close()

	queries = make(map[int]*net.UDPAddr, 4096)

	go runServerHTTP(flag.Arg(1))
	go runServerLocalDNS()
	go runServerUpstreamDNS()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGUSR1)

forever:
	for {
		select {
		case s := <-sig:
			switch s {
			case syscall.SIGUSR1:
				log.Println("SIGUSR1 received, reloading rules")
			default:
				log.Println("Signal received, stopping")
				break forever
			}
		}
	}
}

func parseList(path string) {
	file, err := os.Open(path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR:", err)
		os.Exit(2)
	}
	defer file.Close()

	blocked = make(map[string]bool, 4096)
	counter := 0
	scn := bufio.NewScanner(file)
	for scn.Scan() {
		counter++
		blocked[scn.Text()+"."] = true
	}
	cntRules.Set(int64(counter))
	log.Printf("DNS: Parsed %d entries from list\n", counter)

	return
}

func runServerLocalDNS() {
	log.Println("DNS: Started local server at", proxy.LocalAddr())

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, addr, err := proxy.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		msg := make([]byte, n)
		copy(msg, buf[:n])
		cntMsgs.Add(1)
		go handleDNS(msg, addr)
	}

	panic("not reachable")
}

func runServerUpstreamDNS() {
	log.Println("DNS: Started upstream server")

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, _, err := upstream.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		id := int(uint16(buf[0])<<8 + uint16(buf[1]))
		if to, ok := queries[id]; ok {
			delete(queries, id)
			_, err := proxy.WriteTo(buf[:n], to)
			if err != nil {
				log.Println("DNS ERROR:", err)
				cntErrors.Add(1)
				continue
			}
			cntRelayed.Add(1)
			if *flagVerbose {
				log.Println("DNS: Relayed answer to query", id)
			}
		}
	}

	panic("not reachable")
}

func popPart(host string) string {
	parts := strings.Split(host, ".")
	if len(parts) < 3 {
		return ""
	}

	return strings.Join(parts[1:], ".")
}

func handleDNS(msg []byte, from *net.UDPAddr) {
	var domain bytes.Buffer
	var host string
	var block bool
	var try int

	id := int(uint16(msg[0])<<8 + uint16(msg[1]))
	if *flagVerbose {
		log.Printf("DNS: Query id %d from %s\n", id, from)
	}

	// peak query
	count := uint16(msg[5]) // question counter
	offset := uint16(12)    // point to first domain name

	if count != 1 {
		log.Printf("DNS WARN: Query id %d from %s has %d questions\n", id, from, count)
		return
	}

outer:
	for count > 0 {
	inner:
		for {
			length := int8(msg[offset])
			if length == 0 {
				break inner
			}
			offset++
			domain.WriteString(string(msg[offset : offset+uint16(length)]))
			domain.WriteString(".")
			offset += uint16(length)
		}
		host = domain.String()
		testHost := host
		try = 1

	test:
		for {
			if _, ok := blocked[testHost]; ok {
				block = true
				break outer
			}
			testHost = popPart(testHost)
			if testHost == "" {
				break test
			}
			try++
		}
		domain.Reset()

		offset += 4
		count--
	}
	// end peak query

	if block {
		// fake answer
		cntBlocked.Add(1)
		if *flagVerbose {
			log.Printf("DNS: Blocking (%d) %s\n", try, host)
		}

		msg[2] = uint8(129) // flags upper byte
		msg[3] = uint8(128) // flags lower byte
		msg[7] = uint8(1)   // answer counter

		res := append(msg, msg[12:12+1+len(host)]...)
		res = append(res, answer...)
		_, err := proxy.WriteTo(res, from)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			return
		}
		if *flagVerbose {
			log.Println("DNS: Sent fake answer")
		}
		return
		// end fake answer
	} else {
		if *flagVerbose {
			log.Println("DNS: Asking upstream")
		}
		queries[id] = from
		_, err := upstream.Write(msg)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			goto clean
		}
		return
	}

clean:
	delete(queries, id)

	return
}

func handleHTTP(w http.ResponseWriter, req *http.Request) {
	if *flagVerbose {
		log.Printf("HTTP: Request %s %s %s\n", req.Method, req.Host, req.RequestURI)
	}
	cntServed.Add(1)
	w.Header()["Content-type"] = []string{"image/gif"}
	io.WriteString(w, pixel)

	return
}

func runServerHTTP(host string) {
	addr := fmt.Sprintf("%s:%d", host, *flagHTTPPort)
	http.HandleFunc("/", handleHTTP)
	log.Println("HTTP: Started at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))

	panic("not reachable")
}

// vim: ts=4 sw=4 sts=4
<|endoftext|>


<commit_before>package main

/*
#include <windows.h>
void change_wallpaper(char path[255])
{
    SystemParametersInfo(0x0014, 0, path, 1);
}
*/
import "C"
import (
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"fmt"
	"unsafe"
)

func main() {
	resp, err := http.Get("https://cn.bing.com/")
	if err != nil {
		panic("Address unreachable " + err.Error())
	}
	webpagesrcbyte, _ := ioutil.ReadAll(resp.Body)
	webpagesrc := string(webpagesrcbyte)
	re := regexp.MustCompile("g_img=\\{url:\\s\"(.+\\.jpg)\"")
	imgurl := re.FindStringSubmatch(webpagesrc)[1]
	resp, err = http.Get("https://cn.bing.com" + imgurl)
	if err != nil {
		panic("Unable to download image " + err.Error())
	}
	fileName := imgurl[strings.LastIndex(imgurl, "/")+1:]
	imgdata, _ := ioutil.ReadAll(resp.Body)
	cdir, _ := os.Getwd()
	path := cdir + "/wallpapers/" + time.Now().Format("2006_01_02") + "_" + fileName
	err = ioutil.WriteFile(path, imgdata, os.ModeAppend)
	if err != nil {
		panic(err)
	}
	path = strings.Replace(path, "/", "\\", -1)
	fmt.Println(path)

	cs := C.CString(path)
	defer C.free(unsafe.Pointer(cs))
	C.change_wallpaper(cs)
}
<commit_msg>draw text on wallpaper<commit_after>package main

/*
#include <windows.h>
void change_wallpaper(char path[255])
{
    SystemParametersInfo(0x0014, 0, path, 1);
}
*/
import "C"
import (
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"fmt"
	"html"
	"log"
	"unsafe"
)

const (
	URL  = "https://cn.bing.com/"
	DURL = "https://cn.bing.com/cnhp/life?IID=%s&IG=%s" // detailed page url
)

func fetchWebPage(url string) string {
	log.Println("Connecting...")
	resp, err := http.Get(url)
	if err != nil {
		panic("Address unreachable " + err.Error())
	}
	webpagesrcbyte, _ := ioutil.ReadAll(resp.Body)
	log.Println("Fetched page")
	return string(webpagesrcbyte)
}

func saveImage(src string) string {
	re := regexp.MustCompile("g_img=\\{url:\\s\"(.+\\.jpg)\"")
	imgurl := re.FindStringSubmatch(src)[1]
	log.Println(imgurl)
	resp, err := http.Get("https://cn.bing.com" + imgurl)
	if err != nil {
		log.Println("Unable to download image " + err.Error())
		return ""
	}
	fileName := imgurl[strings.LastIndex(imgurl, "/")+1:]
	imgdata, _ := ioutil.ReadAll(resp.Body)
	cdir, _ := os.Getwd()
	path := cdir + "/wallpapers/" + time.Now().Format("2006_01_02") + "_" + fileName
	err = ioutil.WriteFile(path, imgdata, os.ModeAppend)
	if err != nil {
		panic(err)
	}
	return path
}

func setWindowsWallPaper(path string) {
	path = strings.Replace(path, "/", "\\", -1)
	fmt.Println(path)

	cs := C.CString(path)
	defer C.free(unsafe.Pointer(cs))
	C.change_wallpaper(cs)
}

func getIG(src string) string {
	re := regexp.MustCompile("IG:\"(\\w+?\\d+?)\"")
	found := re.FindStringSubmatch(src)
	if len(found) < 1 {
		return ""
	}
	log.Println(found[1])
	return found[1]
}

func getIID(src string) string {
	re := regexp.MustCompile("_iid=\"(\\w{4}\\.\\d{4})\">")
	found := re.FindStringSubmatch(src)
	if len(found) < 1 {
		return ""
	}
	log.Println(found[1])
	return found[1]
}

func getTitle(src string) string {
	re := regexp.MustCompile("
(.+?)
")
	found := re.FindStringSubmatch(src)
	if len(found) < 1 {
		return ""
	}
	return found[1]
}

func getLocation(src string) string {
	re := regexp.MustCompile("(.+?)")
	found := re.FindStringSubmatch(src)
	if len(found) < 1 {
		return ""
	}
	return found[1]
}

func getArticle(src string) (title string, subtitle string, body string) {
	re := regexp.MustCompile("
(.+?)
")
	m := re.FindStringSubmatch(src)
	if len(m) < 1 {
		return
	}
	title = m[1]
	re = regexp.MustCompile("
(.+?)
")
	subtitle = re.FindStringSubmatch(src)[1]
	re = regexp.MustCompile("
(.+?)
")
	body = re.FindStringSubmatch(src)[1]
	return
}

func main() {
	for {
		webpagesrc := fetchWebPage(URL)
		ioutil.WriteFile("log.html.txt", []byte(webpagesrc), 0666)

		path := saveImage(webpagesrc)
		if path == "" {
			continue
		}
		iid := getIID(webpagesrc)
		ig := getIG(webpagesrc)
		if ig == "" {
			continue
		}

		webpagesrc = fetchWebPage(fmt.Sprintf(DURL, iid, ig))
		fmt.Println(getTitle(webpagesrc))
		fmt.Println(getLocation(webpagesrc))
		_, _, a := getArticle(webpagesrc)
		img := readImage(path)
		a = html.UnescapeString(a)
		t := splitText(a)
		fmt.Println(t)
		drawText(img, t)
		cdir, _ := os.Getwd()
		setWindowsWallPaper(cdir + "/out.jpg")
		log.Println("Done")
		break
	}
}
<|endoftext|>


<commit_before>package main
package io
package encoding/csv

import "fmt"

type Num double

type Splitter struct {
	Numerical bool
	Value     Num
	Left      map[string]bool
	Right     map[string]bool
}

type Node struct {
	Left   *Node
	Right  *Node
	Parent *Node
}

type Feature struct {
	Data      []Num
	Numerical bool
	Mapping   map[string]Num
	Back      map[Num]String
	Name      String
}

func ParseData(input io.Reader) []Feature {
	data := make([]Feature,0,100)
	tsv := csv.NewReader(input)
	tsv.Comma = '\t'

	for {
		line, err := r.ReadString("\n")
		if err != nil {
			last = true
		}

		if last {
			break
		}
	}
}

}

func main() {

}
<commit_msg>building with feature matrix parseing<commit_after>package main

import (
	"encoding/csv"
	"flag"
	"io"
	"log"
	"strings"
	"strconv"
	"os"
)

type Num float64

type Splitter struct {
	Numerical bool
	Value     Num
	Left      map[string]bool
	Right     map[string]bool
}

type Node struct {
	Left     *Node
	Right    *Node
	Splitter *Node
}

//this data strucutre is related to the one in rf-ace. In the future we may wish to represent
//catagorical features usings ints to speed up hash look up etc.
type Feature struct {
	Data      []Num
	Missing   []bool
	Numerical bool
	Map       map[string]Num
	Back      map[Num]string
	Name      string
}

func NewFeature(record []string, capacity int) Feature {
	f := Feature{make([]Num, 0, capacity), make([]bool, 0, capacity), false, make(map[string]Num, capacity), make(map[Num]string, capacity), record[0]}
	switch record[0][0:3] {
	case "N:":
		//Numerical
		f.Numerical = true
		for i := 1; i < len(record); i++ {
			v, err := strconv.ParseFloat(record[i],64)
			if err != nil {
				f.Data = append(f.Data, Num(0))
				f.Missing = append(f.Missing, true)
				continue
			}
			f.Data = append(f.Data, Num(v))
			f.Missing = append(f.Missing, false)

		}

	default:
		//Assume Catagorical
		f.Numerical = false
		fvalue := Num(0.0)
		for i := 1; i < len(record); i++ {
			v := record[i]
			norm := strings.ToLower(v)
			if norm == "?" || norm == "nan" || norm == "na" || norm == "null" {
				f.Data = append(f.Data, Num(0))
				f.Missing = append(f.Missing, true)
				continue
			}
			nv, exsists := f.Map[v]
			if exsists == false {
				f.Map[v] = fvalue
				f.Back[fvalue] = v
				nv = fvalue
				fvalue += 1.0
			}
			f.Data = append(f.Data, Num(nv))
			f.Missing = append(f.Missing, false)

		}

	}
	return f

}

func ParseAFM(input io.Reader) []Feature {
	data := make([]Feature, 0, 100)
	tsv := csv.NewReader(input)
	tsv.Comma = '\t'
	_, err := tsv.Read()
	if err == io.EOF {
		return data
	} else if err != nil {
		log.Print("Error:", err)
		return data
	}
	capacity := tsv.FieldsPerRecord

	for {
		record, err := tsv.Read()
		if err == io.EOF {
			break
		} else if err != nil {
			log.Print("Error:", err)
			break
		}
		data = append(data, NewFeature(record, capacity))
	}
	return data
}

func main() {
	fm := flag.String("fm", "featurematrix.afm", "AFM formated feature matrix to use.")
	flag.Parse()
	datafile, err := os.Open(*fm) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer datafile.Close()

	data := ParseAFM(datafile)
	log.Print("Data file ",len(data)," by ", len(data[0].Data))

}
<|endoftext|>