<commit_before>package main

//----------------------------------------------------
// Gobuster -- by OJ Reeves
//
// A crap attempt at building something that resembles
// dirbuster or dirb using Go. The goal was to build
// a tool that would help learn Go and to actually do
// something useful. The idea of having this compile
// to native code is also appealing.
//
// Run: gobuster -h
//----------------------------------------------------

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
)

// A single result which comes from an individual web
// request.
type Result struct {
	Entity string
	Status int
	Extra  string
}

type PrintResultFunc func(s *State, r *Result)
type ProcessorFunc func(s *State, entity string, resultChan chan<- Result)

// Shim type for "set"
type IntSet struct {
	set map[int]bool
}

// Contains State that are read in from the command
// line when the program is invoked.
type State struct {
	Threads        int
	Wordlist       string
	Url            string
	Cookies        string
	Extensions     []string
	StatusCodes    IntSet
	Verbose        bool
	UseSlash       bool
	FollowRedirect bool
	Quiet          bool
	NoStatus       bool
	Expanded       bool
	Mode           string
	Printer        PrintResultFunc
	Processor      ProcessorFunc
	Client         *http.Client
}

type RedirectHandler struct {
	Transport http.RoundTripper
	State     *State
}

type RedirectError struct {
	StatusCode int
}

// Add an element to a set
func (set *IntSet) Add(i int) bool {
	_, found := set.set[i]
	set.set[i] = true
	return !found
}

// Test if an element is in a set
func (set *IntSet) Contains(i int) bool {
	_, found := set.set[i]
	return found
}

// Stringify the set
func (set *IntSet) Stringify() string {
	values := []string{}
	for s, _ := range set.set {
		values = append(values, strconv.Itoa(s))
	}
	return strings.Join(values, ",")
}

// Make a request to the given URL.
func MakeRequest(client *http.Client, fullUrl, cookie string) *int {
	req, err := http.NewRequest("GET", fullUrl, nil)

	if err != nil {
		return nil
	}

	if cookie != "" {
		req.Header.Set("Cookie", cookie)
	}

	resp, err := client.Do(req)

	if err != nil {
		if ue, ok := err.(*url.Error); ok {
			if re, ok := ue.Err.(*RedirectError); ok {
				return &re.StatusCode
			}
		}
		return nil
	}

	defer resp.Body.Close()

	return &resp.StatusCode
}

// Small helper to combine URL with URI then make a
// request to the generated location.
func GoGet(client *http.Client, url, uri, cookie string) *int {
	return MakeRequest(client, url+uri, cookie)
}

// Parse all the command line options into a settings
// instance for future use.
func ParseCmdLine() *State {
	var extensions string
	var codes string
	valid := true

	s := State{StatusCodes: IntSet{set: map[int]bool{}}}

	// Set up the variables we're interested in parsing.
	flag.IntVar(&s.Threads, "t", 10, "Number of concurrent threads")
	flag.StringVar(&s.Mode, "m", "dir", "Directory/File mode (dir) or DNS mode (dns)")
	flag.StringVar(&s.Wordlist, "w", "", "Path to the wordlist")
	flag.StringVar(&codes, "s", "200,204,301,302,307", "Positive status codes (dir mode only)")
	flag.StringVar(&s.Url, "u", "", "The target URL or Domain")
	flag.StringVar(&s.Cookies, "c", "", "Cookies to use for the requests (dir mode only)")
	flag.StringVar(&extensions, "x", "", "File extension(s) to search for (dir mode only)")
	flag.BoolVar(&s.Verbose, "v", false, "Verbose output (errors and IP addresses")
	flag.BoolVar(&s.FollowRedirect, "r", false, "Follow redirects")
	flag.BoolVar(&s.Quiet, "q", false, "Don't print the banner")
	flag.BoolVar(&s.Expanded, "e", false, "Expanded mode, print full URLs")
	flag.BoolVar(&s.NoStatus, "n", false, "Don't print status codes")
	flag.BoolVar(&s.UseSlash, "f", false, "Append a forward-slash to each directory request (dir mode only)")

	flag.Parse()

	switch strings.ToLower(s.Mode) {
	case "dir":
		s.Printer = PrintDirResult
		s.Processor = ProcessDirEntry
	case "dns":
		s.Printer = PrintDnsResult
		s.Processor = ProcessDnsEntry
	default:
		fmt.Println("Mode (-m): Invalid value:", s.Mode)
		valid = false
	}

	if s.Threads < 0 {
		fmt.Println("Threads (-t): Invalid value:", s.Threads)
		valid = false
	}

	if s.Wordlist == "" {
		fmt.Println("WordList (-w): Must be specified")
		valid = false
	} else if _, err := os.Stat(s.Wordlist); os.IsNotExist(err) {
		fmt.Println("Wordlist (-w): File does not exist:", s.Wordlist)
		valid = false
	}

	if s.Url == "" {
		fmt.Println("Url/Domain (-u): Must be specified")
		valid = false
	}

	if s.Mode == "dir" {
		if strings.HasSuffix(s.Url, "/") == false {
			s.Url = s.Url + "/"
		}

		// extensions are comma seaprated
		if extensions != "" {
			s.Extensions = strings.Split(extensions, ",")
		}

		// status codes are comma seaprated
		if codes != "" {
			for _, c := range strings.Split(codes, ",") {
				i, err := strconv.Atoi(c)
				if err != nil {
					panic("Invalid status code given")
				}
				s.StatusCodes.Add(i)
			}
		}

		if valid {
			s.Client = &http.Client{
				Transport: &RedirectHandler{
					State: &s,
					Transport: &http.Transport{
						TLSClientConfig: &tls.Config{
							InsecureSkipVerify: true,
						},
					},
				}}

			if GoGet(s.Client, s.Url, "", s.Cookies) == nil {
				fmt.Println("[-] Unable to connect:", s.Url)
				valid = false
			}
		}
	}

	if valid {
		return &s
	}

	return nil
}

// Process the busting of the website with the given
// set of settings from the command line.
func Process(s *State) {
	wordlist, err := os.Open(s.Wordlist)
	if err != nil {
		panic("Failed to open wordlist")
	}

	// channels used for comms
	wordChan := make(chan string, s.Threads)
	resultChan := make(chan Result)

	// Use a wait group for waiting for all threads
	// to finish
	processorGroup := new(sync.WaitGroup)
	processorGroup.Add(s.Threads)
	printerGroup := new(sync.WaitGroup)
	printerGroup.Add(1)

	// Create goroutines for each of the number of threads
	// specified.
	for i := 0; i < s.Threads; i++ {
		go func() {
			for {
				word := <-wordChan

				// Did we reach the end? If so break.
				if word == "" {
					break
				}

				// Mode-specific processing
				s.Processor(s, word, resultChan)
			}

			// Indicate to the wait group that the thread
			// has finished.
			processorGroup.Done()
		}()
	}

	// Single goroutine which handles the results as they
	// appear from the worker threads.
	go func() {
		for r := range resultChan {
			s.Printer(s, &r)
		}
		printerGroup.Done()
	}()

	defer wordlist.Close()

	// Lazy reading of the wordlist line by line
	scanner := bufio.NewScanner(wordlist)
	for scanner.Scan() {
		word := scanner.Text()

		// Skip "comment" lines
		if strings.HasPrefix(word, "#") == false {
			wordChan <- word
		}
	}

	close(wordChan)
	processorGroup.Wait()
	close(resultChan)
	printerGroup.Wait()
}

func ProcessDnsEntry(s *State, word string, resultChan chan<- Result) {
	subdomain := word + "." + s.Url
	ips, err := net.LookupHost(subdomain)

	if err == nil {
		result := Result{
			Entity: subdomain,
		}
		if s.Verbose {
			result.Extra = strings.Join(ips, ", ")
		}
		resultChan <- result
	}
}

func ProcessDirEntry(s *State, word string, resultChan chan<- Result) {
	suffix := ""
	if s.UseSlash {
		suffix = "/"
	}

	// Try the DIR first
	dirResp := GoGet(s.Client, s.Url, word+suffix, s.Cookies)
	if dirResp != nil {
		resultChan <- Result{
			Entity: word + suffix,
			Status: *dirResp,
		}
	}

	// Follow up with files using each ext.
	for ext := range s.Extensions {
		file := word + s.Extensions[ext]
		fileResp := GoGet(s.Client, s.Url, file, s.Cookies)
		if fileResp != nil {
			resultChan <- Result{
				Entity: file,
				Status: *fileResp,
			}
		}
	}
}

func PrintDnsResult(s *State, r *Result) {
	if s.Verbose {
		fmt.Printf("Found: %s [%s]\n", r.Entity, r.Extra)
	} else {
		fmt.Printf("Found: %s\n", r.Entity)
	}
}

func PrintDirResult(s *State, r *Result) {
	output := ""

	// Prefix if we're in verbose mode
	if s.Verbose {
		if s.StatusCodes.Contains(r.Status) {
			output += "Found : "
		} else {
			output += "Missed: "
		}
	}

	if s.StatusCodes.Contains(r.Status) || s.Verbose {
		if s.Expanded {
			output += s.Url
		} else {
			output += "/"
		}
		output += r.Entity

		if !s.NoStatus {
			output += fmt.Sprintf(" (%d)", r.Status)
		}

		fmt.Println(output)
	}
}

func (e *RedirectError) Error() string {
	return fmt.Sprintf("Redirect code: %d", e.StatusCode)
}

func (rh *RedirectHandler) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	if rh.State.FollowRedirect {
		return rh.Transport.RoundTrip(req)
	}

	resp, err = rh.Transport.RoundTrip(req)
	if err != nil {
		return resp, err
	}

	switch resp.StatusCode {
	case http.StatusMovedPermanently, http.StatusFound, http.StatusSeeOther,
		http.StatusNotModified, http.StatusUseProxy, http.StatusTemporaryRedirect:
		return nil, &RedirectError{StatusCode: resp.StatusCode}
	}

	return resp, err
}

func Banner(state *State) {
	fmt.Println("\n=====================================================")
	fmt.Println("Gobuster v0.8 (DIR support by OJ Reeves @TheColonial)")
	fmt.Println("              (DNS support by Peleus     @0x42424242)")
	fmt.Println("=====================================================")

	if state != nil {
		fmt.Printf("[+] Mode         : %s\n", state.Mode)
		fmt.Printf("[+] Url/Domain   : %s\n", state.Url)
		fmt.Printf("[+] Threads      : %d\n", state.Threads)
		fmt.Printf("[+] Wordlist     : %s\n", state.Wordlist)

		if state.Mode == "dir" {
			fmt.Printf("[+] Status codes : %s\n", state.StatusCodes.Stringify())

			if state.Cookies != "" {
				fmt.Printf("[+] Cookies      : %s\n", state.Cookies)
			}

			if len(state.Extensions) > 0 {
				fmt.Printf("[+] Extensions   : %s\n", strings.Join(state.Extensions, ","))
			}

			if state.UseSlash {
				fmt.Printf("[+] Add Slash    : true\n")
			}

			if state.FollowRedirect {
				fmt.Printf("[+] Follow Redir : true\n")
			}

			if state.Expanded {
				fmt.Printf("[+] Expanded     : true\n")
			}

			if state.NoStatus {
				fmt.Printf("[+] No status    : true\n")
			}

			if state.Verbose {
				fmt.Printf("[+] Verbose      : true\n")
			}
		}
		fmt.Println("=====================================================")
	}
}

func main() {
	state := ParseCmdLine()
	if state.Quiet {
		Process(state)
	} else {
		Banner(state)
		Process(state)
		fmt.Println("=====================================================")
	}

}
<commit_msg>Re-add missing nil check<commit_after>package main

//----------------------------------------------------
// Gobuster -- by OJ Reeves
//
// A crap attempt at building something that resembles
// dirbuster or dirb using Go. The goal was to build
// a tool that would help learn Go and to actually do
// something useful. The idea of having this compile
// to native code is also appealing.
//
// Run: gobuster -h
//----------------------------------------------------

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
)

// A single result which comes from an individual web
// request.
type Result struct {
	Entity string
	Status int
	Extra  string
}

type PrintResultFunc func(s *State, r *Result)
type ProcessorFunc func(s *State, entity string, resultChan chan<- Result)

// Shim type for "set"
type IntSet struct {
	set map[int]bool
}

// Contains State that are read in from the command
// line when the program is invoked.
type State struct {
	Threads        int
	Wordlist       string
	Url            string
	Cookies        string
	Extensions     []string
	StatusCodes    IntSet
	Verbose        bool
	UseSlash       bool
	FollowRedirect bool
	Quiet          bool
	NoStatus       bool
	Expanded       bool
	Mode           string
	Printer        PrintResultFunc
	Processor      ProcessorFunc
	Client         *http.Client
}

type RedirectHandler struct {
	Transport http.RoundTripper
	State     *State
}

type RedirectError struct {
	StatusCode int
}

// Add an element to a set
func (set *IntSet) Add(i int) bool {
	_, found := set.set[i]
	set.set[i] = true
	return !found
}

// Test if an element is in a set
func (set *IntSet) Contains(i int) bool {
	_, found := set.set[i]
	return found
}

// Stringify the set
func (set *IntSet) Stringify() string {
	values := []string{}
	for s, _ := range set.set {
		values = append(values, strconv.Itoa(s))
	}
	return strings.Join(values, ",")
}

// Make a request to the given URL.
func MakeRequest(client *http.Client, fullUrl, cookie string) *int {
	req, err := http.NewRequest("GET", fullUrl, nil)

	if err != nil {
		return nil
	}

	if cookie != "" {
		req.Header.Set("Cookie", cookie)
	}

	resp, err := client.Do(req)

	if err != nil {
		if ue, ok := err.(*url.Error); ok {
			if re, ok := ue.Err.(*RedirectError); ok {
				return &re.StatusCode
			}
		}
		return nil
	}

	defer resp.Body.Close()

	return &resp.StatusCode
}

// Small helper to combine URL with URI then make a
// request to the generated location.
func GoGet(client *http.Client, url, uri, cookie string) *int {
	return MakeRequest(client, url+uri, cookie)
}

// Parse all the command line options into a settings
// instance for future use.
func ParseCmdLine() *State {
	var extensions string
	var codes string
	valid := true

	s := State{StatusCodes: IntSet{set: map[int]bool{}}}

	// Set up the variables we're interested in parsing.
	flag.IntVar(&s.Threads, "t", 10, "Number of concurrent threads")
	flag.StringVar(&s.Mode, "m", "dir", "Directory/File mode (dir) or DNS mode (dns)")
	flag.StringVar(&s.Wordlist, "w", "", "Path to the wordlist")
	flag.StringVar(&codes, "s", "200,204,301,302,307", "Positive status codes (dir mode only)")
	flag.StringVar(&s.Url, "u", "", "The target URL or Domain")
	flag.StringVar(&s.Cookies, "c", "", "Cookies to use for the requests (dir mode only)")
	flag.StringVar(&extensions, "x", "", "File extension(s) to search for (dir mode only)")
	flag.BoolVar(&s.Verbose, "v", false, "Verbose output (errors and IP addresses")
	flag.BoolVar(&s.FollowRedirect, "r", false, "Follow redirects")
	flag.BoolVar(&s.Quiet, "q", false, "Don't print the banner")
	flag.BoolVar(&s.Expanded, "e", false, "Expanded mode, print full URLs")
	flag.BoolVar(&s.NoStatus, "n", false, "Don't print status codes")
	flag.BoolVar(&s.UseSlash, "f", false, "Append a forward-slash to each directory request (dir mode only)")

	flag.Parse()

	switch strings.ToLower(s.Mode) {
	case "dir":
		s.Printer = PrintDirResult
		s.Processor = ProcessDirEntry
	case "dns":
		s.Printer = PrintDnsResult
		s.Processor = ProcessDnsEntry
	default:
		fmt.Println("Mode (-m): Invalid value:", s.Mode)
		valid = false
	}

	if s.Threads < 0 {
		fmt.Println("Threads (-t): Invalid value:", s.Threads)
		valid = false
	}

	if s.Wordlist == "" {
		fmt.Println("WordList (-w): Must be specified")
		valid = false
	} else if _, err := os.Stat(s.Wordlist); os.IsNotExist(err) {
		fmt.Println("Wordlist (-w): File does not exist:", s.Wordlist)
		valid = false
	}

	if s.Url == "" {
		fmt.Println("Url/Domain (-u): Must be specified")
		valid = false
	}

	if s.Mode == "dir" {
		if strings.HasSuffix(s.Url, "/") == false {
			s.Url = s.Url + "/"
		}

		// extensions are comma seaprated
		if extensions != "" {
			s.Extensions = strings.Split(extensions, ",")
		}

		// status codes are comma seaprated
		if codes != "" {
			for _, c := range strings.Split(codes, ",") {
				i, err := strconv.Atoi(c)
				if err != nil {
					panic("Invalid status code given")
				}
				s.StatusCodes.Add(i)
			}
		}

		if valid {
			s.Client = &http.Client{
				Transport: &RedirectHandler{
					State: &s,
					Transport: &http.Transport{
						TLSClientConfig: &tls.Config{
							InsecureSkipVerify: true,
						},
					},
				}}

			if GoGet(s.Client, s.Url, "", s.Cookies) == nil {
				fmt.Println("[-] Unable to connect:", s.Url)
				valid = false
			}
		}
	}

	if valid {
		return &s
	}

	return nil
}

// Process the busting of the website with the given
// set of settings from the command line.
func Process(s *State) {
	wordlist, err := os.Open(s.Wordlist)
	if err != nil {
		panic("Failed to open wordlist")
	}

	// channels used for comms
	wordChan := make(chan string, s.Threads)
	resultChan := make(chan Result)

	// Use a wait group for waiting for all threads
	// to finish
	processorGroup := new(sync.WaitGroup)
	processorGroup.Add(s.Threads)
	printerGroup := new(sync.WaitGroup)
	printerGroup.Add(1)

	// Create goroutines for each of the number of threads
	// specified.
	for i := 0; i < s.Threads; i++ {
		go func() {
			for {
				word := <-wordChan

				// Did we reach the end? If so break.
				if word == "" {
					break
				}

				// Mode-specific processing
				s.Processor(s, word, resultChan)
			}

			// Indicate to the wait group that the thread
			// has finished.
			processorGroup.Done()
		}()
	}

	// Single goroutine which handles the results as they
	// appear from the worker threads.
	go func() {
		for r := range resultChan {
			s.Printer(s, &r)
		}
		printerGroup.Done()
	}()

	defer wordlist.Close()

	// Lazy reading of the wordlist line by line
	scanner := bufio.NewScanner(wordlist)
	for scanner.Scan() {
		word := scanner.Text()

		// Skip "comment" lines
		if strings.HasPrefix(word, "#") == false {
			wordChan <- word
		}
	}

	close(wordChan)
	processorGroup.Wait()
	close(resultChan)
	printerGroup.Wait()
}

func ProcessDnsEntry(s *State, word string, resultChan chan<- Result) {
	subdomain := word + "." + s.Url
	ips, err := net.LookupHost(subdomain)

	if err == nil {
		result := Result{
			Entity: subdomain,
		}
		if s.Verbose {
			result.Extra = strings.Join(ips, ", ")
		}
		resultChan <- result
	}
}

func ProcessDirEntry(s *State, word string, resultChan chan<- Result) {
	suffix := ""
	if s.UseSlash {
		suffix = "/"
	}

	// Try the DIR first
	dirResp := GoGet(s.Client, s.Url, word+suffix, s.Cookies)
	if dirResp != nil {
		resultChan <- Result{
			Entity: word + suffix,
			Status: *dirResp,
		}
	}

	// Follow up with files using each ext.
	for ext := range s.Extensions {
		file := word + s.Extensions[ext]
		fileResp := GoGet(s.Client, s.Url, file, s.Cookies)
		if fileResp != nil {
			resultChan <- Result{
				Entity: file,
				Status: *fileResp,
			}
		}
	}
}

func PrintDnsResult(s *State, r *Result) {
	if s.Verbose {
		fmt.Printf("Found: %s [%s]\n", r.Entity, r.Extra)
	} else {
		fmt.Printf("Found: %s\n", r.Entity)
	}
}

func PrintDirResult(s *State, r *Result) {
	output := ""

	// Prefix if we're in verbose mode
	if s.Verbose {
		if s.StatusCodes.Contains(r.Status) {
			output += "Found : "
		} else {
			output += "Missed: "
		}
	}

	if s.StatusCodes.Contains(r.Status) || s.Verbose {
		if s.Expanded {
			output += s.Url
		} else {
			output += "/"
		}
		output += r.Entity

		if !s.NoStatus {
			output += fmt.Sprintf(" (%d)", r.Status)
		}

		fmt.Println(output)
	}
}

func (e *RedirectError) Error() string {
	return fmt.Sprintf("Redirect code: %d", e.StatusCode)
}

func (rh *RedirectHandler) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	if rh.State.FollowRedirect {
		return rh.Transport.RoundTrip(req)
	}

	resp, err = rh.Transport.RoundTrip(req)
	if err != nil {
		return resp, err
	}

	switch resp.StatusCode {
	case http.StatusMovedPermanently, http.StatusFound, http.StatusSeeOther,
		http.StatusNotModified, http.StatusUseProxy, http.StatusTemporaryRedirect:
		return nil, &RedirectError{StatusCode: resp.StatusCode}
	}

	return resp, err
}

func Banner(state *State) {
	fmt.Println("\n=====================================================")
	fmt.Println("Gobuster v0.8 (DIR support by OJ Reeves @TheColonial)")
	fmt.Println("              (DNS support by Peleus     @0x42424242)")
	fmt.Println("=====================================================")

	if state != nil {
		fmt.Printf("[+] Mode         : %s\n", state.Mode)
		fmt.Printf("[+] Url/Domain   : %s\n", state.Url)
		fmt.Printf("[+] Threads      : %d\n", state.Threads)
		fmt.Printf("[+] Wordlist     : %s\n", state.Wordlist)

		if state.Mode == "dir" {
			fmt.Printf("[+] Status codes : %s\n", state.StatusCodes.Stringify())

			if state.Cookies != "" {
				fmt.Printf("[+] Cookies      : %s\n", state.Cookies)
			}

			if len(state.Extensions) > 0 {
				fmt.Printf("[+] Extensions   : %s\n", strings.Join(state.Extensions, ","))
			}

			if state.UseSlash {
				fmt.Printf("[+] Add Slash    : true\n")
			}

			if state.FollowRedirect {
				fmt.Printf("[+] Follow Redir : true\n")
			}

			if state.Expanded {
				fmt.Printf("[+] Expanded     : true\n")
			}

			if state.NoStatus {
				fmt.Printf("[+] No status    : true\n")
			}

			if state.Verbose {
				fmt.Printf("[+] Verbose      : true\n")
			}
		}
		fmt.Println("=====================================================")
	}
}

func main() {
	state := ParseCmdLine()
	if state == nil {
		return
	}

	if state.Quiet {
		Process(state)
	} else {
		Banner(state)
		Process(state)
		fmt.Println("=====================================================")
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Dave Collins
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"encoding/csv"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"math/big"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrd/chaincfg/chainhash"
	"github.com/decred/dcrd/wire"
	"github.com/decred/dcrutil"
)

const (
	// fieldsPerRecord defines the number of fields expected in each line
	// of the input CSV data.
	fieldsPerRecord = 3
)

// convertRecord converts the passed record, which is expected to be parsed from
// a CSV file, and thus will be a slice of strings, into a struct with concrete
// types.
func convertRecord(record []string) (*simData, error) {
	headerBytes, err := hex.DecodeString(record[1])
	if err != nil {
		return nil, err
	}

	var header wire.BlockHeader
	if err := header.FromBytes(headerBytes); err != nil {
		return nil, err
	}
	var hashStrings []string
	if record[2] != "" {
		hashStrings = strings.Split(record[2], ":")
	}
	if len(hashStrings) != int(header.FreshStake) {
		return nil, fmt.Errorf("%d ticket hashes in CSV for %d new tickets",
			len(hashStrings), header.FreshStake)
	}
	ticketHashes := make([]chainhash.Hash, 0, len(hashStrings))
	for _, hashString := range hashStrings {
		hash, err := chainhash.NewHashFromStr(hashString)
		if err != nil {
			return nil, err
		}
		ticketHashes = append(ticketHashes, *hash)
	}

	return &simData{
		header:       headerBytes,
		voters:       header.Voters,
		prevValid:    dcrutil.IsFlagSet16(header.VoteBits, dcrutil.BlockValid),
		newTickets:   header.FreshStake,
		ticketHashes: ticketHashes,
		revocations:  uint16(header.Revocations),
	}, nil
}

// reportProgress periodically prints out the current simulator height to
// stdout.
func (s *simulator) reportProgress() {
	if s.tip.height%10000 == 0 && s.tip.height != 0 {
		fmt.Println()
	}
	if s.tip.height%1000 == 0 && s.tip.height != 0 {
		fmt.Printf("..%d", s.tip.height)
	}
}

// simulateFromCSV runs the simulation using input data from a CSV file.  It is
// realistically only intended to be used with data extracted from mainnet in
// order to exactly replicate its live ticket pool.
func (s *simulator) simulateFromCSV(csvPath string) error {
	// Open the simulation CSV data which is expected to be in the following
	// format:
	//
	// Block Header,Winning Ticket Hashes
	csvFile, err := os.Open(csvPath)
	if err != nil {
		return err
	}

	// Create a new simulator using input from the CSV file.
	r := csv.NewReader(csvFile)
	r.FieldsPerRecord = fieldsPerRecord
	var handledHeader bool
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		// Skip header fields if they exist.
		if !handledHeader {
			handledHeader = true
			_, err := strconv.Atoi(record[0])
			if err != nil {
				continue
			}
		}

		// Convert the CSV to concrete data.
		data, err := convertRecord(record)
		if err != nil {
			return err
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

// calcYieldDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// estimated yield purchasing a ticket would produce.
func calcYieldDemand(ticketPrice, perVoteSubsidy int64) float64 {
	// 100% demand when the yield is over 5%.
	yield := float64(perVoteSubsidy) / float64(ticketPrice)
	if yield > 0.05 {
		return 1.0
	}

	// No demand when the yield is under 2%.
	if yield < 0.02 {
		return 0.0
	}

	// The yield is in between 2% and 5%, so create a linear demand
	// accordingly.
	return (yield - 0.02) / 0.03
}

// calcVWAPDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// volume-weighted average ticket purchase of the previous ticket price windows.
func calcVWAPDemand(ticketPrice, ticketVWAP int64) float64 {
	// 100% demand when the ticket price is under 80% of the VWAP.
	eightyPercentVWAP := (ticketVWAP * 8) / 10
	if ticketPrice < eightyPercentVWAP {
		return 1.0
	}

	// No demand when the ticket price is over 120% of the VWAP.
	if ticketPrice > (ticketVWAP*12)/10 {
		return 0.0
	}

	// The ticket price is in between 80% and 120% of the VWAP, so create
	// a linear demand accordingly.
	fortyPercentVWAP := (ticketVWAP * 4) / 10
	return 1 - float64(ticketPrice-eightyPercentVWAP)/float64(fortyPercentVWAP)
}

// calcVWAP calculates and return the volume-weighted average ticket purchase
// price for up to 'StakeDiffWindows' worth of the previous ticket price
// windows.
func (s *simulator) calcPrevVWAP(prevNode *blockNode) int64 {
	windowSize := int32(s.params.StakeDiffWindowSize)
	stakeDiffWindows := int32(s.params.StakeDiffWindows)

	// Calculate the height the block just before the most recent ticket
	// price change.
	wantHeight := prevNode.height - (prevNode.height+1)%windowSize
	prevNode = s.ancestorNode(prevNode, wantHeight, nil)

	// Loop through previous required number of previous blocks and tally up
	// all of the weighted ticket purchase prices as well as the total
	// number of purchased tickets.
	numTickets, weightedVal := new(big.Int), new(big.Int)
	weightedSum, totalTickets := new(big.Int), new(big.Int)
	blocksToIterate := stakeDiffWindows * windowSize
	for i := int32(0); i < blocksToIterate && prevNode != nil; i++ {
		// weightedSum += numTickets*ticketPrice
		// totalTickets += numTickets
		numTickets.SetInt64(int64(len(prevNode.ticketsAdded)))
		weightedVal.SetInt64(prevNode.ticketPrice)
		weightedVal.Mul(weightedVal, numTickets)
		weightedSum.Add(weightedSum, weightedVal)
		totalTickets.Add(totalTickets, numTickets)
		prevNode = prevNode.parent
	}

	// Return minimum ticket price if there are were not any ticket
	// purchases at all in the entire period being examined.
	if totalTickets.Sign() == 0 {
		return s.params.MinimumStakeDiff
	}

	return new(big.Int).Div(weightedSum, totalTickets).Int64()
}

// calcDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval).
func (s *simulator) calcDemand(nextHeight int32, ticketPrice int64) float64 {
	// There is always 100% demand for minimum price tickets.
	if ticketPrice == s.params.MinimumStakeDiff {
		return 1.0
	}

	// Calculate the demand based on yield.
	ticketsPerBlock := s.params.TicketsPerBlock
	posSubsidy := s.calcPoSSubsidy(nextHeight - 1)
	perVoteSubsidy := posSubsidy / dcrutil.Amount(ticketsPerBlock)
	yieldDemand := calcYieldDemand(ticketPrice, int64(perVoteSubsidy))

	// Calculate the demand based on the volume-weighted average ticket
	// purchase price.
	currentVWAP := s.calcPrevVWAP(s.tip)
	vwapDemand := calcVWAPDemand(ticketPrice, currentVWAP)

	// The demand is the combination of the two unless there is full demand
	// based on yield and no demand based on the VWAP, in which case there
	// is 100% demand.
	demand := yieldDemand * vwapDemand
	if yieldDemand == 1.0 && vwapDemand == 0.0 {
		demand = 1.0
	}
	return demand
}

// simulate runs the simulation using a calculated demand curve which models
// how ticket purchasing would typically proceed based upon the price and the
// VWAP.
func (s *simulator) simulate(numBlocks uint64) error {
	// Shorter versions of some params for convenience.
	ticketMaturity := int32(s.params.TicketMaturity)
	ticketsPerBlock := s.params.TicketsPerBlock
	stakeValidationHeight := int32(s.params.StakeValidationHeight)
	stakeDiffWindowSize := int32(s.params.StakeDiffWindowSize)
	maxNewTicketsPerBlock := int32(s.params.MaxFreshStakePerBlock)
	maxTicketsPerWindow := maxNewTicketsPerBlock * stakeDiffWindowSize

	demandPerWindow := maxTicketsPerWindow
	for i := uint64(0); i < numBlocks; i++ {
		var nextHeight int32
		if s.tip != nil {
			nextHeight = s.tip.height + 1
		}

		// Purchase tickets according to simulated demand curve.
		//
		// When the height is prior to the stake validation height, just
		// use a 50% demand rate to ramp up the simulation.
		var newTickets uint8
		if nextHeight < stakeValidationHeight {
			if nextHeight >= ticketMaturity+1 {
				newTickets = uint8(maxNewTicketsPerBlock / 2)
			}
		} else {
			nextTicketPrice := s.nextTicketPriceFunc()
			if nextHeight%stakeDiffWindowSize == 0 && nextHeight != 0 {
				demand := s.calcDemand(nextHeight, nextTicketPrice)
				demandPerWindow = int32(float64(maxTicketsPerWindow) * demand)
			}

			newTickets = uint8(demandPerWindow / stakeDiffWindowSize)
			maxPossible := int64(s.spendableSupply) / nextTicketPrice
			if int64(newTickets) > maxPossible {
				newTickets = uint8(maxPossible)
			}
		}

		// TODO(davec): Account for tickets being purchased.
		// Limit the total staked coins to 40% of the total supply.
		stakedCoins := s.totalSupply - s.spendableSupply
		if newTickets > 0 && stakedCoins > (s.totalSupply*4/10) {
			newTickets = 0
		}

		// Start voting once stake validation height is reached.  This
		// assumes no votes are missed and revokes all expired tickets
		// as soon as possible which isn't very realistic, but it
		// doesn't have any effect on the ticket prices, so it's good
		// enough.  It could be useful to make this more realistic for
		// other simulation purposes though.
		var numVotes uint16
		if nextHeight >= stakeValidationHeight {
			numVotes = ticketsPerBlock
		}
		data := &simData{
			newTickets:  newTickets,
			prevValid:   true,
			revocations: uint16(len(s.unrevokedTickets)),
			voters:      numVotes,
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

func main() {
	var cpuProfilePath = flag.String("cpuprofile", "",
		"Write CPU profile to the specified file")
	var csvPath = flag.String("inputcsv", "",
		"Path to simulation CSV input data -- This overrides numblocks")
	var numBlocks = flag.Uint64("numblocks", 100000, "Number of blocks to simulate")
	flag.Parse()

	// Generate a CPU profile if requested.
	if *cpuProfilePath != "" {
		f, err := os.Create(*cpuProfilePath)
		if err != nil {
			fmt.Println("Unable to create cpu profile:", err)
			return
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// *********************************************************************
	// NOTE: Set a different function to calculate the next required stake
	// difficulty (aka ticket price) here.
	// *********************************************************************
	sim := newSimulator(&chaincfg.MainNetParams)
	sim.nextTicketPriceFunc = sim.curCalcNextStakeDiff

	startTime := time.Now()
	if *csvPath != "" {
		fmt.Printf("Running simulation from %q.\n", *csvPath)
		fmt.Printf("Height")
		if err := sim.simulateFromCSV(*csvPath); err != nil {
			fmt.Println(err)
			return
		}
	} else {
		fmt.Printf("Running simulation for %d blocks.\n", *numBlocks)
		fmt.Printf("Height")
		if err := sim.simulate(*numBlocks); err != nil {
			fmt.Println(err)
			return
		}
	}
	fmt.Println("..done")
	fmt.Println("Simulation took", time.Since(startTime))

	// Generate the simulation results and open them in a browser.
	if err := generateResults(sim); err != nil {
		fmt.Println(err)
		return
	}
}
<commit_msg>Simulate sudden surge and drop in stake coins.<commit_after>// Copyright (c) 2017 Dave Collins
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"encoding/csv"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"math/big"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrd/chaincfg/chainhash"
	"github.com/decred/dcrd/wire"
	"github.com/decred/dcrutil"
)

const (
	// fieldsPerRecord defines the number of fields expected in each line
	// of the input CSV data.
	fieldsPerRecord = 3
)

// convertRecord converts the passed record, which is expected to be parsed from
// a CSV file, and thus will be a slice of strings, into a struct with concrete
// types.
func convertRecord(record []string) (*simData, error) {
	headerBytes, err := hex.DecodeString(record[1])
	if err != nil {
		return nil, err
	}

	var header wire.BlockHeader
	if err := header.FromBytes(headerBytes); err != nil {
		return nil, err
	}
	var hashStrings []string
	if record[2] != "" {
		hashStrings = strings.Split(record[2], ":")
	}
	if len(hashStrings) != int(header.FreshStake) {
		return nil, fmt.Errorf("%d ticket hashes in CSV for %d new tickets",
			len(hashStrings), header.FreshStake)
	}
	ticketHashes := make([]chainhash.Hash, 0, len(hashStrings))
	for _, hashString := range hashStrings {
		hash, err := chainhash.NewHashFromStr(hashString)
		if err != nil {
			return nil, err
		}
		ticketHashes = append(ticketHashes, *hash)
	}

	return &simData{
		header:       headerBytes,
		voters:       header.Voters,
		prevValid:    dcrutil.IsFlagSet16(header.VoteBits, dcrutil.BlockValid),
		newTickets:   header.FreshStake,
		ticketHashes: ticketHashes,
		revocations:  uint16(header.Revocations),
	}, nil
}

// reportProgress periodically prints out the current simulator height to
// stdout.
func (s *simulator) reportProgress() {
	if s.tip.height%10000 == 0 && s.tip.height != 0 {
		fmt.Println()
	}
	if s.tip.height%1000 == 0 && s.tip.height != 0 {
		fmt.Printf("..%d", s.tip.height)
	}
}

// simulateFromCSV runs the simulation using input data from a CSV file.  It is
// realistically only intended to be used with data extracted from mainnet in
// order to exactly replicate its live ticket pool.
func (s *simulator) simulateFromCSV(csvPath string) error {
	// Open the simulation CSV data which is expected to be in the following
	// format:
	//
	// Block Header,Winning Ticket Hashes
	csvFile, err := os.Open(csvPath)
	if err != nil {
		return err
	}

	// Create a new simulator using input from the CSV file.
	r := csv.NewReader(csvFile)
	r.FieldsPerRecord = fieldsPerRecord
	var handledHeader bool
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		// Skip header fields if they exist.
		if !handledHeader {
			handledHeader = true
			_, err := strconv.Atoi(record[0])
			if err != nil {
				continue
			}
		}

		// Convert the CSV to concrete data.
		data, err := convertRecord(record)
		if err != nil {
			return err
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

// calcYieldDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// estimated yield purchasing a ticket would produce.
func calcYieldDemand(ticketPrice, perVoteSubsidy int64) float64 {
	// 100% demand when the yield is over 5%.
	yield := float64(perVoteSubsidy) / float64(ticketPrice)
	if yield > 0.05 {
		return 1.0
	}

	// No demand when the yield is under 2%.
	if yield < 0.02 {
		return 0.0
	}

	// The yield is in between 2% and 5%, so create a linear demand
	// accordingly.
	return (yield - 0.02) / 0.03
}

// calcVWAPDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// volume-weighted average ticket purchase of the previous ticket price windows.
func calcVWAPDemand(ticketPrice, ticketVWAP int64) float64 {
	// 100% demand when the ticket price is under 80% of the VWAP.
	eightyPercentVWAP := (ticketVWAP * 8) / 10
	if ticketPrice < eightyPercentVWAP {
		return 1.0
	}

	// No demand when the ticket price is over 120% of the VWAP.
	if ticketPrice > (ticketVWAP*12)/10 {
		return 0.0
	}

	// The ticket price is in between 80% and 120% of the VWAP, so create
	// a linear demand accordingly.
	fortyPercentVWAP := (ticketVWAP * 4) / 10
	return 1 - float64(ticketPrice-eightyPercentVWAP)/float64(fortyPercentVWAP)
}

// calcVWAP calculates and return the volume-weighted average ticket purchase
// price for up to 'StakeDiffWindows' worth of the previous ticket price
// windows.
func (s *simulator) calcPrevVWAP(prevNode *blockNode) int64 {
	windowSize := int32(s.params.StakeDiffWindowSize)
	stakeDiffWindows := int32(s.params.StakeDiffWindows)

	// Calculate the height the block just before the most recent ticket
	// price change.
	wantHeight := prevNode.height - (prevNode.height+1)%windowSize
	prevNode = s.ancestorNode(prevNode, wantHeight, nil)

	// Loop through previous required number of previous blocks and tally up
	// all of the weighted ticket purchase prices as well as the total
	// number of purchased tickets.
	numTickets, weightedVal := new(big.Int), new(big.Int)
	weightedSum, totalTickets := new(big.Int), new(big.Int)
	blocksToIterate := stakeDiffWindows * windowSize
	for i := int32(0); i < blocksToIterate && prevNode != nil; i++ {
		// weightedSum += numTickets*ticketPrice
		// totalTickets += numTickets
		numTickets.SetInt64(int64(len(prevNode.ticketsAdded)))
		weightedVal.SetInt64(prevNode.ticketPrice)
		weightedVal.Mul(weightedVal, numTickets)
		weightedSum.Add(weightedSum, weightedVal)
		totalTickets.Add(totalTickets, numTickets)
		prevNode = prevNode.parent
	}

	// Return minimum ticket price if there are were not any ticket
	// purchases at all in the entire period being examined.
	if totalTickets.Sign() == 0 {
		return s.params.MinimumStakeDiff
	}

	return new(big.Int).Div(weightedSum, totalTickets).Int64()
}

// calcDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval).
func (s *simulator) calcDemand(nextHeight int32, ticketPrice int64) float64 {
	// There is always 100% demand for minimum price tickets.
	if ticketPrice == s.params.MinimumStakeDiff {
		return 1.0
	}

	// Calculate the demand based on yield.
	ticketsPerBlock := s.params.TicketsPerBlock
	posSubsidy := s.calcPoSSubsidy(nextHeight - 1)
	perVoteSubsidy := posSubsidy / dcrutil.Amount(ticketsPerBlock)
	yieldDemand := calcYieldDemand(ticketPrice, int64(perVoteSubsidy))

	// Calculate the demand based on the volume-weighted average ticket
	// purchase price.
	currentVWAP := s.calcPrevVWAP(s.tip)
	vwapDemand := calcVWAPDemand(ticketPrice, currentVWAP)

	// The demand is the combination of the two unless there is full demand
	// based on yield and no demand based on the VWAP, in which case there
	// is 100% demand.
	demand := yieldDemand * vwapDemand
	if yieldDemand == 1.0 && vwapDemand == 0.0 {
		demand = 1.0
	}
	return demand
}

// simulate runs the simulation using a calculated demand curve which models
// how ticket purchasing would typically proceed based upon the price and the
// VWAP.
func (s *simulator) simulate(numBlocks uint64) error {
	// Shorter versions of some params for convenience.
	ticketMaturity := int32(s.params.TicketMaturity)
	ticketsPerBlock := s.params.TicketsPerBlock
	stakeValidationHeight := int32(s.params.StakeValidationHeight)
	stakeDiffWindowSize := int32(s.params.StakeDiffWindowSize)
	maxNewTicketsPerBlock := int32(s.params.MaxFreshStakePerBlock)
	maxTicketsPerWindow := maxNewTicketsPerBlock * stakeDiffWindowSize

	demandPerWindow := maxTicketsPerWindow
	for i := uint64(0); i < numBlocks; i++ {
		var nextHeight int32
		if s.tip != nil {
			nextHeight = s.tip.height + 1
		}

		// Purchase tickets according to simulated demand curve.
		//
		// When the height is prior to the stake validation height, just
		// use a 50% demand rate to ramp up the simulation.
		var newTickets uint8
		if nextHeight < stakeValidationHeight {
			if nextHeight >= ticketMaturity+1 {
				newTickets = uint8(maxNewTicketsPerBlock / 2)
			}
		} else {
			nextTicketPrice := s.nextTicketPriceFunc()
			if nextHeight%stakeDiffWindowSize == 0 && nextHeight != 0 {
				demand := s.calcDemand(nextHeight, nextTicketPrice)
				demandPerWindow = int32(float64(maxTicketsPerWindow) * demand)
			}

			newTickets = uint8(demandPerWindow / stakeDiffWindowSize)
			maxPossible := int64(s.spendableSupply) / nextTicketPrice
			if int64(newTickets) > maxPossible {
				newTickets = uint8(maxPossible)
			}
		}

		// TODO(davec): Account for tickets being purchased.
		// Limit the total staked coins to 40% of the total supply
		// except for in between blocks 80k and 100k which limit to 50%
		// of the total supply in order to simulate a sudden surge and
		// drop the amount of stake coins.
		stakedCoins := s.totalSupply - s.spendableSupply
		if nextHeight < 80000 || nextHeight > 100000 {
			if newTickets > 0 && stakedCoins > (s.totalSupply*4/10) {
				newTickets = 0
			}
		} else {
			if newTickets > 0 && stakedCoins > (s.totalSupply*5/10) {
				newTickets = 0
			}
		}

		// Start voting once stake validation height is reached.  This
		// assumes no votes are missed and revokes all expired tickets
		// as soon as possible which isn't very realistic, but it
		// doesn't have any effect on the ticket prices, so it's good
		// enough.  It could be useful to make this more realistic for
		// other simulation purposes though.
		var numVotes uint16
		if nextHeight >= stakeValidationHeight {
			numVotes = ticketsPerBlock
		}
		data := &simData{
			newTickets:  newTickets,
			prevValid:   true,
			revocations: uint16(len(s.unrevokedTickets)),
			voters:      numVotes,
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

func main() {
	var cpuProfilePath = flag.String("cpuprofile", "",
		"Write CPU profile to the specified file")
	var csvPath = flag.String("inputcsv", "",
		"Path to simulation CSV input data -- This overrides numblocks")
	var numBlocks = flag.Uint64("numblocks", 100000, "Number of blocks to simulate")
	flag.Parse()

	// Generate a CPU profile if requested.
	if *cpuProfilePath != "" {
		f, err := os.Create(*cpuProfilePath)
		if err != nil {
			fmt.Println("Unable to create cpu profile:", err)
			return
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// *********************************************************************
	// NOTE: Set a different function to calculate the next required stake
	// difficulty (aka ticket price) here.
	// *********************************************************************
	sim := newSimulator(&chaincfg.MainNetParams)
	sim.nextTicketPriceFunc = sim.curCalcNextStakeDiff

	startTime := time.Now()
	if *csvPath != "" {
		fmt.Printf("Running simulation from %q.\n", *csvPath)
		fmt.Printf("Height")
		if err := sim.simulateFromCSV(*csvPath); err != nil {
			fmt.Println(err)
			return
		}
	} else {
		fmt.Printf("Running simulation for %d blocks.\n", *numBlocks)
		fmt.Printf("Height")
		if err := sim.simulate(*numBlocks); err != nil {
			fmt.Println(err)
			return
		}
	}
	fmt.Println("..done")
	fmt.Println("Simulation took", time.Since(startTime))

	// Generate the simulation results and open them in a browser.
	if err := generateResults(sim); err != nil {
		fmt.Println(err)
		return
	}
}
<|endoftext|>


<commit_before>package libvirt

import (
	"fmt"
	"log"
	"net"
	"strings"
	"time"

	libvirt "github.com/dmacvicar/libvirt-go"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

const (
	netModeIsolated = "none"
	netModeNat      = "nat"
	netModeRoute    = "route"
	netModeBridge   = "bridge"
)

// a libvirt network resource
//
// Resource example:
//
// resource "libvirt_network" "k8snet" {
//    name = "k8snet"
//    domain = "k8s.local"
//    mode = "nat"
//    addresses = ["10.17.3.0/24"]
// }
//
// "addresses" can contain (0 or 1) ipv4 and (0 or 1) ipv6 ranges
// "mode" can be one of: "nat" (default), "isolated"
//
func resourceLibvirtNetwork() *schema.Resource {
	return &schema.Resource{
		Create: resourceLibvirtNetworkCreate,
		Read:   resourceLibvirtNetworkRead,
		Delete: resourceLibvirtNetworkDelete,
		Exists: resourceLibvirtNetworkExists,
		Update: resourceLibvirtNetworkUpdate,
		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"domain": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"mode": &schema.Schema{ // can be "none", "nat" (default), "route", "bridge"
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  netModeNat,
			},
			"bridge": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},
			"addresses": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Required: false,
				ForceNew: true,
				Elem: &schema.Schema{
					Type: schema.TypeString,
				},
			},
			"running": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
				ForceNew: false,
			},
		},
	}
}

func resourceLibvirtNetworkExists(d *schema.ResourceData, meta interface{}) (bool, error) {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return false, fmt.Errorf("The libvirt connection was nil.")
	}
	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	defer network.Free()
	return err == nil, err
}

func resourceLibvirtNetworkUpdate(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}
	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	defer network.Free()

	active, err := network.IsActive()
	if err != nil {
		return err
	}

	if !active {
		log.Printf("[DEBUG] Activating network")
		if err := network.Create(); err != nil {
			return err
		}
		d.Set("running", true)
	}

	return nil
}

func resourceLibvirtNetworkCreate(d *schema.ResourceData, meta interface{}) error {
	// see https://libvirt.org/formatnetwork.html
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	networkDef := newNetworkDef()
	networkDef.Name = d.Get("name").(string)
	networkDef.Domain = &defNetworkDomain{
		Name: d.Get("domain").(string),
	}

	// use a bridge provided by the user, or create one otherwise (libvirt will assign on automatically when empty)
	bridgeName := ""
	if b, ok := d.GetOk("bridge"); ok {
		bridgeName = b.(string)
	}
	networkDef.Bridge = &defNetworkBridge{
		Name: bridgeName,
		Stp:  "on",
	}

	// check the network mode
	networkDef.Forward.Mode = strings.ToLower(d.Get("mode").(string))
	if networkDef.Forward.Mode == netModeIsolated || networkDef.Forward.Mode == netModeNat || networkDef.Forward.Mode == netModeRoute {

		// there is no mode when using an isolated network
		if networkDef.Forward.Mode == netModeIsolated {
			networkDef.Forward = nil
		}

		// some network modes require a DHCP/DNS server
		// set the addresses for DHCP
		if addresses, ok := d.GetOk("addresses"); ok {
			ipsPtrsLst := []*defNetworkIp{}
			for _, addressI := range addresses.([]interface{}) {
				address := addressI.(string)
				_, ipNet, err := net.ParseCIDR(address)
				if err != nil {
					return fmt.Errorf("Error parsing addresses definition '%s': %s", address, err)
				}
				ones, bits := ipNet.Mask.Size()
				family := "ipv4"
				if bits == (net.IPv6len * 8) {
					family = "ipv6"
				}
				ipsRange := 2 ^ bits - 2 ^ ones
				if ipsRange < 4 {
					return fmt.Errorf("Netmask seems to be too strict: only %d IPs available (%s)", ipsRange-3, family)
				}

				// we should calculate the range served by DHCP. For example, for
				// 192.168.121.0/24 we will serve 192.168.121.2 - 192.168.121.254
				start, end := NetworkRange(ipNet)

				// skip the .0, (for the network),
				start[len(start)-1]++

				// assign the .1 to the host interface
				dni := defNetworkIp{
					Address: start.String(),
					Prefix:  ones,
					Family:  family,
				}

				start[len(start)-1]++ // then skip the .1
				end[len(end)-1]--     // and skip the .255 (for broadcast)

				dni.Dhcp = &defNetworkIpDhcp{
					Ranges: []*defNetworkIpDhcpRange{
						&defNetworkIpDhcpRange{
							Start: start.String(),
							End:   end.String(),
						},
					},
				}
				ipsPtrsLst = append(ipsPtrsLst, &dni)
			}
			networkDef.Ips = ipsPtrsLst
		}
	} else if networkDef.Forward.Mode == netModeBridge {
		if bridgeName == "" {
			return fmt.Errorf("'bridge' must be provided when using the bridged network mode")
		}
	} else {
		return fmt.Errorf("unsuppoorted network mode '%s'", networkDef.Forward.Mode)
	}

	// once we have the network defined, connect to libvirt and create it from the XML serialization
	connectURI, err := virConn.GetURI()
	if err != nil {
		return fmt.Errorf("Error retrieving libvirt connection URI: %s", err)
	}
	log.Printf("[INFO] Creating libvirt network at %s", connectURI)

	data, err := xmlMarshallIndented(networkDef)
	if err != nil {
		return fmt.Errorf("Error serializing libvirt network: %s", err)
	}

	log.Printf("[DEBUG] Creating libvirt network at %s: %s", connectURI, data)
	network, err := virConn.NetworkDefineXML(data)
	if err != nil {
		return fmt.Errorf("Error defining libvirt network: %s - %s", err, data)
	}
	err = network.Create()
	if err != nil {
		return fmt.Errorf("Error crearing libvirt network: %s", err)
	}
	defer network.Free()

	id, err := network.GetUUIDString()
	if err != nil {
		return fmt.Errorf("Error retrieving libvirt network id: %s", err)
	}
	d.SetId(id)

	log.Printf("[INFO] Created network %s [%s]", networkDef.Name, d.Id())

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"BUILD"},
		Target:     []string{"ACTIVE"},
		Refresh:    waitForNetworkActive(network),
		Timeout:    1 * time.Minute,
		Delay:      5 * time.Second,
		MinTimeout: 3 * time.Second,
	}
	_, err = stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf("Error waiting for network to reach ACTIVE state: %s", err)
	}

	return resourceLibvirtNetworkRead(d, meta)
}

func resourceLibvirtNetworkRead(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	if err != nil {
		return fmt.Errorf("Error retrieving libvirt network: %s", err)
	}
	defer network.Free()

	networkDef, err := newDefNetworkfromLibvirt(&network)
	if err != nil {
		return fmt.Errorf("Error reading libvirt network XML description: %s", err)
	}

	d.Set("name", networkDef.Name)
	d.Set("domain", networkDef.Domain.Name)
	d.Set("bridge", networkDef.Bridge.Name)

	active, err := network.IsActive()
	if err != nil {
		return err
	}
	d.Set("running", active)

	addresses := []string{}
	for _, address := range networkDef.Ips {
		// we get the host interface IP (ie, 10.10.8.1) but we want the network CIDR (ie, 10.10.8.0/24)
		// so we need some transformations...
		addr := net.ParseIP(address.Address)
		if addr == nil {
			return fmt.Errorf("Error parsing IP '%s': %s", address, err)
		}
		bits := net.IPv6len * 8
		if addr.To4() != nil {
			bits = net.IPv4len * 8
		}
		mask := net.CIDRMask(address.Prefix, bits)
		network := addr.Mask(mask)
		addresses = append(addresses, fmt.Sprintf("%s/%d", network, address.Prefix))
	}
	if len(addresses) > 0 {
		d.Set("addresses", addresses)
	}

	// TODO: get any other parameters from the network and save them

	log.Printf("[DEBUG] Network ID %s successfully read", d.Id())
	return nil
}

func resourceLibvirtNetworkDelete(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}
	log.Printf("[DEBUG] Deleting network ID %s", d.Id())

	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	if err != nil {
		return fmt.Errorf("When destroying libvirt network: error retrieving %s", err)
	}
	defer network.Free()

	active, err := network.IsActive()
	if err != nil {
		return fmt.Errorf("Couldn't determine if network is active: %s", err)
	}
	if !active {
		// we have to restart an inactive network, otherwise it won't be
		// possible to remove it.
		if err := network.Create(); err != nil {
			return fmt.Errorf("Cannot restart an inactive network %s", err)
		}
	}

	if err := network.Destroy(); err != nil {
		return fmt.Errorf("When destroying libvirt network: %s", err)
	}

	if err := network.Undefine(); err != nil {
		return fmt.Errorf("Couldn't undefine libvirt network: %s", err)
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"ACTIVE"},
		Target:     []string{"NOT-EXISTS"},
		Refresh:    waitForNetworkDestroyed(virConn, d.Id()),
		Timeout:    1 * time.Minute,
		Delay:      5 * time.Second,
		MinTimeout: 3 * time.Second,
	}
	_, err = stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf("Error waiting for network to reach NOT-EXISTS state: %s", err)
	}
	return nil
}

func waitForNetworkActive(network libvirt.VirNetwork) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		active, err := network.IsActive()
		if err != nil {
			return nil, "", err
		}
		if active {
			return network, "ACTIVE", nil
		}
		return network, "BUILD", err
	}
}

// wait for network to be up and timeout after 5 minutes.
func waitForNetworkDestroyed(virConn *libvirt.VirConnection, uuid string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		log.Printf("Waiting for network %s to be destroyed", uuid)
		network, err := virConn.LookupNetworkByUUIDString(uuid)
		if err.(libvirt.VirError).Code == libvirt.VIR_ERR_NO_NETWORK {
			return virConn, "NOT-EXISTS", nil
		}
		defer network.Free()
		return virConn, "ACTIVE", err
	}
}
<commit_msg>Do not specify NAT when using routed network<commit_after>package libvirt

import (
	"fmt"
	"log"
	"net"
	"strings"
	"time"

	libvirt "github.com/dmacvicar/libvirt-go"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

const (
	netModeIsolated = "none"
	netModeNat      = "nat"
	netModeRoute    = "route"
	netModeBridge   = "bridge"
)

// a libvirt network resource
//
// Resource example:
//
// resource "libvirt_network" "k8snet" {
//    name = "k8snet"
//    domain = "k8s.local"
//    mode = "nat"
//    addresses = ["10.17.3.0/24"]
// }
//
// "addresses" can contain (0 or 1) ipv4 and (0 or 1) ipv6 ranges
// "mode" can be one of: "nat" (default), "isolated"
//
func resourceLibvirtNetwork() *schema.Resource {
	return &schema.Resource{
		Create: resourceLibvirtNetworkCreate,
		Read:   resourceLibvirtNetworkRead,
		Delete: resourceLibvirtNetworkDelete,
		Exists: resourceLibvirtNetworkExists,
		Update: resourceLibvirtNetworkUpdate,
		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"domain": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"mode": &schema.Schema{ // can be "none", "nat" (default), "route", "bridge"
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  netModeNat,
			},
			"bridge": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},
			"addresses": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Required: false,
				ForceNew: true,
				Elem: &schema.Schema{
					Type: schema.TypeString,
				},
			},
			"running": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
				ForceNew: false,
			},
		},
	}
}

func resourceLibvirtNetworkExists(d *schema.ResourceData, meta interface{}) (bool, error) {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return false, fmt.Errorf("The libvirt connection was nil.")
	}
	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	defer network.Free()
	return err == nil, err
}

func resourceLibvirtNetworkUpdate(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}
	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	defer network.Free()

	active, err := network.IsActive()
	if err != nil {
		return err
	}

	if !active {
		log.Printf("[DEBUG] Activating network")
		if err := network.Create(); err != nil {
			return err
		}
		d.Set("running", true)
	}

	return nil
}

func resourceLibvirtNetworkCreate(d *schema.ResourceData, meta interface{}) error {
	// see https://libvirt.org/formatnetwork.html
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	networkDef := newNetworkDef()
	networkDef.Name = d.Get("name").(string)
	networkDef.Domain = &defNetworkDomain{
		Name: d.Get("domain").(string),
	}

	// use a bridge provided by the user, or create one otherwise (libvirt will assign on automatically when empty)
	bridgeName := ""
	if b, ok := d.GetOk("bridge"); ok {
		bridgeName = b.(string)
	}
	networkDef.Bridge = &defNetworkBridge{
		Name: bridgeName,
		Stp:  "on",
	}

	// check the network mode
	networkDef.Forward.Mode = strings.ToLower(d.Get("mode").(string))
	if networkDef.Forward.Mode == netModeIsolated || networkDef.Forward.Mode == netModeNat || networkDef.Forward.Mode == netModeRoute {

		if networkDef.Forward.Mode == netModeIsolated {
			// there is no forwarding when using an isolated network
			networkDef.Forward = nil
		} else if networkDef.Forward.Mode == netModeRoute {
			// there is no NAT when using a routed network
			networkDef.Forward.Nat = nil
		}

		// some network modes require a DHCP/DNS server
		// set the addresses for DHCP
		if addresses, ok := d.GetOk("addresses"); ok {
			ipsPtrsLst := []*defNetworkIp{}
			for _, addressI := range addresses.([]interface{}) {
				address := addressI.(string)
				_, ipNet, err := net.ParseCIDR(address)
				if err != nil {
					return fmt.Errorf("Error parsing addresses definition '%s': %s", address, err)
				}
				ones, bits := ipNet.Mask.Size()
				family := "ipv4"
				if bits == (net.IPv6len * 8) {
					family = "ipv6"
				}
				ipsRange := 2 ^ bits - 2 ^ ones
				if ipsRange < 4 {
					return fmt.Errorf("Netmask seems to be too strict: only %d IPs available (%s)", ipsRange-3, family)
				}

				// we should calculate the range served by DHCP. For example, for
				// 192.168.121.0/24 we will serve 192.168.121.2 - 192.168.121.254
				start, end := NetworkRange(ipNet)

				// skip the .0, (for the network),
				start[len(start)-1]++

				// assign the .1 to the host interface
				dni := defNetworkIp{
					Address: start.String(),
					Prefix:  ones,
					Family:  family,
				}

				start[len(start)-1]++ // then skip the .1
				end[len(end)-1]--     // and skip the .255 (for broadcast)

				dni.Dhcp = &defNetworkIpDhcp{
					Ranges: []*defNetworkIpDhcpRange{
						&defNetworkIpDhcpRange{
							Start: start.String(),
							End:   end.String(),
						},
					},
				}
				ipsPtrsLst = append(ipsPtrsLst, &dni)
			}
			networkDef.Ips = ipsPtrsLst
		}
	} else if networkDef.Forward.Mode == netModeBridge {
		if bridgeName == "" {
			return fmt.Errorf("'bridge' must be provided when using the bridged network mode")
		}
	} else {
		return fmt.Errorf("unsupported network mode '%s'", networkDef.Forward.Mode)
	}

	// once we have the network defined, connect to libvirt and create it from the XML serialization
	connectURI, err := virConn.GetURI()
	if err != nil {
		return fmt.Errorf("Error retrieving libvirt connection URI: %s", err)
	}
	log.Printf("[INFO] Creating libvirt network at %s", connectURI)

	data, err := xmlMarshallIndented(networkDef)
	if err != nil {
		return fmt.Errorf("Error serializing libvirt network: %s", err)
	}

	log.Printf("[DEBUG] Creating libvirt network at %s: %s", connectURI, data)
	network, err := virConn.NetworkDefineXML(data)
	if err != nil {
		return fmt.Errorf("Error defining libvirt network: %s - %s", err, data)
	}
	err = network.Create()
	if err != nil {
		return fmt.Errorf("Error crearing libvirt network: %s", err)
	}
	defer network.Free()

	id, err := network.GetUUIDString()
	if err != nil {
		return fmt.Errorf("Error retrieving libvirt network id: %s", err)
	}
	d.SetId(id)

	log.Printf("[INFO] Created network %s [%s]", networkDef.Name, d.Id())

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"BUILD"},
		Target:     []string{"ACTIVE"},
		Refresh:    waitForNetworkActive(network),
		Timeout:    1 * time.Minute,
		Delay:      5 * time.Second,
		MinTimeout: 3 * time.Second,
	}
	_, err = stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf("Error waiting for network to reach ACTIVE state: %s", err)
	}

	return resourceLibvirtNetworkRead(d, meta)
}

func resourceLibvirtNetworkRead(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}

	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	if err != nil {
		return fmt.Errorf("Error retrieving libvirt network: %s", err)
	}
	defer network.Free()

	networkDef, err := newDefNetworkfromLibvirt(&network)
	if err != nil {
		return fmt.Errorf("Error reading libvirt network XML description: %s", err)
	}

	d.Set("name", networkDef.Name)
	d.Set("domain", networkDef.Domain.Name)
	d.Set("bridge", networkDef.Bridge.Name)

	active, err := network.IsActive()
	if err != nil {
		return err
	}
	d.Set("running", active)

	addresses := []string{}
	for _, address := range networkDef.Ips {
		// we get the host interface IP (ie, 10.10.8.1) but we want the network CIDR (ie, 10.10.8.0/24)
		// so we need some transformations...
		addr := net.ParseIP(address.Address)
		if addr == nil {
			return fmt.Errorf("Error parsing IP '%s': %s", address, err)
		}
		bits := net.IPv6len * 8
		if addr.To4() != nil {
			bits = net.IPv4len * 8
		}
		mask := net.CIDRMask(address.Prefix, bits)
		network := addr.Mask(mask)
		addresses = append(addresses, fmt.Sprintf("%s/%d", network, address.Prefix))
	}
	if len(addresses) > 0 {
		d.Set("addresses", addresses)
	}

	// TODO: get any other parameters from the network and save them

	log.Printf("[DEBUG] Network ID %s successfully read", d.Id())
	return nil
}

func resourceLibvirtNetworkDelete(d *schema.ResourceData, meta interface{}) error {
	virConn := meta.(*Client).libvirt
	if virConn == nil {
		return fmt.Errorf("The libvirt connection was nil.")
	}
	log.Printf("[DEBUG] Deleting network ID %s", d.Id())

	network, err := virConn.LookupNetworkByUUIDString(d.Id())
	if err != nil {
		return fmt.Errorf("When destroying libvirt network: error retrieving %s", err)
	}
	defer network.Free()

	active, err := network.IsActive()
	if err != nil {
		return fmt.Errorf("Couldn't determine if network is active: %s", err)
	}
	if !active {
		// we have to restart an inactive network, otherwise it won't be
		// possible to remove it.
		if err := network.Create(); err != nil {
			return fmt.Errorf("Cannot restart an inactive network %s", err)
		}
	}

	if err := network.Destroy(); err != nil {
		return fmt.Errorf("When destroying libvirt network: %s", err)
	}

	if err := network.Undefine(); err != nil {
		return fmt.Errorf("Couldn't undefine libvirt network: %s", err)
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"ACTIVE"},
		Target:     []string{"NOT-EXISTS"},
		Refresh:    waitForNetworkDestroyed(virConn, d.Id()),
		Timeout:    1 * time.Minute,
		Delay:      5 * time.Second,
		MinTimeout: 3 * time.Second,
	}
	_, err = stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf("Error waiting for network to reach NOT-EXISTS state: %s", err)
	}
	return nil
}

func waitForNetworkActive(network libvirt.VirNetwork) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		active, err := network.IsActive()
		if err != nil {
			return nil, "", err
		}
		if active {
			return network, "ACTIVE", nil
		}
		return network, "BUILD", err
	}
}

// wait for network to be up and timeout after 5 minutes.
func waitForNetworkDestroyed(virConn *libvirt.VirConnection, uuid string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		log.Printf("Waiting for network %s to be destroyed", uuid)
		network, err := virConn.LookupNetworkByUUIDString(uuid)
		if err.(libvirt.VirError).Code == libvirt.VIR_ERR_NO_NETWORK {
			return virConn, "NOT-EXISTS", nil
		}
		defer network.Free()
		return virConn, "ACTIVE", err
	}
}
<|endoftext|>


<commit_before>/*
@author: xuchengxuan(bigpyer@126.com)
@brief: logger module
*/
package alog

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"
)

const (
	DATEFORMAT        = "2006-01-02"
	DEFAULT_LOG_SCAN  = 300
	DEFAULT_LOG_LEVEL = ERROR
)

type LEVEL byte

const (
	INFO LEVEL = iota
	DEBUG
	WARN
	ERROR
)

type logger struct {
	mu       *sync.RWMutex
	fileDir  string
	fileName string

	date *time.Time

	logFile  *os.File
	lger     *log.Logger
	timeScan int64

	logChan  chan string
	logLevel LEVEL
}

// logger handler constructor
func NewLogger(dir string, name string) *logger {
	dailyLogger := &logger{
		mu:       new(sync.RWMutex),
		fileDir:  dir,
		fileName: name,
		logChan:  make(chan string, 500),
		logLevel: DEFAULT_LOG_LEVEL,
	}

	dailyLogger.initDailyLogger()

	return dailyLogger
}

/*
* @param level(0,1,2,3)
* @desc  if output level is larger than log level,
*       the content will be outputed
 */
func (f *logger) SetLogLevel(level LEVEL) {
	f.logLevel = level
}

func (f *logger) initDailyLogger() {

	t, _ := time.Parse(DATEFORMAT, time.Now().Format(DATEFORMAT))

	f.date = &t
	f.mu.Lock()
	defer f.mu.Unlock()

	logFile := joinFilePath(f.fileDir, f.fileName)
	f.logFile, _ = os.OpenFile(logFile, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
	f.lger = log.New(f.logFile, "", log.LstdFlags|log.Lmicroseconds)

	go f.writeLog()
	go f.monitorFile()
}

func (f *logger) isNeedRotate() bool {
	t, _ := time.Parse(DATEFORMAT, time.Now().Format(DATEFORMAT))
	if t.After(*f.date) {
		return true
	}
	return false
}

func (f *logger) rotate() {
	logFile := joinFilePath(f.fileDir, f.fileName)
	logFileBak := logFile + "." + f.date.Format(DATEFORMAT)

	if !isExist(logFileBak) && f.isNeedRotate() {
		if f.logFile != nil {
			f.logFile.Close()
		}

		err := os.Rename(logFile, logFileBak)
		if err != nil {
			f.lger.Printf("logger rename error: %v", err.Error())
		}

		t, _ := time.Parse(DATEFORMAT, time.Now().Format(DATEFORMAT))
		f.date = &t
		f.logFile, _ = os.Create(logFile)
		f.lger = log.New(f.logFile, "", log.LstdFlags|log.Lmicroseconds)
	}
}

func (f *logger) monitorFile() {
	defer func() {
		if err := recover(); err != nil {
			f.lger.Panic("logger's FileMonitor() catch panic: %v\n", err)
		}
	}()

	// check frequency
	logScan := DEFAULT_LOG_SCAN

	timer := time.NewTicker(time.Duration(logScan) * time.Second)
	for {
		select {
		case <-timer.C:
			f.checkFile()
		}
	}
}

func (f *logger) checkFile() {
	defer func() {
		if err := recover(); err != nil {
			f.lger.Printf("logger's FileCheck() catch panic: %v\n", err)
		}
	}()
	if f.isNeedRotate() {
		f.mu.Lock()
		defer f.mu.Unlock()

		f.rotate()
	}
}

// passive to close filelogger
func (f *logger) Close() error {

	close(f.logChan)
	f.lger = nil

	return f.logFile.Close()
}

// Receive logStr from f's logChan and print logstr to file
func (f *logger) writeLog() {
	defer func() {
		if err := recover(); err != nil {
			log.Printf(" writeLog catch panic: %v\n", err)
		}
	}()

	for {
		select {
		case str := <-f.logChan:
			f.outPut(str)
		}
	}
}

// print log
func (f *logger) outPut(str string) {
	f.mu.RLock()
	defer f.mu.RUnlock()

	f.lger.Output(2, str)
}

func isExist(path string) bool {
	_, err := os.Stat(path)
	return err == nil || os.IsExist(err)
}

func joinFilePath(path, file string) string {
	return filepath.Join(path, file)
}

func shortFileName(file string) string {
	return filepath.Base(file)
}

// info log
func (f *logger) Info(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= INFO {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[INFO] "+format, v...)
	}
}

// debug log
func (f *logger) Debug(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= DEBUG {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[DEBUG] "+format, v...)
	}
}

// warn log
func (f *logger) Warn(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= WARN {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[WARN] "+format, v...)
	}
}

// error log
func (f *logger) Error(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= ERROR {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[ERROR] "+format, v...)
	}
}
<commit_msg>按照日期生成日志文件<commit_after>/*
@author: xuchengxuan(bigpyer@126.com)
@brief: logger module
*/
package alog

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"
)

const (
	DATEFORMAT        = "2006-01-02"
	DEFAULT_LOG_SCAN  = 300
	DEFAULT_LOG_LEVEL = ERROR
)

type LEVEL byte

const (
	INFO LEVEL = iota
	DEBUG
	WARN
	ERROR
)

type logger struct {
	mu       *sync.RWMutex
	fileDir  string
	fileName string

	date *time.Time

	logFile  *os.File
	lger     *log.Logger
	timeScan int64

	logChan  chan string
	logLevel LEVEL
}

// logger handler constructor
func NewLogger(dir string, name string) *logger {
	dailyLogger := &logger{
		mu:       new(sync.RWMutex),
		fileDir:  dir,
		fileName: name,
		logChan:  make(chan string, 500),
		logLevel: DEFAULT_LOG_LEVEL,
	}

	dailyLogger.initDailyLogger()

	return dailyLogger
}

/*
* @param level(0,1,2,3)
* @desc  if output level is larger than log level,
*       the content will be outputed
 */
func (f *logger) SetLogLevel(level LEVEL) {
	f.logLevel = level
}

func (f *logger) initDailyLogger() {

	t, _ := time.Parse(DATEFORMAT, time.Now().Format(DATEFORMAT))

	f.date = &t
	f.fileName = f.fileName + "." + f.date.Format(DATEFORMAT)
	f.mu.Lock()
	defer f.mu.Unlock()

	logFile := joinFilePath(f.fileDir, f.fileName)
	f.logFile, _ = os.OpenFile(logFile, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
	f.lger = log.New(f.logFile, "", log.LstdFlags|log.Lmicroseconds)

	go f.writeLog()
	go f.monitorFile()
}

func (f *logger) isNeedRotate() bool {
	t, _ := time.Parse(DATEFORMAT, time.Now().Format(DATEFORMAT))
	if t.After(*f.date) {
		return true
	}
	return false
}

// rotate file by date
func (f *logger) rotate() {
	t, _ := time.Parse(DATEFORMAT, time.Now().Format(DATEFORMAT))
	f.date = &t

	f.fileName = f.fileName + "." + f.date.Format(DATEFORMAT)
	logFile := joinFilePath(f.fileDir, f.fileName)

	f.logFile.Close()
	f.logFile, _ = os.Create(logFile)
	f.lger = log.New(f.logFile, "", log.LstdFlags|log.Lmicroseconds)
}

func (f *logger) monitorFile() {
	defer func() {
		if err := recover(); err != nil {
			f.lger.Panic("logger's FileMonitor() catch panic: %v\n", err)
		}
	}()

	// check frequency
	logScan := DEFAULT_LOG_SCAN

	timer := time.NewTicker(time.Duration(logScan) * time.Second)
	for {
		select {
		case <-timer.C:
			f.checkFile()
		}
	}
}

func (f *logger) checkFile() {
	defer func() {
		if err := recover(); err != nil {
			f.lger.Printf("logger's FileCheck() catch panic: %v\n", err)
		}
	}()
	if f.isNeedRotate() {
		f.mu.Lock()
		defer f.mu.Unlock()

		f.rotate()
	}
}

// passive to close filelogger
func (f *logger) Close() error {

	close(f.logChan)
	f.lger = nil

	return f.logFile.Close()
}

// Receive logStr from f's logChan and print logstr to file
func (f *logger) writeLog() {
	defer func() {
		if err := recover(); err != nil {
			log.Printf(" writeLog catch panic: %v\n", err)
		}
	}()

	for {
		select {
		case str := <-f.logChan:
			f.outPut(str)
		}
	}
}

// print log
func (f *logger) outPut(str string) {
	f.mu.RLock()
	defer f.mu.RUnlock()

	f.lger.Output(2, str)
}

func joinFilePath(path, file string) string {
	return filepath.Join(path, file)
}

func shortFileName(file string) string {
	return filepath.Base(file)
}

// info log
func (f *logger) Info(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= INFO {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[INFO] "+format, v...)
	}
}

// debug log
func (f *logger) Debug(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= DEBUG {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[DEBUG] "+format, v...)
	}
}

// warn log
func (f *logger) Warn(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= WARN {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[WARN] "+format, v...)
	}
}

// error log
func (f *logger) Error(format string, v ...interface{}) {
	_, file, line, _ := runtime.Caller(1) //calldepth=1
	if f.logLevel <= ERROR {
		f.logChan <- fmt.Sprintf("[%v:%v]", shortFileName(file), line) + fmt.Sprintf("[ERROR] "+format, v...)
	}
}
<|endoftext|>


<commit_before>package ansi

import (
	"bytes"
	"fmt"
	"strings"
	"sync"
)

const (
	black = iota
	red
	green
	yellow
	blue
	magenta
	cyan
	white

	normalIntensityFG = 30
	highIntensityFG   = 90
	normalIntensityBG = 40
	highIntensityBG   = 100

	start     = "\033["
	bold      = "1;"
	blink     = "5;"
	underline = "4;"
	inverse   = "7;"

	// Reset is the ANSI reset escape sequence
	Reset = "\033[0m"
)

var (
	plain = false
	// Colors maps common color names to their ANSI color code.
	Colors = map[string]int{
		"black":   black,
		"red":     red,
		"green":   green,
		"yellow":  yellow,
		"blue":    blue,
		"magenta": magenta,
		"cyan":    cyan,
		"white":   white,
	}
)

// ColorCode returns the ANSI color color code for style.
func ColorCode(style string) string {
	return colorCode(style).String()
}

// Gets the ANSI color code for a style.
func colorCode(style string) *bytes.Buffer {
	buf := bytes.NewBufferString("")
	if plain || style == "" {
		return buf
	}
	if style == "reset" {
		buf.WriteString(Reset)
		return buf
	}

	foregroundBackground := strings.Split(style, ":")
	foreground := strings.Split(foregroundBackground[0], "+")
	fg := Colors[foreground[0]]
	fgStyle := ""
	if len(foreground) > 1 {
		fgStyle = foreground[1]
	}

	bg, bgStyle := "", ""

	if len(foregroundBackground) > 1 {
		background := strings.Split(foregroundBackground[1], "+")
		bg = background[0]
		if len(background) > 1 {
			bgStyle = background[1]
		}
	}

	buf.WriteString(start)
	base := normalIntensityFG
	if len(fgStyle) > 0 {
		if strings.Contains(fgStyle, "b") {
			buf.WriteString(bold)
		}
		if strings.Contains(fgStyle, "B") {
			buf.WriteString(blink)
		}
		if strings.Contains(fgStyle, "u") {
			buf.WriteString(underline)
		}
		if strings.Contains(fgStyle, "i") {
			buf.WriteString(inverse)
		}
		if strings.Contains(fgStyle, "h") {
			base = highIntensityFG
		}
	}
	fmt.Fprintf(buf, "%d;", base+fg)

	base = normalIntensityBG
	if len(bg) > 0 {
		if strings.Contains(bgStyle, "h") {
			base = highIntensityBG
		}
		fmt.Fprintf(buf, "%d;", base+Colors[bg])
	}

	// remove last ";"
	buf.Truncate(buf.Len() - 1)
	buf.WriteRune('m')
	return buf
}

// Color colors a string based on the ANSI color code for style.
func Color(s, style string) string {
	if plain || len(style) < 1 {
		return s
	}
	buf := colorCode(style)
	buf.WriteString(s)
	buf.WriteString(Reset)
	return buf.String()
}

// ColorFunc creates a closureto avoid ANSI color code calculation.
func ColorFunc(style string) func(string) string {
	if style == "" {
		return func(s string) string {
			return s
		}
	}
	buf := colorCode(style)
	lenBuf := buf.Len()
	var mutex sync.Mutex

	return func(s string) string {
		if plain || s == "" {
			return s
		}
		mutex.Lock()
		buf.WriteString(s)
		buf.WriteString(Reset)
		result := buf.String()
		buf.Truncate(lenBuf)
		mutex.Unlock()
		return result
	}
}

// DisableColors disables ANSI color codes. On by default.
func DisableColors(disable bool) {
	plain = disable
}
<commit_msg>simplify<commit_after>package ansi

import (
	"bytes"
	"fmt"
	"strings"
)

const (
	black = iota
	red
	green
	yellow
	blue
	magenta
	cyan
	white

	normalIntensityFG = 30
	highIntensityFG   = 90
	normalIntensityBG = 40
	highIntensityBG   = 100

	start     = "\033["
	bold      = "1;"
	blink     = "5;"
	underline = "4;"
	inverse   = "7;"

	// Reset is the ANSI reset escape sequence
	Reset = "\033[0m"
)

var (
	plain = false
	// Colors maps common color names to their ANSI color code.
	Colors = map[string]int{
		"black":   black,
		"red":     red,
		"green":   green,
		"yellow":  yellow,
		"blue":    blue,
		"magenta": magenta,
		"cyan":    cyan,
		"white":   white,
	}
)

// ColorCode returns the ANSI color color code for style.
func ColorCode(style string) string {
	return colorCode(style).String()
}

// Gets the ANSI color code for a style.
func colorCode(style string) *bytes.Buffer {
	buf := bytes.NewBufferString("")
	if plain || style == "" {
		return buf
	}
	if style == "reset" {
		buf.WriteString(Reset)
		return buf
	}

	foregroundBackground := strings.Split(style, ":")
	foreground := strings.Split(foregroundBackground[0], "+")
	fg := Colors[foreground[0]]
	fgStyle := ""
	if len(foreground) > 1 {
		fgStyle = foreground[1]
	}

	bg, bgStyle := "", ""

	if len(foregroundBackground) > 1 {
		background := strings.Split(foregroundBackground[1], "+")
		bg = background[0]
		if len(background) > 1 {
			bgStyle = background[1]
		}
	}

	buf.WriteString(start)
	base := normalIntensityFG
	if len(fgStyle) > 0 {
		if strings.Contains(fgStyle, "b") {
			buf.WriteString(bold)
		}
		if strings.Contains(fgStyle, "B") {
			buf.WriteString(blink)
		}
		if strings.Contains(fgStyle, "u") {
			buf.WriteString(underline)
		}
		if strings.Contains(fgStyle, "i") {
			buf.WriteString(inverse)
		}
		if strings.Contains(fgStyle, "h") {
			base = highIntensityFG
		}
	}
	fmt.Fprintf(buf, "%d;", base+fg)

	base = normalIntensityBG
	if len(bg) > 0 {
		if strings.Contains(bgStyle, "h") {
			base = highIntensityBG
		}
		fmt.Fprintf(buf, "%d;", base+Colors[bg])
	}

	// remove last ";"
	buf.Truncate(buf.Len() - 1)
	buf.WriteRune('m')
	return buf
}

// Color colors a string based on the ANSI color code for style.
func Color(s, style string) string {
	if plain || len(style) < 1 {
		return s
	}
	buf := colorCode(style)
	buf.WriteString(s)
	buf.WriteString(Reset)
	return buf.String()
}

// ColorFunc creates a closureto avoid ANSI color code calculation.
func ColorFunc(style string) func(string) string {
	if style == "" {
		return func(s string) string {
			return s
		}
	}
	color := ColorCode(style)
	return func(s string) string {
		if plain || s == "" {
			return s
		}
		buf := bytes.NewBufferString(color)
		buf.WriteString(s)
		buf.WriteString(Reset)
		result := buf.String()
		return result
	}
}

// DisableColors disables ANSI color codes. On by default.
func DisableColors(disable bool) {
	plain = disable
}
<|endoftext|>