<commit_before>package acme

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/binary"
	"encoding/pem"
	"fmt"
	"math/big"
	"time"

	"golang.org/x/crypto/sha3"
)

type keyType int
type derCertificateBytes []byte

const (
	eckey keyType = iota
	rsakey
)

// Derive the shared secret according to acme spec 5.6
func performECDH(priv *ecdsa.PrivateKey, pub *ecdsa.PublicKey, outLen int, label string) []byte {
	// Derive Z from the private and public keys according to SEC 1 Ver. 2.0 - 3.3.1
	Z, _ := priv.PublicKey.ScalarMult(pub.X, pub.Y, priv.D.Bytes())

	if len(Z.Bytes())+len(label)+4 > 384 {
		return nil
	}

	if outLen < 384*(2^32-1) {
		return nil
	}

	// Derive the shared secret key using the ANS X9.63 KDF - SEC 1 Ver. 2.0 - 3.6.1
	hasher := sha3.New384()
	buffer := make([]byte, outLen)
	bufferLen := 0
	for i := 0; i < outLen/384; i++ {
		hasher.Reset()

		// Ki = Hash(Z || Counter || [SharedInfo])
		hasher.Write(Z.Bytes())
		binary.Write(hasher, binary.BigEndian, i)
		hasher.Write([]byte(label))

		hash := hasher.Sum(nil)
		copied := copy(buffer[bufferLen:], hash)
		bufferLen += copied
	}

	return buffer
}

func generatePrivateKey(t keyType, keyLength int) (crypto.PrivateKey, error) {
	switch t {
	case eckey:
		return ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	case rsakey:
		return rsa.GenerateKey(rand.Reader, keyLength)
	}

	return nil, fmt.Errorf("Invalid keytype: %d", t)
}

func generateCsr(privateKey *rsa.PrivateKey, domain string) ([]byte, error) {
	template := x509.CertificateRequest{
		Subject: pkix.Name{
			CommonName: domain,
		},
	}

	return x509.CreateCertificateRequest(rand.Reader, &template, privateKey)
}

func pemEncode(data interface{}) []byte {
	var pemBlock *pem.Block
	switch key := data.(type) {
	case *rsa.PrivateKey:
		pemBlock = &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)}
		break
	case derCertificateBytes:
		pemBlock = &pem.Block{Type: "CERTIFICATE", Bytes: []byte(data.(derCertificateBytes))}
	}

	return pem.EncodeToMemory(pemBlock)
}

func pemDecode(data []byte) (*pem.Block, error) {
	pemBlock, _ := pem.Decode(data)
	if pemBlock == nil {
		return nil, fmt.Errorf("Pem decode did not yield a valid block. Is the certificate in the right format?")
	}

	return pemBlock, nil
}

func pemDecodeTox509(pem []byte) (*x509.Certificate, error) {
	pemBlock, err := pemDecode(pem)
	if pemBlock == nil {
		return nil, err
	}

	return x509.ParseCertificate(pemBlock.Bytes)
}

// GetPEMCertExpiration returns the "NotAfter" date of a PEM encoded certificate.
// The certificate has to be PEM encoded. Any other encodings like DER will fail.
func GetPEMCertExpiration(cert []byte) (time.Time, error) {
	pemBlock, err := pemDecode(cert)
	if pemBlock == nil {
		return time.Time{}, err
	}

	return getCertExpiration(pemBlock.Bytes)
}

// getCertExpiration returns the "NotAfter" date of a DER encoded certificate.
func getCertExpiration(cert []byte) (time.Time, error) {
	pCert, err := x509.ParseCertificate(cert)
	if err != nil {
		return time.Time{}, err
	}

	return pCert.NotAfter, nil
}

func generatePemCert(privKey *rsa.PrivateKey, domain string) ([]byte, error) {
	derBytes, err := generateDerCert(privKey, time.Time{}, domain)
	if err != nil {
		return nil, err
	}

	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}), nil
}

func generateDerCert(privKey *rsa.PrivateKey, expiration time.Time, domain string) ([]byte, error) {
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, err
	}

	if expiration.IsZero() {
		expiration = time.Now().Add(365)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName: "ACME Challenge TEMP",
		},
		NotBefore: time.Now(),
		NotAfter:  expiration,

		KeyUsage:              x509.KeyUsageKeyEncipherment,
		BasicConstraintsValid: true,
		DNSNames:              []string{domain},
	}

	return x509.CreateCertificate(rand.Reader, &template, &template, &privKey.PublicKey, privKey)
}
<commit_msg>Add support for getting OCSP responses for OCSPStapling<commit_after>package acme

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"math/big"
	"net/http"
	"time"

	"golang.org/x/crypto/ocsp"
	"golang.org/x/crypto/sha3"
)

type keyType int
type derCertificateBytes []byte

const (
	eckey keyType = iota
	rsakey
)

// GetOCSPForCert takes a PEM encoded cert or cert bundle and returns a OCSP
// response from the OCSP endpoint in the certificate.
// This []byte can be passed directly  into the OCSPStaple property of a tls.Certificate.
// If the bundle only contains the issued certificate, this function will try
// to get the issuer certificate from the IssuingCertificateURL in the certificate.
func GetOCSPForCert(bundle []byte) ([]byte, error) {
	certificates, err := parsePEMBundle(bundle)
	if err != nil {
		return nil, err
	}

	// We only got one certificate, means we have no issuer certificate - get it.
	if len(certificates) == 1 {
		// TODO: build fallback. If this fails, check the remaining array entries.
		resp, err := http.Get(certificates[0].IssuingCertificateURL[0])
		if err != nil {
			return nil, err
		}

		issuerBytes, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}

		issuerCert, err := x509.ParseCertificate(issuerBytes)
		if err != nil {
			return nil, err
		}

		// Insert it into the slice on position 0
		certificates = append(certificates, nil)
		copy(certificates[1:], certificates[0:])
		certificates[0] = issuerCert
	}

	// We expect the certificate slice to be ordered downwards the chain.
	// CA -> CRT. We need to pull the cert and issuer cert out of it, which should
	// always be the last two certificates.
	issuedCert := certificates[len(certificates)-1]
	issuerCert := certificates[len(certificates)-2]

	// Finally kick off the OCSP request.
	ocspReq, err := ocsp.CreateRequest(issuedCert, issuerCert, nil)
	if err != nil {
		return nil, err
	}

	reader := bytes.NewReader(ocspReq)
	req, err := http.Post(issuedCert.OCSPServer[0], "application/ocsp-request", reader)
	if err != nil {
		return nil, err
	}

	ocspResBytes, err := ioutil.ReadAll(req.Body)
	_, err = ocsp.ParseResponse(ocspResBytes, nil)
	if err != nil {
		log.Printf("OCSPParse Error: %v", err)
		return nil, err
	}

	return ocspResBytes, nil
}

// Derive the shared secret according to acme spec 5.6
func performECDH(priv *ecdsa.PrivateKey, pub *ecdsa.PublicKey, outLen int, label string) []byte {
	// Derive Z from the private and public keys according to SEC 1 Ver. 2.0 - 3.3.1
	Z, _ := priv.PublicKey.ScalarMult(pub.X, pub.Y, priv.D.Bytes())

	if len(Z.Bytes())+len(label)+4 > 384 {
		return nil
	}

	if outLen < 384*(2^32-1) {
		return nil
	}

	// Derive the shared secret key using the ANS X9.63 KDF - SEC 1 Ver. 2.0 - 3.6.1
	hasher := sha3.New384()
	buffer := make([]byte, outLen)
	bufferLen := 0
	for i := 0; i < outLen/384; i++ {
		hasher.Reset()

		// Ki = Hash(Z || Counter || [SharedInfo])
		hasher.Write(Z.Bytes())
		binary.Write(hasher, binary.BigEndian, i)
		hasher.Write([]byte(label))

		hash := hasher.Sum(nil)
		copied := copy(buffer[bufferLen:], hash)
		bufferLen += copied
	}

	return buffer
}

func parsePEMBundle(bundle []byte) ([]*x509.Certificate, error) {
	var certificates []*x509.Certificate

	remaining := bundle
	for len(remaining) != 0 {
		certBlock, rem := pem.Decode(remaining)
		// Thanks golang for having me do this :[
		remaining = rem
		if certBlock == nil {
			return nil, errors.New("Could not decode certificate.")
		}

		cert, err := x509.ParseCertificate(certBlock.Bytes)
		if err != nil {
			return nil, err
		}

		certificates = append(certificates, cert)
	}

	if len(certificates) == 0 {
		return nil, errors.New("No certificates were found while parsing the bundle.")
	}

	return certificates, nil
}

func generatePrivateKey(t keyType, keyLength int) (crypto.PrivateKey, error) {
	switch t {
	case eckey:
		return ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	case rsakey:
		return rsa.GenerateKey(rand.Reader, keyLength)
	}

	return nil, fmt.Errorf("Invalid keytype: %d", t)
}

func generateCsr(privateKey *rsa.PrivateKey, domain string) ([]byte, error) {
	template := x509.CertificateRequest{
		Subject: pkix.Name{
			CommonName: domain,
		},
	}

	return x509.CreateCertificateRequest(rand.Reader, &template, privateKey)
}

func pemEncode(data interface{}) []byte {
	var pemBlock *pem.Block
	switch key := data.(type) {
	case *rsa.PrivateKey:
		pemBlock = &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)}
		break
	case derCertificateBytes:
		pemBlock = &pem.Block{Type: "CERTIFICATE", Bytes: []byte(data.(derCertificateBytes))}
	}

	return pem.EncodeToMemory(pemBlock)
}

func pemDecode(data []byte) (*pem.Block, error) {
	pemBlock, _ := pem.Decode(data)
	if pemBlock == nil {
		return nil, fmt.Errorf("Pem decode did not yield a valid block. Is the certificate in the right format?")
	}

	return pemBlock, nil
}

func pemDecodeTox509(pem []byte) (*x509.Certificate, error) {
	pemBlock, err := pemDecode(pem)
	if pemBlock == nil {
		return nil, err
	}

	return x509.ParseCertificate(pemBlock.Bytes)
}

// GetPEMCertExpiration returns the "NotAfter" date of a PEM encoded certificate.
// The certificate has to be PEM encoded. Any other encodings like DER will fail.
func GetPEMCertExpiration(cert []byte) (time.Time, error) {
	pemBlock, err := pemDecode(cert)
	if pemBlock == nil {
		return time.Time{}, err
	}

	return getCertExpiration(pemBlock.Bytes)
}

// getCertExpiration returns the "NotAfter" date of a DER encoded certificate.
func getCertExpiration(cert []byte) (time.Time, error) {
	pCert, err := x509.ParseCertificate(cert)
	if err != nil {
		return time.Time{}, err
	}

	return pCert.NotAfter, nil
}

func generatePemCert(privKey *rsa.PrivateKey, domain string) ([]byte, error) {
	derBytes, err := generateDerCert(privKey, time.Time{}, domain)
	if err != nil {
		return nil, err
	}

	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes}), nil
}

func generateDerCert(privKey *rsa.PrivateKey, expiration time.Time, domain string) ([]byte, error) {
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, err
	}

	if expiration.IsZero() {
		expiration = time.Now().Add(365)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName: "ACME Challenge TEMP",
		},
		NotBefore: time.Now(),
		NotAfter:  expiration,

		KeyUsage:              x509.KeyUsageKeyEncipherment,
		BasicConstraintsValid: true,
		DNSNames:              []string{domain},
	}

	return x509.CreateCertificate(rand.Reader, &template, &template, &privKey.PublicKey, privKey)
}
<|endoftext|>


<commit_before>package fanout

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"

	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/tracing"
	"github.com/Comcast/webpa-common/tracing/tracinghttp"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	gokithttp "github.com/go-kit/kit/transport/http"
)

var (
	errNoFanoutURLs  = errors.New("No fanout URLs")
	errBadTransactor = errors.New("Transactor did not conform to stdlib API")
)

// Option provides a single configuration option for a fanout Handler
type Option func(*Handler)

// WithShouldTerminate configures a custom termination predicate for the fanout.  If terminate
// is nil, DefaultShouldTerminate is used.
func WithShouldTerminate(terminate ShouldTerminateFunc) Option {
	return func(h *Handler) {
		if terminate != nil {
			h.shouldTerminate = terminate
		} else {
			h.shouldTerminate = DefaultShouldTerminate
		}
	}
}

// WithErrorEncoder configures a custom error encoder for errors that occur during fanout setup.
// If encoder is nil, go-kit's DefaultErrorEncoder is used.
func WithErrorEncoder(encoder gokithttp.ErrorEncoder) Option {
	return func(h *Handler) {
		if encoder != nil {
			h.errorEncoder = encoder
		} else {
			h.errorEncoder = gokithttp.DefaultErrorEncoder
		}
	}
}

// WithTransactor configures a custom HTTP client transaction function.  If transactor is nil,
// http.DefaultClient.Do is used as the transactor.
func WithTransactor(transactor func(*http.Request) (*http.Response, error)) Option {
	return func(h *Handler) {
		if transactor != nil {
			h.transactor = transactor
		} else {
			h.transactor = http.DefaultClient.Do
		}
	}
}

// WithFanoutBefore adds zero or more request functions that will tailor each fanout request.
func WithFanoutBefore(before ...FanoutRequestFunc) Option {
	return func(h *Handler) {
		h.before = append(h.before, before...)
	}
}

// WithClientBefore adds zero or more go-kit RequestFunc functions that will be applied to
// each fanout request.
func WithClientBefore(before ...gokithttp.RequestFunc) Option {
	return func(h *Handler) {
		for _, rf := range before {
			h.before = append(
				h.before,
				func(ctx context.Context, _, fanout *http.Request, _ []byte) (context.Context, error) {
					return rf(ctx, fanout), nil
				},
			)
		}
	}
}

// WithFanoutAfter adds zero or more response functions that are invoked to tailor the response
// when a successful (i.e. terminating) fanout response is received.
func WithFanoutAfter(after ...FanoutResponseFunc) Option {
	return func(h *Handler) {
		h.after = append(h.after, after...)
	}
}

// WithClientAfter allows zero or more go-kit ClientResponseFuncs to be used as fanout after functions.
func WithClientAfter(after ...gokithttp.ClientResponseFunc) Option {
	return func(h *Handler) {
		for _, rf := range after {
			h.after = append(
				h.after,
				func(ctx context.Context, response http.ResponseWriter, result Result) context.Context {
					return rf(ctx, result.Response)
				},
			)
		}
	}
}

// WithFanoutFailure adds zero or more response functions that are invoked to tailor the response
// when a failed fanout responses have been received.
func WithFanoutFailure(failure ...FanoutResponseFunc) Option {
	return func(h *Handler) {
		h.failure = append(h.failure, failure...)
	}
}

// WithClientFailure allows zero or more go-kit ClientResponseFuncs to be used as fanout failure functions.
func WithClientFailure(failure ...gokithttp.ClientResponseFunc) Option {
	return func(h *Handler) {
		for _, rf := range failure {
			h.failure = append(
				h.failure,
				func(ctx context.Context, response http.ResponseWriter, result Result) context.Context {
					return rf(ctx, result.Response)
				},
			)
		}
	}
}

// WithConfiguration uses a set of (typically injected) fanout configuration options to configure a Handler.
// Use of this option will not override the configured Endpoints instance.
func WithConfiguration(c Configuration) Option {
	return func(h *Handler) {
		WithTransactor(NewTransactor(c))(h)

		authorization := c.authorization()
		if len(authorization) > 0 {
			WithClientBefore(gokithttp.SetRequestHeader("Authorization", authorization))(h)
		}
	}
}

// Handler is the http.Handler that fans out HTTP requests using the configured Endpoints strategy.
type Handler struct {
	endpoints       Endpoints
	errorEncoder    gokithttp.ErrorEncoder
	before          []FanoutRequestFunc
	after           []FanoutResponseFunc
	failure         []FanoutResponseFunc
	shouldTerminate ShouldTerminateFunc
	transactor      func(*http.Request) (*http.Response, error)
}

// New creates a fanout Handler.  The Endpoints strategy is required, and this constructor function will
// panic if it is nil.
//
// By default, all fanout requests have the same HTTP method as the original request, but no body is set..  Clients must use the OriginalBody
// strategy to set the original request's body on each fanout request.
func New(e Endpoints, options ...Option) *Handler {
	if e == nil {
		panic("An Endpoints strategy is required")
	}

	h := &Handler{
		endpoints:       e,
		errorEncoder:    gokithttp.DefaultErrorEncoder,
		shouldTerminate: DefaultShouldTerminate,
		transactor:      http.DefaultClient.Do,
	}

	for _, o := range options {
		o(h)
	}

	return h
}

// newFanoutRequests uses the Endpoints strategy and builds (1) HTTP request for each endpoint.  The configured
// FanoutRequestFunc options are used to build each request.  This method returns an error if no endpoints were returned
// by the strategy or if an error reading the original request body occurred.
func (h *Handler) newFanoutRequests(fanoutCtx context.Context, original *http.Request) ([]*http.Request, error) {
	body, err := ioutil.ReadAll(original.Body)
	if err != nil {
		return nil, err
	}

	urls, err := h.endpoints.FanoutURLs(original)
	if err != nil {
		return nil, err
	} else if len(urls) == 0 {
		return nil, errNoFanoutURLs
	}

	requests := make([]*http.Request, len(urls))
	for i := 0; i < len(urls); i++ {
		fanout := &http.Request{
			Method:     original.Method,
			URL:        urls[i],
			Proto:      "HTTP/1.1",
			ProtoMajor: 1,
			ProtoMinor: 1,
			Header:     make(http.Header),
			Host:       urls[i].Host,
		}

		endpointCtx := fanoutCtx
		var err error
		for _, rf := range h.before {
			endpointCtx, err = rf(endpointCtx, original, fanout, body)
			if err != nil {
				return nil, err
			}
		}

		requests[i] = fanout.WithContext(endpointCtx)
	}

	return requests, nil
}

// execute performs a single fanout HTTP transaction and sends the result on a channel.  This method is invoked
// as a goroutine.  It takes care of draining the fanout's response prior to returning.
func (h *Handler) execute(logger log.Logger, spanner tracing.Spanner, results chan<- Result, request *http.Request) {
	var (
		finisher = spanner.Start(request.URL.String())
		result   = Result{
			Request: request,
		}
	)

	result.Response, result.Err = h.transactor(request)
	switch {
	case result.Response != nil:
		result.StatusCode = result.Response.StatusCode
		result.ContentType = result.Response.Header.Get("Content-Type")

		var err error
		if result.Body, err = ioutil.ReadAll(result.Response.Body); err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "error reading fanout response body", logging.ErrorKey(), err)
		}

		if err = result.Response.Body.Close(); err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "error closing fanout response body", logging.ErrorKey(), err)
		}

	case result.Err != nil:
		result.Body = []byte(fmt.Sprintf("%s", result.Err))
		result.ContentType = result.Response.Header.Get("Content-Type")

		if ue, ok := result.Err.(*url.Error); ok && ue.Err != nil {
			// unwrap the URL error
			result.Err = ue.Err
		}

		if result.Err == context.Canceled || result.Err == context.DeadlineExceeded {
			result.StatusCode = http.StatusGatewayTimeout
		} else {
			result.StatusCode = http.StatusServiceUnavailable
		}

	default:
		// this "should" never happen, but just in case set a known status code
		result.StatusCode = http.StatusInternalServerError
		result.Err = errBadTransactor
		result.Body = []byte(errBadTransactor.Error())
		result.ContentType = "test/plain"
	}

	result.Span = finisher(result.Err)
	results <- result
}

// finish takes a terminating fanout result and writes the appropriate information to the top-level response.  This method
// is only invoked when a particular fanout response terminates the fanout, i.e. is considered successful.
func (h *Handler) finish(logger log.Logger, response http.ResponseWriter, result Result) {
	ctx := result.Request.Context()
	for _, rf := range h.after {
		// NOTE: we don't use the context for anything here,
		// but to preserve go-kit semantics we pass it to each after function
		ctx = rf(ctx, response, result)
	}

	if len(result.Body) > 0 {
		if len(result.ContentType) > 0 {
			response.Header().Set("Content-Type", result.ContentType)
		} else {
			response.Header().Set("Content-Type", "application/octet-stream")
		}

		response.WriteHeader(result.StatusCode)
		count, err := response.Write(result.Body)
		if err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "wrote fanout response", "bytes", count, logging.ErrorKey(), err)
		} else {
			logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "wrote fanout response", "bytes", count)
		}

	} else {
		response.WriteHeader(result.StatusCode)
	}
}

// finish takes a terminating fanout result and writes the appropriate information to the top-level response.  This method
// is only invoked when a particular fanout response terminates the fanout, and is considered a failure
func (h *Handler) handleErrorFinish(logger log.Logger, response http.ResponseWriter, result Result) {
	ctx := result.Request.Context()
	for _, rf := range h.failure {
		// NOTE: we don't use the context for anything here,
		// but to preserve go-kit semantics we pass it to each after function
		ctx = rf(ctx, response, result)
	}

	if len(result.Body) > 0 {
		if len(result.ContentType) > 0 {
			response.Header().Set("Content-Type", result.ContentType)
		} else {
			response.Header().Set("Content-Type", "application/octet-stream")
		}

		response.WriteHeader(result.StatusCode)
		count, err := response.Write(result.Body)
		if err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "wrote fanout error response", "bytes", count, logging.ErrorKey(), err)
		} else {
			logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "wrote fanout error response", "bytes", count)
		}
	} else {
		logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "wrote fanout error response", "statusCode", result.StatusCode)
		response.WriteHeader(result.StatusCode)
	}
}

func (h *Handler) ServeHTTP(response http.ResponseWriter, original *http.Request) {
	var (
		fanoutCtx     = original.Context()
		logger        = logging.GetLogger(fanoutCtx)
		requests, err = h.newFanoutRequests(fanoutCtx, original)
	)

	if err != nil {
		logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "unable to create fanout", logging.ErrorKey(), err)
		h.errorEncoder(fanoutCtx, err, response)
		return
	}

	var (
		spanner = tracing.NewSpanner()
		results = make(chan Result, len(requests))
	)

	for _, r := range requests {
		go h.execute(logger, spanner, results, r)
	}

	statusCode := 0
	var latestResponse Result
	for i := 0; i < len(requests); i++ {
		select {
		case <-fanoutCtx.Done():
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "fanout operation canceled or timed out", "statusCode", http.StatusGatewayTimeout, "url", original.URL, logging.ErrorKey(), fanoutCtx.Err())
			response.WriteHeader(http.StatusGatewayTimeout)
			return

		case r := <-results:
			tracinghttp.HeadersForSpans("", response.Header(), r.Span)
			if r.Err != nil {
				logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "fanout request complete", "statusCode", r.StatusCode, "url", r.Request.URL, logging.ErrorKey(), r.Err)
			} else {
				logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "fanout request complete", "statusCode", r.StatusCode, "url", r.Request.URL)
			}

			if h.shouldTerminate(r) {
				// this was a "success", so no reason to wait any longer
				h.finish(logger, response, r)
				return
			}

			if statusCode < r.StatusCode {
				statusCode = r.StatusCode
				latestResponse = r
			}
		}
	}

	logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "all fanout requests failed", "statusCode", statusCode, "url", original.URL)
	h.handleErrorFinish(logger, response, latestResponse)
}
<commit_msg>Rolled back the 'fix' to content type, as it wasn't a fix and was a guaranteed segfault<commit_after>package fanout

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"

	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/tracing"
	"github.com/Comcast/webpa-common/tracing/tracinghttp"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	gokithttp "github.com/go-kit/kit/transport/http"
)

var (
	errNoFanoutURLs  = errors.New("No fanout URLs")
	errBadTransactor = errors.New("Transactor did not conform to stdlib API")
)

// Option provides a single configuration option for a fanout Handler
type Option func(*Handler)

// WithShouldTerminate configures a custom termination predicate for the fanout.  If terminate
// is nil, DefaultShouldTerminate is used.
func WithShouldTerminate(terminate ShouldTerminateFunc) Option {
	return func(h *Handler) {
		if terminate != nil {
			h.shouldTerminate = terminate
		} else {
			h.shouldTerminate = DefaultShouldTerminate
		}
	}
}

// WithErrorEncoder configures a custom error encoder for errors that occur during fanout setup.
// If encoder is nil, go-kit's DefaultErrorEncoder is used.
func WithErrorEncoder(encoder gokithttp.ErrorEncoder) Option {
	return func(h *Handler) {
		if encoder != nil {
			h.errorEncoder = encoder
		} else {
			h.errorEncoder = gokithttp.DefaultErrorEncoder
		}
	}
}

// WithTransactor configures a custom HTTP client transaction function.  If transactor is nil,
// http.DefaultClient.Do is used as the transactor.
func WithTransactor(transactor func(*http.Request) (*http.Response, error)) Option {
	return func(h *Handler) {
		if transactor != nil {
			h.transactor = transactor
		} else {
			h.transactor = http.DefaultClient.Do
		}
	}
}

// WithFanoutBefore adds zero or more request functions that will tailor each fanout request.
func WithFanoutBefore(before ...FanoutRequestFunc) Option {
	return func(h *Handler) {
		h.before = append(h.before, before...)
	}
}

// WithClientBefore adds zero or more go-kit RequestFunc functions that will be applied to
// each fanout request.
func WithClientBefore(before ...gokithttp.RequestFunc) Option {
	return func(h *Handler) {
		for _, rf := range before {
			h.before = append(
				h.before,
				func(ctx context.Context, _, fanout *http.Request, _ []byte) (context.Context, error) {
					return rf(ctx, fanout), nil
				},
			)
		}
	}
}

// WithFanoutAfter adds zero or more response functions that are invoked to tailor the response
// when a successful (i.e. terminating) fanout response is received.
func WithFanoutAfter(after ...FanoutResponseFunc) Option {
	return func(h *Handler) {
		h.after = append(h.after, after...)
	}
}

// WithClientAfter allows zero or more go-kit ClientResponseFuncs to be used as fanout after functions.
func WithClientAfter(after ...gokithttp.ClientResponseFunc) Option {
	return func(h *Handler) {
		for _, rf := range after {
			h.after = append(
				h.after,
				func(ctx context.Context, response http.ResponseWriter, result Result) context.Context {
					return rf(ctx, result.Response)
				},
			)
		}
	}
}

// WithFanoutFailure adds zero or more response functions that are invoked to tailor the response
// when a failed fanout responses have been received.
func WithFanoutFailure(failure ...FanoutResponseFunc) Option {
	return func(h *Handler) {
		h.failure = append(h.failure, failure...)
	}
}

// WithClientFailure allows zero or more go-kit ClientResponseFuncs to be used as fanout failure functions.
func WithClientFailure(failure ...gokithttp.ClientResponseFunc) Option {
	return func(h *Handler) {
		for _, rf := range failure {
			h.failure = append(
				h.failure,
				func(ctx context.Context, response http.ResponseWriter, result Result) context.Context {
					return rf(ctx, result.Response)
				},
			)
		}
	}
}

// WithConfiguration uses a set of (typically injected) fanout configuration options to configure a Handler.
// Use of this option will not override the configured Endpoints instance.
func WithConfiguration(c Configuration) Option {
	return func(h *Handler) {
		WithTransactor(NewTransactor(c))(h)

		authorization := c.authorization()
		if len(authorization) > 0 {
			WithClientBefore(gokithttp.SetRequestHeader("Authorization", authorization))(h)
		}
	}
}

// Handler is the http.Handler that fans out HTTP requests using the configured Endpoints strategy.
type Handler struct {
	endpoints       Endpoints
	errorEncoder    gokithttp.ErrorEncoder
	before          []FanoutRequestFunc
	after           []FanoutResponseFunc
	failure         []FanoutResponseFunc
	shouldTerminate ShouldTerminateFunc
	transactor      func(*http.Request) (*http.Response, error)
}

// New creates a fanout Handler.  The Endpoints strategy is required, and this constructor function will
// panic if it is nil.
//
// By default, all fanout requests have the same HTTP method as the original request, but no body is set..  Clients must use the OriginalBody
// strategy to set the original request's body on each fanout request.
func New(e Endpoints, options ...Option) *Handler {
	if e == nil {
		panic("An Endpoints strategy is required")
	}

	h := &Handler{
		endpoints:       e,
		errorEncoder:    gokithttp.DefaultErrorEncoder,
		shouldTerminate: DefaultShouldTerminate,
		transactor:      http.DefaultClient.Do,
	}

	for _, o := range options {
		o(h)
	}

	return h
}

// newFanoutRequests uses the Endpoints strategy and builds (1) HTTP request for each endpoint.  The configured
// FanoutRequestFunc options are used to build each request.  This method returns an error if no endpoints were returned
// by the strategy or if an error reading the original request body occurred.
func (h *Handler) newFanoutRequests(fanoutCtx context.Context, original *http.Request) ([]*http.Request, error) {
	body, err := ioutil.ReadAll(original.Body)
	if err != nil {
		return nil, err
	}

	urls, err := h.endpoints.FanoutURLs(original)
	if err != nil {
		return nil, err
	} else if len(urls) == 0 {
		return nil, errNoFanoutURLs
	}

	requests := make([]*http.Request, len(urls))
	for i := 0; i < len(urls); i++ {
		fanout := &http.Request{
			Method:     original.Method,
			URL:        urls[i],
			Proto:      "HTTP/1.1",
			ProtoMajor: 1,
			ProtoMinor: 1,
			Header:     make(http.Header),
			Host:       urls[i].Host,
		}

		endpointCtx := fanoutCtx
		var err error
		for _, rf := range h.before {
			endpointCtx, err = rf(endpointCtx, original, fanout, body)
			if err != nil {
				return nil, err
			}
		}

		requests[i] = fanout.WithContext(endpointCtx)
	}

	return requests, nil
}

// execute performs a single fanout HTTP transaction and sends the result on a channel.  This method is invoked
// as a goroutine.  It takes care of draining the fanout's response prior to returning.
func (h *Handler) execute(logger log.Logger, spanner tracing.Spanner, results chan<- Result, request *http.Request) {
	var (
		finisher = spanner.Start(request.URL.String())
		result   = Result{
			Request: request,
		}
	)

	result.Response, result.Err = h.transactor(request)
	switch {
	case result.Response != nil:
		result.StatusCode = result.Response.StatusCode
		result.ContentType = result.Response.Header.Get("Content-Type")

		var err error
		if result.Body, err = ioutil.ReadAll(result.Response.Body); err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "error reading fanout response body", logging.ErrorKey(), err)
		}

		if err = result.Response.Body.Close(); err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "error closing fanout response body", logging.ErrorKey(), err)
		}

	case result.Err != nil:
		result.Body = []byte(fmt.Sprintf("%s", result.Err))
		result.ContentType = "text/plain"

		if ue, ok := result.Err.(*url.Error); ok && ue.Err != nil {
			// unwrap the URL error
			result.Err = ue.Err
		}

		if result.Err == context.Canceled || result.Err == context.DeadlineExceeded {
			result.StatusCode = http.StatusGatewayTimeout
		} else {
			result.StatusCode = http.StatusServiceUnavailable
		}

	default:
		// this "should" never happen, but just in case set a known status code
		result.StatusCode = http.StatusInternalServerError
		result.Err = errBadTransactor
		result.Body = []byte(errBadTransactor.Error())
		result.ContentType = "test/plain"
	}

	result.Span = finisher(result.Err)
	results <- result
}

// finish takes a terminating fanout result and writes the appropriate information to the top-level response.  This method
// is only invoked when a particular fanout response terminates the fanout, i.e. is considered successful.
func (h *Handler) finish(logger log.Logger, response http.ResponseWriter, result Result) {
	ctx := result.Request.Context()
	for _, rf := range h.after {
		// NOTE: we don't use the context for anything here,
		// but to preserve go-kit semantics we pass it to each after function
		ctx = rf(ctx, response, result)
	}

	if len(result.Body) > 0 {
		if len(result.ContentType) > 0 {
			response.Header().Set("Content-Type", result.ContentType)
		} else {
			response.Header().Set("Content-Type", "application/octet-stream")
		}

		response.WriteHeader(result.StatusCode)
		count, err := response.Write(result.Body)
		if err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "wrote fanout response", "bytes", count, logging.ErrorKey(), err)
		} else {
			logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "wrote fanout response", "bytes", count)
		}

	} else {
		response.WriteHeader(result.StatusCode)
	}
}

// finish takes a terminating fanout result and writes the appropriate information to the top-level response.  This method
// is only invoked when a particular fanout response terminates the fanout, and is considered a failure
func (h *Handler) handleErrorFinish(logger log.Logger, response http.ResponseWriter, result Result) {
	ctx := result.Request.Context()
	for _, rf := range h.failure {
		// NOTE: we don't use the context for anything here,
		// but to preserve go-kit semantics we pass it to each after function
		ctx = rf(ctx, response, result)
	}

	if len(result.Body) > 0 {
		if len(result.ContentType) > 0 {
			response.Header().Set("Content-Type", result.ContentType)
		} else {
			response.Header().Set("Content-Type", "application/octet-stream")
		}

		response.WriteHeader(result.StatusCode)
		count, err := response.Write(result.Body)
		if err != nil {
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "wrote fanout error response", "bytes", count, logging.ErrorKey(), err)
		} else {
			logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "wrote fanout error response", "bytes", count)
		}
	} else {
		logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "wrote fanout error response", "statusCode", result.StatusCode)
		response.WriteHeader(result.StatusCode)
	}
}

func (h *Handler) ServeHTTP(response http.ResponseWriter, original *http.Request) {
	var (
		fanoutCtx     = original.Context()
		logger        = logging.GetLogger(fanoutCtx)
		requests, err = h.newFanoutRequests(fanoutCtx, original)
	)

	if err != nil {
		logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "unable to create fanout", logging.ErrorKey(), err)
		h.errorEncoder(fanoutCtx, err, response)
		return
	}

	var (
		spanner = tracing.NewSpanner()
		results = make(chan Result, len(requests))
	)

	for _, r := range requests {
		go h.execute(logger, spanner, results, r)
	}

	statusCode := 0
	var latestResponse Result
	for i := 0; i < len(requests); i++ {
		select {
		case <-fanoutCtx.Done():
			logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "fanout operation canceled or timed out", "statusCode", http.StatusGatewayTimeout, "url", original.URL, logging.ErrorKey(), fanoutCtx.Err())
			response.WriteHeader(http.StatusGatewayTimeout)
			return

		case r := <-results:
			tracinghttp.HeadersForSpans("", response.Header(), r.Span)
			if r.Err != nil {
				logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "fanout request complete", "statusCode", r.StatusCode, "url", r.Request.URL, logging.ErrorKey(), r.Err)
			} else {
				logger.Log(level.Key(), level.DebugValue(), logging.MessageKey(), "fanout request complete", "statusCode", r.StatusCode, "url", r.Request.URL)
			}

			if h.shouldTerminate(r) {
				// this was a "success", so no reason to wait any longer
				h.finish(logger, response, r)
				return
			}

			if statusCode < r.StatusCode {
				statusCode = r.StatusCode
				latestResponse = r
			}
		}
	}

	logger.Log(level.Key(), level.ErrorValue(), logging.MessageKey(), "all fanout requests failed", "statusCode", statusCode, "url", original.URL)
	h.handleErrorFinish(logger, response, latestResponse)
}
<|endoftext|>


<commit_before>package util

import (
	"github.com/ghts/ghts/lib"
	"github.com/ghts/ghts/lib/daily_price_data"
	"github.com/ghts/ghts/lib/krx_time"
	xt "github.com/ghts/ghts/xing/base"
	xing "github.com/ghts/ghts/xing/go"

	"database/sql"
	"time"
)

func F당일_일일_가격정보_수집(db *sql.DB) (에러 error) {
	if krx.F한국증시_동시호가_시간임() || krx.F한국증시_정규_거래_시간임() {
		lib.F문자열_출력("장 중에는 정확한 데이터를 수집할 수 없습니다.")
		return
	}

	daily_price_data.F일일_가격정보_테이블_생성(db)

	종목코드_모음_전체 := xing.F종목코드_모음_전체()
	당일 := lib.F일자2정수(xing.F당일())

	for i := 0; i <= len(종목코드_모음_전체)/50; i++ {
		var 종목코드_모음 []string

		if len(종목코드_모음_전체) > (i+1)*50 {
			종목코드_모음 = 종목코드_모음_전체[i*50 : (i+1)*50]
		} else {
			종목코드_모음 = 종목코드_모음_전체[i*50:]
		}

		응답값_맵, 에러 := xing.TrT8407_현물_멀티_현재가_조회(종목코드_모음)
		lib.F확인(에러)

		for 종목코드, 값 := range 응답값_맵 {
			s := new(daily_price_data.S일일_가격정보)
			s.M종목코드 = 종목코드
			s.M일자 = 당일
			s.M시가 = float64(값.M시가)
			s.M고가 = float64(값.M고가)
			s.M저가 = float64(값.M저가)
			s.M종가 = float64(값.M현재가)
			s.M거래량 = float64(값.M누적_거래량)

			종목별_일일_가격정보_모음, 에러 := daily_price_data.New종목별_일일_가격정보_모음([]*daily_price_data.S일일_가격정보{s})
			lib.F확인(에러)

			lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
		}
	}

	lib.F문자열_출력("당일 가격정보 수집 완료.")

	return nil
}

func F일개월_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	var 시작일, 종료일 time.Time
	var 종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음

	daily_price_data.F일일_가격정보_테이블_생성(db)

	for i, 종목코드 := range 종목코드_모음 {
		for {
			if xing.C32_재시작_실행_중.G값() {
				lib.F대기(lib.P10초)
				continue
			}

			break
		}

		종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음_DB읽기(db, 종목코드)
		lib.F확인(에러)

		// 시작일 설정. 데이터 수량이 1개이나 100개이나 소요 시간은 비슷함.
		시작일 = lib.F지금().AddDate(0, 0, -31)

		// 종료일 설정
		if lib.F지금().After(xing.F당일().Add(15*lib.P1시간 + lib.P30분)) {
			종료일 = xing.F당일()
		} else {
			종료일 = xing.F전일()
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, 종료일, 종목별_일일_가격정보_모음, i)
	}

	return nil
}

func F일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	var 시작일, 종료일, 마지막_저장일 time.Time
	var 종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음

	daily_price_data.F일일_가격정보_테이블_생성(db)

	for i, 종목코드 := range 종목코드_모음 {
		for {
			if xing.C32_재시작_실행_중.G값() {
				lib.F대기(lib.P10초)
				continue
			}

			break
		}

		종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음_DB읽기(db, 종목코드)
		lib.F확인(에러)

		// 시작일 설정
		시작일 = lib.F지금().AddDate(-30, 0, 0)
		if 에러 == nil && len(종목별_일일_가격정보_모음.M저장소) > 0 {
			// lib.S종목별_일일_가격정보_모음 는 일자 순서로 정렬되어 있음.
			마지막_저장일 = 종목별_일일_가격정보_모음.M저장소[len(종목별_일일_가격정보_모음.M저장소)-1].G일자2()
			시작일 = 마지막_저장일.AddDate(0, 0, 1)
		}

		if 시작일.After(xing.F당일()) {
			lib.F문자열_출력("%v [%v] : 최신 데이터 업데이트.", i, 종목코드)
			continue
		}

		// 종료일 설정
		if lib.F지금().After(xing.F당일().Add(15*lib.P1시간 + lib.P30분)) {
			종료일 = xing.F당일()
		} else {
			종료일 = xing.F전일()
		}

		if 시작일.After(종료일) {
			continue
		} else if 시작일.Equal(종료일) { // 시작일과 종료일이 같으면 수천 개의 데이터를 불러오는 현상이 있음.
			시작일 = 시작일.AddDate(0, 0, -1)
		}

		// 데이터 수량이 1개이나 100개이나 소요 시간은 비슷함.
		if 시작일.After(lib.F금일().AddDate(0, 0, -14)) {
			시작일 = lib.F금일().AddDate(0, 0, -14)
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, 종료일, 종목별_일일_가격정보_모음, i)
	}

	return nil
}

func f일일_가격정보_수집_도우미(db *sql.DB,
	종목코드 string, 시작일, 종료일 time.Time,
	종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음, i int) {
	// 데이터 수집
	값_모음, 에러 := xing.TrT8413_현물_차트_일주월(종목코드, 시작일, 종료일, xt.P일주월_일)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	} else if len(값_모음) == 0 {
		lib.F체크포인트(i, 종목코드, "추가 저장할 데이터가 없음.")
		return // 추가 저장할 데이터가 없음.
	}

	일일_가격정보_슬라이스 := make([]*daily_price_data.S일일_가격정보, len(값_모음))

	for i, 일일_데이터 := range 값_모음 {
		일일_가격정보_슬라이스[i] = daily_price_data.New일일_가격정보(
			일일_데이터.M종목코드,
			일일_데이터.M일자,
			일일_데이터.M시가,
			일일_데이터.M고가,
			일일_데이터.M저가,
			일일_데이터.M종가,
			일일_데이터.M거래량)
	}

	lib.F문자열_출력("%v %v %v~%v %v개", i, 종목코드, 시작일.Format(lib.P일자_형식), 종료일.Format(lib.P일자_형식), len(값_모음))

	종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음(일일_가격정보_슬라이스)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	}

	lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
}
<commit_msg>코드 정리<commit_after>package util

import (
	"github.com/ghts/ghts/lib"
	"github.com/ghts/ghts/lib/daily_price_data"
	"github.com/ghts/ghts/lib/krx_time"
	xt "github.com/ghts/ghts/xing/base"
	xing "github.com/ghts/ghts/xing/go"

	"database/sql"
	"time"
)

func F당일_일일_가격정보_수집(db *sql.DB) (에러 error) {
	if krx.F한국증시_동시호가_시간임() || krx.F한국증시_정규_거래_시간임() {
		lib.F문자열_출력("장 중에는 정확한 데이터를 수집할 수 없습니다.")
		return
	}

	daily_price_data.F일일_가격정보_테이블_생성(db)

	당일 := lib.F일자2정수(xing.F당일())
	현재가_맵, 에러 := xing.TrT8407_현물_멀티_현재가_조회_전종목()
	lib.F확인(에러)

	for 종목코드, 값 := range 현재가_맵 {
		s := new(daily_price_data.S일일_가격정보)
		s.M종목코드 = 종목코드
		s.M일자 = 당일
		s.M시가 = float64(값.M시가)
		s.M고가 = float64(값.M고가)
		s.M저가 = float64(값.M저가)
		s.M종가 = float64(값.M현재가)
		s.M거래량 = float64(값.M누적_거래량)

		종목별_일일_가격정보_모음, 에러 := daily_price_data.New종목별_일일_가격정보_모음([]*daily_price_data.S일일_가격정보{s})
		lib.F확인(에러)

		lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
	}

	lib.F문자열_출력("당일 가격정보 수집 완료.")

	return nil
}

func F일개월_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	var 시작일, 종료일 time.Time
	var 종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음

	daily_price_data.F일일_가격정보_테이블_생성(db)

	for i, 종목코드 := range 종목코드_모음 {
		for {
			if xing.C32_재시작_실행_중.G값() {
				lib.F대기(lib.P10초)
				continue
			}

			break
		}

		종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음_DB읽기(db, 종목코드)
		lib.F확인(에러)

		// 시작일 설정. 데이터 수량이 1개이나 100개이나 소요 시간은 비슷함.
		시작일 = lib.F지금().AddDate(0, 0, -31)

		// 종료일 설정
		if lib.F지금().After(xing.F당일().Add(15*lib.P1시간 + lib.P30분)) {
			종료일 = xing.F당일()
		} else {
			종료일 = xing.F전일()
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, 종료일, 종목별_일일_가격정보_모음, i)
	}

	return nil
}

func F일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	var 시작일, 종료일, 마지막_저장일 time.Time
	var 종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음

	daily_price_data.F일일_가격정보_테이블_생성(db)

	for i, 종목코드 := range 종목코드_모음 {
		for {
			if xing.C32_재시작_실행_중.G값() {
				lib.F대기(lib.P10초)
				continue
			}

			break
		}

		종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음_DB읽기(db, 종목코드)
		lib.F확인(에러)

		// 시작일 설정
		시작일 = lib.F지금().AddDate(-30, 0, 0)
		if 에러 == nil && len(종목별_일일_가격정보_모음.M저장소) > 0 {
			// lib.S종목별_일일_가격정보_모음 는 일자 순서로 정렬되어 있음.
			마지막_저장일 = 종목별_일일_가격정보_모음.M저장소[len(종목별_일일_가격정보_모음.M저장소)-1].G일자2()
			시작일 = 마지막_저장일.AddDate(0, 0, 1)
		}

		if 시작일.After(xing.F당일()) {
			lib.F문자열_출력("%v [%v] : 최신 데이터 업데이트.", i, 종목코드)
			continue
		}

		// 종료일 설정
		if lib.F지금().After(xing.F당일().Add(15*lib.P1시간 + lib.P30분)) {
			종료일 = xing.F당일()
		} else {
			종료일 = xing.F전일()
		}

		if 시작일.After(종료일) {
			continue
		} else if 시작일.Equal(종료일) { // 시작일과 종료일이 같으면 수천 개의 데이터를 불러오는 현상이 있음.
			시작일 = 시작일.AddDate(0, 0, -1)
		}

		// 데이터 수량이 1개이나 100개이나 소요 시간은 비슷함.
		if 시작일.After(lib.F금일().AddDate(0, 0, -14)) {
			시작일 = lib.F금일().AddDate(0, 0, -14)
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, 종료일, 종목별_일일_가격정보_모음, i)
	}

	return nil
}

func f일일_가격정보_수집_도우미(db *sql.DB,
	종목코드 string, 시작일, 종료일 time.Time,
	종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음, i int) {
	// 데이터 수집
	값_모음, 에러 := xing.TrT8413_현물_차트_일주월(종목코드, 시작일, 종료일, xt.P일주월_일)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	} else if len(값_모음) == 0 {
		lib.F체크포인트(i, 종목코드, "추가 저장할 데이터가 없음.")
		return // 추가 저장할 데이터가 없음.
	}

	일일_가격정보_슬라이스 := make([]*daily_price_data.S일일_가격정보, len(값_모음))

	for i, 일일_데이터 := range 값_모음 {
		일일_가격정보_슬라이스[i] = daily_price_data.New일일_가격정보(
			일일_데이터.M종목코드,
			일일_데이터.M일자,
			일일_데이터.M시가,
			일일_데이터.M고가,
			일일_데이터.M저가,
			일일_데이터.M종가,
			일일_데이터.M거래량)
	}

	lib.F문자열_출력("%v %v %v~%v %v개", i, 종목코드, 시작일.Format(lib.P일자_형식), 종료일.Format(lib.P일자_형식), len(값_모음))

	종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음(일일_가격정보_슬라이스)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	}

	lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
}
<|endoftext|>


<commit_before>package util

import (
	"github.com/ghts/ghts/lib"
	"github.com/ghts/ghts/lib/daily_price_data"
	"github.com/ghts/ghts/lib/krx_time"
	xt "github.com/ghts/ghts/xing/base"
	xing "github.com/ghts/ghts/xing/go"

	"bytes"
	"database/sql"
	"time"
)

func F당일_일일_가격정보_수집(db *sql.DB) (에러 error) {
	if krx.F한국증시_동시호가_시간임() || krx.F한국증시_정규_거래_시간임() {
		lib.F문자열_출력("장 중에는 정확한 데이터를 수집할 수 없습니다.")
		return
	}

	daily_price_data.F일일_가격정보_테이블_생성(db)

	당일 := lib.F일자2정수(xing.F당일())
	현재가_맵, 에러 := xing.TrT8407_현물_멀티_현재가_조회_전종목()
	lib.F확인(에러)

	for 종목코드, 값 := range 현재가_맵 {
		s := new(daily_price_data.S일일_가격정보)
		s.M종목코드 = 종목코드
		s.M일자 = 당일
		s.M시가 = float64(값.M시가)
		s.M고가 = float64(값.M고가)
		s.M저가 = float64(값.M저가)
		s.M종가 = float64(값.M현재가)
		s.M거래량 = float64(값.M누적_거래량)

		종목별_일일_가격정보_모음, 에러 := daily_price_data.New종목별_일일_가격정보_모음([]*daily_price_data.S일일_가격정보{s})
		lib.F확인(에러)

		lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
	}

	lib.F문자열_출력("당일 가격정보 수집 완료.")

	return nil
}

func F일개월_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	return F고정_기간_일일_가격정보_수집(db, 종목코드_모음, 31*lib.P1일)
}

func F일년_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	return F고정_기간_일일_가격정보_수집(db, 종목코드_모음, lib.P1년)
}

func F고정_기간_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string, 기간 time.Duration, 추가_인수 ...bool) (에러 error) {
	if len(종목코드_모음) == 0 {
		return nil
	}

	daily_price_data.F일일_가격정보_테이블_생성(db)

	시작일 := lib.F금일().Add(-1 * 기간)

	출력_여부 := true
	if len(추가_인수) > 0 {
		출력_여부 = 추가_인수[0]
	}

	종목코드_맵 := make(map[string]lib.S비어있음) // 종목 순서를 랜덤화

	for _, 종목코드 := range 종목코드_모음 {
		종목코드_맵[종목코드] = lib.F비어있는_값()
	}

	i := 0

	for 종목코드 := range 종목코드_맵 {
		if lib.F공통_종료_채널_닫힘() {
			return nil
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, i, len(종목코드_맵), 출력_여부)
		i++

		lib.F대기(lib.P4초) // TR 한도 초과 관련.
	}

	return nil
}

func F일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string, 추가_인수 ...bool) (에러 error) {
	var 시작일, 마지막_저장일 time.Time
	var 종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음

	daily_price_data.F일일_가격정보_테이블_생성(db)

	출력_여부 := lib.F조건부_참거짓(len(추가_인수) > 0, 추가_인수[0], true)
	출력_문자열_버퍼 := new(bytes.Buffer)

	for i, 종목코드 := range 종목코드_모음 {
		종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음_DB읽기(db, 종목코드)
		lib.F확인(에러)

		// 시작일 설정
		시작일 = lib.F지금().AddDate(-30, 0, 0)
		if 에러 == nil && len(종목별_일일_가격정보_모음.M저장소) > 0 {
			// lib.S종목별_일일_가격정보_모음 는 일자 순서로 정렬되어 있음.
			마지막_저장일 = 종목별_일일_가격정보_모음.M저장소[len(종목별_일일_가격정보_모음.M저장소)-1].G일자2()
			시작일 = 마지막_저장일.AddDate(0, 0, 1)
		}

		if 시작일.After(xing.F당일()) {
			//lib.F문자열_출력("%v [%v] : 최신 데이터 업데이트.", i, 종목코드)
			continue
		} else if 시작일.After(lib.F금일().AddDate(0, 0, -14)) {
			// 데이터 수량이 1개이나 100개이나 소요 시간은 비슷함.
			시작일 = lib.F금일().AddDate(0, 0, -14)
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, i, len(종목코드_모음), 출력_여부, 출력_문자열_버퍼)
	}

	lib.F문자열_출력(출력_문자열_버퍼.String())

	return nil
}

func f일일_가격정보_수집_도우미(db *sql.DB, 종목코드 string, 시작일 time.Time, i, 전체_수량 int, 출력_여부 bool, 버퍼 ...*bytes.Buffer) {
	var 종료일 time.Time

	// 종료일 설정
	if lib.F지금().After(xing.F당일().Add(15*lib.P1시간 + lib.P30분)) {
		종료일 = xing.F당일()
	} else {
		종료일 = xing.F전일()
	}

	// 시작일 오류 확인
	if 시작일 = lib.F2일자(시작일); 시작일.After(종료일) {
		return
	} else if 시작일.Equal(종료일) { // 시작일과 종료일이 같으면 수천 개의 데이터를 불러오는 현상이 있음.
		시작일 = 시작일.AddDate(0, 0, -1)
	}

	// 데이터 수집
	값_모음, 에러 := xing.TrT8413_현물_차트_일주월(종목코드, 시작일, 종료일, xt.P일주월_일)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	} else if len(값_모음) == 0 {
		lib.F체크포인트(i, 종목코드, "추가 저장할 데이터가 없음.")
		return // 추가 저장할 데이터가 없음.
	}

	일일_가격정보_슬라이스 := make([]*daily_price_data.S일일_가격정보, len(값_모음))

	for i, 일일_데이터 := range 값_모음 {
		일일_가격정보_슬라이스[i] = daily_price_data.New일일_가격정보(
			일일_데이터.M종목코드,
			일일_데이터.M일자,
			일일_데이터.M시가,
			일일_데이터.M고가,
			일일_데이터.M저가,
			일일_데이터.M종가,
			일일_데이터.M거래량)
	}

	if 출력_여부 && len(버퍼) > 0 && 버퍼[0] != nil {
		// 버퍼가 존재하면 버퍼에 출력
		버퍼[0].WriteString(lib.F2문자열("%v 일일 가격 정보 수집 (%v/%v) : %v %v~%v %v개\n",
			lib.F지금().Format("15:04"), i+1, 전체_수량,
			xing.F종목_식별_문자열(종목코드), 시작일.Format(lib.P일자_형식), 종료일.Format(lib.P일자_형식), len(값_모음)))
	} else if 출력_여부 {
		lib.F문자열_출력("%v 일일 가격 정보 수집 (%v/%v) %.1f%% : %v %v~%v %v개\n",
			lib.F지금().Format("15:04"), i+1, 전체_수량, float64(i+1)/float64(전체_수량)*100,
			xing.F종목_식별_문자열(종목코드), 시작일.Format(lib.P일자_형식), 종료일.Format(lib.P일자_형식), len(값_모음))
	} else if i%100 == 0 {
		lib.F문자열_출력("%v 일일 가격 정보 수집 (%v/%v) %.1f%%",
			lib.F지금().Format("15:04"), i+1, 전체_수량, float64(i+1)/float64(전체_수량)*100)
	}

	종목별_일일_가격정보_모음, 에러 := daily_price_data.New종목별_일일_가격정보_모음(일일_가격정보_슬라이스)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	}

	lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
}
<commit_msg>일일 가격정보 수집 출력 메시지 수정.<commit_after>package util

import (
	"github.com/ghts/ghts/lib"
	"github.com/ghts/ghts/lib/daily_price_data"
	"github.com/ghts/ghts/lib/krx_time"
	xt "github.com/ghts/ghts/xing/base"
	xing "github.com/ghts/ghts/xing/go"

	"bytes"
	"database/sql"
	"time"
)

func F당일_일일_가격정보_수집(db *sql.DB) (에러 error) {
	if krx.F한국증시_동시호가_시간임() || krx.F한국증시_정규_거래_시간임() {
		lib.F문자열_출력("장 중에는 정확한 데이터를 수집할 수 없습니다.")
		return
	}

	daily_price_data.F일일_가격정보_테이블_생성(db)

	당일 := lib.F일자2정수(xing.F당일())
	현재가_맵, 에러 := xing.TrT8407_현물_멀티_현재가_조회_전종목()
	lib.F확인(에러)

	for 종목코드, 값 := range 현재가_맵 {
		s := new(daily_price_data.S일일_가격정보)
		s.M종목코드 = 종목코드
		s.M일자 = 당일
		s.M시가 = float64(값.M시가)
		s.M고가 = float64(값.M고가)
		s.M저가 = float64(값.M저가)
		s.M종가 = float64(값.M현재가)
		s.M거래량 = float64(값.M누적_거래량)

		종목별_일일_가격정보_모음, 에러 := daily_price_data.New종목별_일일_가격정보_모음([]*daily_price_data.S일일_가격정보{s})
		lib.F확인(에러)

		lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
	}

	lib.F문자열_출력("당일 가격정보 수집 완료.")

	return nil
}

func F일개월_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	return F고정_기간_일일_가격정보_수집(db, 종목코드_모음, 31*lib.P1일)
}

func F일년_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string) (에러 error) {
	return F고정_기간_일일_가격정보_수집(db, 종목코드_모음, lib.P1년)
}

func F고정_기간_일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string, 기간 time.Duration, 추가_인수 ...bool) (에러 error) {
	if len(종목코드_모음) == 0 {
		return nil
	}

	daily_price_data.F일일_가격정보_테이블_생성(db)

	시작일 := lib.F금일().Add(-1 * 기간)

	출력_여부 := true
	if len(추가_인수) > 0 {
		출력_여부 = 추가_인수[0]
	}

	종목코드_맵 := make(map[string]lib.S비어있음) // 종목 순서를 랜덤화

	for _, 종목코드 := range 종목코드_모음 {
		종목코드_맵[종목코드] = lib.F비어있는_값()
	}

	i := 0

	for 종목코드 := range 종목코드_맵 {
		if lib.F공통_종료_채널_닫힘() {
			return nil
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, i, len(종목코드_맵), 출력_여부)
		i++

		lib.F대기(lib.P4초) // TR 한도 초과 관련.
	}

	return nil
}

func F일일_가격정보_수집(db *sql.DB, 종목코드_모음 []string, 추가_인수 ...bool) (에러 error) {
	var 시작일, 마지막_저장일 time.Time
	var 종목별_일일_가격정보_모음 *daily_price_data.S종목별_일일_가격정보_모음

	daily_price_data.F일일_가격정보_테이블_생성(db)

	출력_여부 := lib.F조건부_참거짓(len(추가_인수) > 0, 추가_인수[0], true)
	출력_문자열_버퍼 := new(bytes.Buffer)

	for i, 종목코드 := range 종목코드_모음 {
		종목별_일일_가격정보_모음, 에러 = daily_price_data.New종목별_일일_가격정보_모음_DB읽기(db, 종목코드)
		lib.F확인(에러)

		// 시작일 설정
		시작일 = lib.F지금().AddDate(-30, 0, 0)
		if 에러 == nil && len(종목별_일일_가격정보_모음.M저장소) > 0 {
			// lib.S종목별_일일_가격정보_모음 는 일자 순서로 정렬되어 있음.
			마지막_저장일 = 종목별_일일_가격정보_모음.M저장소[len(종목별_일일_가격정보_모음.M저장소)-1].G일자2()
			시작일 = 마지막_저장일.AddDate(0, 0, 1)
		}

		if 시작일.After(xing.F당일()) {
			//lib.F문자열_출력("%v [%v] : 최신 데이터 업데이트.", i, 종목코드)
			continue
		} else if 시작일.After(lib.F금일().AddDate(0, 0, -14)) {
			// 데이터 수량이 1개이나 100개이나 소요 시간은 비슷함.
			시작일 = lib.F금일().AddDate(0, 0, -14)
		}

		f일일_가격정보_수집_도우미(db, 종목코드, 시작일, i, len(종목코드_모음), 출력_여부, 출력_문자열_버퍼)
	}

	lib.F문자열_출력(출력_문자열_버퍼.String())

	return nil
}

func f일일_가격정보_수집_도우미(db *sql.DB, 종목코드 string, 시작일 time.Time, i, 전체_수량 int, 출력_여부 bool, 버퍼 ...*bytes.Buffer) {
	var 종료일 time.Time

	// 종료일 설정
	if lib.F지금().After(xing.F당일().Add(15*lib.P1시간 + lib.P30분)) {
		종료일 = xing.F당일()
	} else {
		종료일 = xing.F전일()
	}

	// 시작일 오류 확인
	if 시작일 = lib.F2일자(시작일); 시작일.After(종료일) {
		return
	} else if 시작일.Equal(종료일) { // 시작일과 종료일이 같으면 수천 개의 데이터를 불러오는 현상이 있음.
		시작일 = 시작일.AddDate(0, 0, -1)
	}

	// 데이터 수집
	값_모음, 에러 := xing.TrT8413_현물_차트_일주월(종목코드, 시작일, 종료일, xt.P일주월_일)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	} else if len(값_모음) == 0 {
		lib.F체크포인트(i, 종목코드, "추가 저장할 데이터가 없음.")
		return // 추가 저장할 데이터가 없음.
	}

	일일_가격정보_슬라이스 := make([]*daily_price_data.S일일_가격정보, len(값_모음))

	for i, 일일_데이터 := range 값_모음 {
		일일_가격정보_슬라이스[i] = daily_price_data.New일일_가격정보(
			일일_데이터.M종목코드,
			일일_데이터.M일자,
			일일_데이터.M시가,
			일일_데이터.M고가,
			일일_데이터.M저가,
			일일_데이터.M종가,
			일일_데이터.M거래량)
	}

	if 출력_여부 && len(버퍼) > 0 && 버퍼[0] != nil {
		// 버퍼가 존재하면 버퍼에 출력
		버퍼[0].WriteString(lib.F2문자열("%v 일일 가격 정보 수집 (%v/%v) : %v %v~%v %v개\n",
			lib.F지금().Format("15:04"), i+1, 전체_수량,
			xing.F종목_식별_문자열(종목코드), 시작일.Format(lib.P일자_형식), 종료일.Format(lib.P일자_형식), len(값_모음)))
	} else if 출력_여부 {
		lib.F문자열_출력("%v 일일 가격 정보 수집 (%v/%v) %.1f%% : %v %v~%v %v개\n",
			lib.F지금().Format("15:04"), i+1, 전체_수량, float64(i+1)/float64(전체_수량)*100,
			xing.F종목_식별_문자열(종목코드), 시작일.Format(lib.P일자_형식), 종료일.Format(lib.P일자_형식), len(값_모음))
	} else if i>0 && i%100 == 0 {
		lib.F문자열_출력("%v 일일 가격 정보 수집 (%v/%v) %.1f%%",
			lib.F지금().Format("15:04"), i+1, 전체_수량, float64(i+1)/float64(전체_수량)*100)
	}

	종목별_일일_가격정보_모음, 에러 := daily_price_data.New종목별_일일_가격정보_모음(일일_가격정보_슬라이스)
	if 에러 != nil {
		lib.F에러_출력(에러)
		return
	}

	lib.F확인(종목별_일일_가격정보_모음.DB저장(db))
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"io/ioutil"
	"path/filepath"

	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/discovery"
	"github.com/TheThingsNetwork/ttn/core/types"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"golang.org/x/net/context" // See https://github.com/grpc/grpc-go/issues/711"
	"google.golang.org/grpc/metadata"
)

// brokerRegisterPrefixCmd represents the secure command
var brokerRegisterPrefixCmd = &cobra.Command{
	Use:   "register-prefix [prefix ...]",
	Short: "Register a prefix to this Broker",
	Long:  `ttn broker register prefix registers a prefix to this Broker`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) == 0 {
			cmd.UsageFunc()(cmd)
			return
		}

		path := filepath.Clean(viper.GetString("key-dir") + "/ca.cert")
		cert, err := ioutil.ReadFile(path)
		if err == nil && !api.RootCAs.AppendCertsFromPEM(cert) {
			ctx.Warnf("Could not add root certificates from %s", path)
		}

		conn, err := api.Dial(viper.GetString("discovery-address"))
		if err != nil {
			ctx.WithError(err).Fatal("Could not connect to Discovery server")
		}
		client := discovery.NewDiscoveryClient(conn)

		client.GetAll(context.Background(), &discovery.GetServiceRequest{})

		md := metadata.Pairs(
			"service-name", "broker",
			"id", viper.GetString("id"),
			"token", viper.GetString("auth-token"),
		)
		dscContext := metadata.NewContext(context.Background(), md)

		for _, prefixString := range args {
			ctx := ctx.WithField("Prefix", prefixString)
			prefix, err := types.ParseDevAddrPrefix(prefixString)
			if err != nil {
				ctx.WithError(err).Error("Could not register prefix")
				continue
			}
			_, err = client.AddMetadata(dscContext, &discovery.MetadataRequest{
				ServiceName: "broker",
				Id:          viper.GetString("id"),
				Metadata: &discovery.Metadata{Metadata: &discovery.Metadata_DevAddrPrefix{
					DevAddrPrefix: prefix.Bytes(),
				}},
			})
			if err != nil {
				ctx.WithError(err).Error("Could not register prefix")
			}
			ctx.Info("Registered prefix")
		}
	},
}

func init() {
	brokerCmd.AddCommand(brokerRegisterPrefixCmd)
}
<commit_msg>Don't say OK when it's not OK<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/discovery"
	"github.com/TheThingsNetwork/ttn/core/types"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"golang.org/x/net/context" // See https://github.com/grpc/grpc-go/issues/711"
	"google.golang.org/grpc/metadata"
)

// brokerRegisterPrefixCmd represents the secure command
var brokerRegisterPrefixCmd = &cobra.Command{
	Use:   "register-prefix [prefix ...]",
	Short: "Register a prefix to this Broker",
	Long:  `ttn broker register prefix registers a prefix to this Broker`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) == 0 {
			cmd.UsageFunc()(cmd)
			return
		}

		path := filepath.Clean(viper.GetString("key-dir") + "/ca.cert")
		cert, err := ioutil.ReadFile(path)
		if err == nil && !api.RootCAs.AppendCertsFromPEM(cert) {
			ctx.Warnf("Could not add root certificates from %s", path)
		}

		conn, err := api.Dial(viper.GetString("discovery-address"))
		if err != nil {
			ctx.WithError(err).Fatal("Could not connect to Discovery server")
		}
		client := discovery.NewDiscoveryClient(conn)

		client.GetAll(context.Background(), &discovery.GetServiceRequest{})

		md := metadata.Pairs(
			"service-name", "broker",
			"id", viper.GetString("id"),
			"token", viper.GetString("auth-token"),
		)
		dscContext := metadata.NewContext(context.Background(), md)

		success := true
		for _, prefixString := range args {
			ctx := ctx.WithField("Prefix", prefixString)
			prefix, err := types.ParseDevAddrPrefix(prefixString)
			if err != nil {
				ctx.WithError(err).Error("Could not register prefix")
				success = false
				continue
			}
			_, err = client.AddMetadata(dscContext, &discovery.MetadataRequest{
				ServiceName: "broker",
				Id:          viper.GetString("id"),
				Metadata: &discovery.Metadata{Metadata: &discovery.Metadata_DevAddrPrefix{
					DevAddrPrefix: prefix.Bytes(),
				}},
			})
			if err != nil {
				ctx.WithError(err).Error("Could not register prefix")
				success = false
				continue
			}
			ctx.Info("Registered prefix")
		}

		if !success {
			os.Exit(1)
		}
	},
}

func init() {
	brokerCmd.AddCommand(brokerRegisterPrefixCmd)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/oauth2"

	"github.com/gorilla/sessions"
	"github.com/nsheridan/cashier/lib"
	"github.com/nsheridan/cashier/server/auth"
	"github.com/nsheridan/cashier/server/auth/testprovider"
	"github.com/nsheridan/cashier/server/config"
	"github.com/nsheridan/cashier/server/signer"
	"github.com/nsheridan/cashier/server/store"
	"github.com/nsheridan/cashier/testdata"
)

func newContext(t *testing.T) *appContext {
	f, err := ioutil.TempFile(os.TempDir(), "signing_key_")
	if err != nil {
		t.Error(err)
	}
	defer os.Remove(f.Name())
	f.Write(testdata.Priv)
	f.Close()
	signer, err := signer.New(&config.SSH{
		SigningKey: f.Name(),
		MaxAge:     "1h",
	})
	if err != nil {
		t.Error(err)
	}
	return &appContext{
		cookiestore:  sessions.NewCookieStore([]byte("secret")),
		authprovider: testprovider.New(),
		certstore:    store.NewMemoryStore(),
		authsession:  &auth.Session{AuthURL: "https://www.example.com/auth"},
		sshKeySigner: signer,
	}
}

func TestLoginHandler(t *testing.T) {
	req, _ := http.NewRequest("GET", "/auth/login", nil)
	resp := httptest.NewRecorder()
	loginHandler(newContext(t), resp, req)
	if resp.Code != http.StatusFound && resp.Header().Get("Location") != "https://www.example.com/auth" {
		t.Error("Unexpected response")
	}
}

func TestCallbackHandler(t *testing.T) {
	req, _ := http.NewRequest("GET", "/auth/callback", nil)
	req.Form = url.Values{"state": []string{"state"}, "code": []string{"abcdef"}}
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	ctx.setAuthStateCookie(resp, req, "state")
	callbackHandler(ctx, resp, req)
	if resp.Code != http.StatusFound && resp.Header().Get("Location") != "/" {
		t.Error("Unexpected response")
	}
}

func TestRootHandler(t *testing.T) {
	req, _ := http.NewRequest("GET", "/", nil)
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	tok := &oauth2.Token{
		AccessToken: "XXX_TEST_TOKEN_STRING_XXX",
		Expiry:      time.Now().Add(1 * time.Hour),
	}
	ctx.setAuthTokenCookie(resp, req, tok)
	rootHandler(ctx, resp, req)
	if resp.Code != http.StatusOK && !strings.Contains(resp.Body.String(), "XXX_TEST_TOKEN_STRING_XXX") {
		t.Error("Unable to find token in response")
	}
}

func TestRootHandlerNoSession(t *testing.T) {
	req, _ := http.NewRequest("GET", "/", nil)
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	rootHandler(ctx, resp, req)
	if resp.Code != http.StatusSeeOther {
		t.Errorf("Unexpected status: %s, wanted %s", http.StatusText(resp.Code), http.StatusText(http.StatusSeeOther))
	}
}

func TestSignRevoke(t *testing.T) {
	t.Skip()
	s, _ := json.Marshal(&lib.SignRequest{
		Key: string(testdata.Pub),
	})
	req, _ := http.NewRequest("POST", "/sign", bytes.NewReader(s))
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	req.Header.Set("Authorization", "Bearer abcdef")
	signHandler(ctx, resp, req)
	if resp.Code != http.StatusOK {
		t.Error("Unexpected response")
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Error(err)
	}
	r := &lib.SignResponse{}
	if err := json.Unmarshal(b, r); err != nil {
		t.Error(err)
	}
	if r.Status != "ok" {
		t.Error("Unexpected response")
	}
	k, _, _, _, err := ssh.ParseAuthorizedKey([]byte(r.Response))
	if err != nil {
		t.Error(err)
	}
	cert, ok := k.(*ssh.Certificate)
	if !ok {
		t.Error("Did not receive a certificate")
	}
	// Revoke the cert and verify
	req, _ = http.NewRequest("POST", "/revoke", nil)
	req.Form = url.Values{"cert_id": []string{cert.KeyId}}
	revokeCertHandler(ctx, resp, req)
	req, _ = http.NewRequest("GET", "/revoked", nil)
	listRevokedCertsHandler(ctx, resp, req)
	revoked, _ := ioutil.ReadAll(resp.Body)
	if string(revoked[:len(revoked)-1]) != r.Response {
		t.Error("omg")
	}
}
<commit_msg>Fix and enable handers test<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/oauth2"

	"github.com/gorilla/sessions"
	"github.com/nsheridan/cashier/lib"
	"github.com/nsheridan/cashier/server/auth"
	"github.com/nsheridan/cashier/server/auth/testprovider"
	"github.com/nsheridan/cashier/server/config"
	"github.com/nsheridan/cashier/server/signer"
	"github.com/nsheridan/cashier/server/store"
	"github.com/nsheridan/cashier/testdata"
	"github.com/stripe/krl"
)

func newContext(t *testing.T) *appContext {
	f, err := ioutil.TempFile(os.TempDir(), "signing_key_")
	if err != nil {
		t.Error(err)
	}
	defer os.Remove(f.Name())
	f.Write(testdata.Priv)
	f.Close()
	signer, err := signer.New(&config.SSH{
		SigningKey: f.Name(),
		MaxAge:     "1h",
	})
	if err != nil {
		t.Error(err)
	}
	return &appContext{
		cookiestore:  sessions.NewCookieStore([]byte("secret")),
		authprovider: testprovider.New(),
		certstore:    store.NewMemoryStore(),
		authsession:  &auth.Session{AuthURL: "https://www.example.com/auth"},
		sshKeySigner: signer,
	}
}

func TestLoginHandler(t *testing.T) {
	req, _ := http.NewRequest("GET", "/auth/login", nil)
	resp := httptest.NewRecorder()
	loginHandler(newContext(t), resp, req)
	if resp.Code != http.StatusFound && resp.Header().Get("Location") != "https://www.example.com/auth" {
		t.Error("Unexpected response")
	}
}

func TestCallbackHandler(t *testing.T) {
	req, _ := http.NewRequest("GET", "/auth/callback", nil)
	req.Form = url.Values{"state": []string{"state"}, "code": []string{"abcdef"}}
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	ctx.setAuthStateCookie(resp, req, "state")
	callbackHandler(ctx, resp, req)
	if resp.Code != http.StatusFound && resp.Header().Get("Location") != "/" {
		t.Error("Unexpected response")
	}
}

func TestRootHandler(t *testing.T) {
	req, _ := http.NewRequest("GET", "/", nil)
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	tok := &oauth2.Token{
		AccessToken: "XXX_TEST_TOKEN_STRING_XXX",
		Expiry:      time.Now().Add(1 * time.Hour),
	}
	ctx.setAuthTokenCookie(resp, req, tok)
	rootHandler(ctx, resp, req)
	if resp.Code != http.StatusOK && !strings.Contains(resp.Body.String(), "XXX_TEST_TOKEN_STRING_XXX") {
		t.Error("Unable to find token in response")
	}
}

func TestRootHandlerNoSession(t *testing.T) {
	req, _ := http.NewRequest("GET", "/", nil)
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	rootHandler(ctx, resp, req)
	if resp.Code != http.StatusSeeOther {
		t.Errorf("Unexpected status: %s, wanted %s", http.StatusText(resp.Code), http.StatusText(http.StatusSeeOther))
	}
}

func TestSignRevoke(t *testing.T) {
	s, _ := json.Marshal(&lib.SignRequest{
		Key:        string(testdata.Pub),
		ValidUntil: time.Now().UTC().Add(1 * time.Hour),
	})
	req, _ := http.NewRequest("POST", "/sign", bytes.NewReader(s))
	resp := httptest.NewRecorder()
	ctx := newContext(t)
	req.Header.Set("Authorization", "Bearer abcdef")
	signHandler(ctx, resp, req)
	if resp.Code != http.StatusOK {
		t.Error("Unexpected response")
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Error(err)
	}
	r := &lib.SignResponse{}
	if err := json.Unmarshal(b, r); err != nil {
		t.Error(err)
	}
	if r.Status != "ok" {
		t.Error("Unexpected response")
	}
	k, _, _, _, err := ssh.ParseAuthorizedKey([]byte(r.Response))
	if err != nil {
		t.Error(err)
	}
	cert, ok := k.(*ssh.Certificate)
	if !ok {
		t.Error("Did not receive a certificate")
	}
	// Revoke the cert and verify
	req, _ = http.NewRequest("POST", "/revoke", nil)
	req.Form = url.Values{"cert_id": []string{cert.KeyId}}
	tok := &oauth2.Token{
		AccessToken: "authenticated",
		Expiry:      time.Now().Add(1 * time.Hour),
	}
	ctx.setAuthTokenCookie(resp, req, tok)
	revokeCertHandler(ctx, resp, req)
	req, _ = http.NewRequest("GET", "/revoked", nil)
	listRevokedCertsHandler(ctx, resp, req)
	revoked, _ := ioutil.ReadAll(resp.Body)
	rl, _ := krl.ParseKRL(revoked)
	if !rl.IsRevoked(cert) {
		t.Errorf("cert %s was not revoked", cert.KeyId)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
	"os"
	"syscall"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/libcontainerd"
	"github.com/docker/docker/pkg/system"
)

var defaultDaemonConfigFile = os.Getenv("programdata") + string(os.PathSeparator) + "docker" + string(os.PathSeparator) + "config" + string(os.PathSeparator) + "daemon.json"

// currentUserIsOwner checks whether the current user is the owner of the given
// file.
func currentUserIsOwner(f string) bool {
	return false
}

// setDefaultUmask doesn't do anything on windows
func setDefaultUmask() error {
	return nil
}

func getDaemonConfDir() string {
	return os.Getenv("PROGRAMDATA") + `\docker\config`
}

// notifySystem sends a message to the host when the server is ready to be used
func notifySystem() {
	if service != nil {
		err := service.started()
		if err != nil {
			logrus.Fatal(err)
		}
	}
}

// notifyShutdown is called after the daemon shuts down but before the process exits.
func notifyShutdown(err error) {
	if service != nil {
		service.stopped(err)
	}
}

// setupConfigReloadTrap configures a Win32 event to reload the configuration.
func (cli *DaemonCli) setupConfigReloadTrap() {
	go func() {
		sa := syscall.SecurityAttributes{
			Length: 0,
		}
		ev := "Global\\docker-daemon-config-" + fmt.Sprint(os.Getpid())
		if h, _ := system.CreateEvent(&sa, false, false, ev); h != 0 {
			logrus.Debugf("Config reload - waiting signal at %s", ev)
			for {
				syscall.WaitForSingleObject(h, syscall.INFINITE)
				cli.reloadConfig()
			}
		}
	}()
}

func (cli *DaemonCli) getPlatformRemoteOptions() []libcontainerd.RemoteOption {
	return nil
}

// getLibcontainerdRoot gets the root directory for libcontainerd to store its
// state. The Windows libcontainerd implementation does not need to write a spec
// or state to disk, so this is a no-op.
func (cli *DaemonCli) getLibcontainerdRoot() string {
	return ""
}

// getSwarmRunRoot gets the root directory for swarm to store runtime state
// For example, the control socket
func (cli *DaemonCli) getSwarmRunRoot() string {
	return ""
}

func allocateDaemonPort(addr string) error {
	return nil
}

func wrapListeners(proto string, ls []net.Listener) []net.Listener {
	return ls
}
<commit_msg>Write fatal error to event log if running as service<commit_after>package main

import (
	"fmt"
	"net"
	"os"
	"syscall"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/libcontainerd"
	"github.com/docker/docker/pkg/system"
)

var defaultDaemonConfigFile = os.Getenv("programdata") + string(os.PathSeparator) + "docker" + string(os.PathSeparator) + "config" + string(os.PathSeparator) + "daemon.json"

// currentUserIsOwner checks whether the current user is the owner of the given
// file.
func currentUserIsOwner(f string) bool {
	return false
}

// setDefaultUmask doesn't do anything on windows
func setDefaultUmask() error {
	return nil
}

func getDaemonConfDir() string {
	return os.Getenv("PROGRAMDATA") + `\docker\config`
}

// notifySystem sends a message to the host when the server is ready to be used
func notifySystem() {
	if service != nil {
		err := service.started()
		if err != nil {
			logrus.Fatal(err)
		}
	}
}

// notifyShutdown is called after the daemon shuts down but before the process exits.
func notifyShutdown(err error) {
	if service != nil {
		if err != nil {
			logrus.Fatal(err)
		}
		service.stopped(err)
	}
}

// setupConfigReloadTrap configures a Win32 event to reload the configuration.
func (cli *DaemonCli) setupConfigReloadTrap() {
	go func() {
		sa := syscall.SecurityAttributes{
			Length: 0,
		}
		ev := "Global\\docker-daemon-config-" + fmt.Sprint(os.Getpid())
		if h, _ := system.CreateEvent(&sa, false, false, ev); h != 0 {
			logrus.Debugf("Config reload - waiting signal at %s", ev)
			for {
				syscall.WaitForSingleObject(h, syscall.INFINITE)
				cli.reloadConfig()
			}
		}
	}()
}

func (cli *DaemonCli) getPlatformRemoteOptions() []libcontainerd.RemoteOption {
	return nil
}

// getLibcontainerdRoot gets the root directory for libcontainerd to store its
// state. The Windows libcontainerd implementation does not need to write a spec
// or state to disk, so this is a no-op.
func (cli *DaemonCli) getLibcontainerdRoot() string {
	return ""
}

// getSwarmRunRoot gets the root directory for swarm to store runtime state
// For example, the control socket
func (cli *DaemonCli) getSwarmRunRoot() string {
	return ""
}

func allocateDaemonPort(addr string) error {
	return nil
}

func wrapListeners(proto string, ls []net.Listener) []net.Listener {
	return ls
}
<|endoftext|>


<commit_before>package mysql

import (
	"hash/adler32"
	"strings"
)

// JobTable converts a topic name to a mysql table name.
func JobTable(topic string) string {
	return strings.Replace(topic, ".", "_", -1)
}

// HistoryTable converts a topic name to a mysql history table name.
func HistoryTable(topic string) string {
	return JobTable(topic) + "_archive"
}

// App_id convert a string appid to hash int which is used to locate shard.
func App_id(appid string) int {
	return int(adler32.Checksum([]byte(appid)))
}
<commit_msg>job table will all prefix with job_<commit_after>package mysql

import (
	"hash/adler32"
	"strings"
)

const jobTablePrefix = "job_"

// JobTable converts a topic name to a mysql table name.
func JobTable(topic string) string {
	return jobTablePrefix + strings.Replace(topic, ".", "_", -1)
}

// HistoryTable converts a topic name to a mysql history table name.
func HistoryTable(topic string) string {
	return JobTable(topic) + "_archive"
}

// App_id convert a string appid to hash int which is used to locate shard.
func App_id(appid string) int {
	return int(adler32.Checksum([]byte(appid)))
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All rights reserved.
//
// 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"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/google/slothfs/cache"
	"github.com/google/slothfs/fs"
	"github.com/google/slothfs/gitiles"
	"github.com/hanwen/go-fuse/fuse/nodefs"
)

func main() {
	repo := flag.String("repo", "", "Set the repository name.")
	debug := flag.Bool("debug", false, "Print FUSE debug info.")
	cacheDir := flag.String("cache", filepath.Join(os.Getenv("HOME"), ".cache", "slothfs"),
		"Set directory for file system cache.")
	gitilesOptions := gitiles.DefineFlags()
	flag.Parse()

	if *cacheDir == "" {
		log.Fatal("must set --cache")
	}
	if len(flag.Args()) < 1 {
		log.Fatal("usage: main -repo REPO MOUNT-POINT")
	}

	mntDir := flag.Arg(0)
	cache, err := cache.NewCache(*cacheDir, cache.Options{})
	if err != nil {
		log.Printf("NewCache: %v", err)
	}

	service, err := gitiles.NewService(*gitilesOptions)
	if err != nil {
		log.Printf("NewService: %v", err)
	}

	repoService := service.NewRepoService(*repo)
	project, err := repoService.Get()
	if err != nil {
		log.Fatalf("GetProject(%s): %v", *repo, err)
	}

	opts := fs.GitilesOptions{
		CloneURL: project.CloneURL,
	}

	root := fs.NewGitilesConfigFSRoot(cache, repoService, &opts)
	server, _, err := nodefs.MountRoot(mntDir, root, &nodefs.Options{
		EntryTimeout:    time.Hour,
		NegativeTimeout: time.Hour,
		AttrTimeout:     time.Hour,
		Debug:           *debug,
	})
	if err != nil {
		log.Fatalf("MountFileSystem: %v", err)
	}
	log.Printf("Started gitiles fs FUSE on %s", mntDir)
	server.Serve()
}
<commit_msg>slothfs-gitilesfs: exit on errors.<commit_after>// Copyright 2016 Google Inc. All rights reserved.
//
// 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"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/google/slothfs/cache"
	"github.com/google/slothfs/fs"
	"github.com/google/slothfs/gitiles"
	"github.com/hanwen/go-fuse/fuse/nodefs"
)

func main() {
	repo := flag.String("repo", "", "Set the repository name.")
	debug := flag.Bool("debug", false, "Print FUSE debug info.")
	cacheDir := flag.String("cache", filepath.Join(os.Getenv("HOME"), ".cache", "slothfs"),
		"Set directory for file system cache.")
	gitilesOptions := gitiles.DefineFlags()
	flag.Parse()

	if *cacheDir == "" {
		log.Fatal("must set --cache")
	}
	if len(flag.Args()) < 1 {
		log.Fatal("usage: main -repo REPO MOUNT-POINT")
	}

	mntDir := flag.Arg(0)
	cache, err := cache.NewCache(*cacheDir, cache.Options{})
	if err != nil {
		log.Fatalf("NewCache: %v", err)
	}

	service, err := gitiles.NewService(*gitilesOptions)
	if err != nil {
		log.Fatalf("NewService: %v", err)
	}

	repoService := service.NewRepoService(*repo)
	project, err := repoService.Get()
	if err != nil {
		log.Fatalf("GetProject(%s): %v", *repo, err)
	}

	opts := fs.GitilesOptions{
		CloneURL: project.CloneURL,
	}

	root := fs.NewGitilesConfigFSRoot(cache, repoService, &opts)
	server, _, err := nodefs.MountRoot(mntDir, root, &nodefs.Options{
		EntryTimeout:    time.Hour,
		NegativeTimeout: time.Hour,
		AttrTimeout:     time.Hour,
		Debug:           *debug,
	})
	if err != nil {
		log.Fatalf("MountFileSystem: %v", err)
	}
	log.Printf("Started gitiles fs FUSE on %s", mntDir)
	server.Serve()
}
<|endoftext|>


<commit_before>package agent

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"unicode/utf8"

	"github.com/mitchellh/go-homedir"
	"github.com/nsf/termbox-go"
	"github.com/weekface/togo/ui"
	"github.com/weekface/togo/util"
)

var help = `
Usage:

  list        List all the todo s.
  add <task>  Add a new todo task.
  help|h      Show help info.
	quit|exit|q Quit TOGO.

Press Ctr+C to quit.`

// type Agent present the togo terminate's all attributes.
type Agent struct {
	Ui *ui.DefaultUi

	// Now, it contain togo's prompt strings.
	Chars string

	// version
	Version string

	// new files path
	NewPath string

	// old files path
	OldPath string
}

// return a default Agent.
func New() *Agent {
	dir, _ := homedir.Dir()
	return &Agent{
		Ui: &ui.DefaultUi{
			Fg:          termbox.ColorWhite,
			Bg:          termbox.ColorBlack,
			AlertString: "Press what you want. Press Ctr+C to quit.",
		},
		Version: "0.3.0",
		NewPath: filepath.Join(dir, ".togo/new"),
		OldPath: filepath.Join(dir, ".togo/old"),
	}
}

// backspace key support
func (agent *Agent) DeletePromp() {
	_, size := utf8.DecodeLastRuneInString(agent.Chars)
	agent.Chars = agent.Chars[:len(agent.Chars)-size]
	agent.DrawPromp("")
}

// sugar
func (agent *Agent) Clear() {
	agent.Ui.Clear()
}

// sugar
func (agent *Agent) Flush() {
	agent.Ui.Flush()
}

// display help info
func (agent *Agent) ShowHelp() {
	agent.Clear()

	agent.Chars = ""
	agent.DrawPromp("")

	lines := strings.Split(help, "\n")

	x := 1
	y := 1

	for _, line := range lines {
		agent.Ui.PrintLine(line, x, y)
		x = 0
		y++
	}

	agent.Flush()
}

// list all new todos.
func (agent *Agent) ListNew() {
	files, err := filepath.Glob(filepath.Join(agent.NewPath, "*"))
	list := make([]string, 0)

	if err != nil {
		panic("Read todos fail!")
	}

	for _, file := range files {
		data, err := ioutil.ReadFile(file)
		if err != nil {
			panic(err)
		}

		list = append(list, string(data))
	}

	agent.Clear()
	agent.Ui.PrintLines(list, 0, 2)
	agent.Chars = ""
	agent.DrawPromp("")
	agent.Flush()
}

// add a todo.
func (agent *Agent) Add() {
	slice := strings.SplitN(agent.Chars, " ", 2)

	if len(slice) == 2 {
		filename := util.Hash(slice[1])
		ioutil.WriteFile(filepath.Join(agent.NewPath, filename), []byte(slice[1]), 0644)
		agent.ListNew()
	}
}

// quit.
func (agent *Agent) Quit() {
	termbox.Close()
	os.Exit(0)
}

// parse command. Now, just support help command to show the help info.
func (agent *Agent) ParseCmd() {
	command := strings.Split(agent.Chars, " ")[0]

	switch command {

	// help command.
	case "help", "h":
		agent.ShowHelp()

	// quit command.
	case "quit", "Quit", "Exit", "exit", "q":
		agent.Quit()

	// list command.
	case "list":
		agent.ListNew()

	// add command.
	case "add", "Add":
		agent.Add()

	// default, print what you press.
	default:
		agent.ShowHelp()
	}
}

// draw the promp line.
func (agent *Agent) DrawPromp(str string) {
	agent.Chars = agent.Chars + string(str)

	prompStr := "TOGO> " + agent.Chars

	x := 0
	y := 0

	width := agent.Ui.PrintLine(prompStr, x, y)
	agent.Ui.Cursor(width, 0)
	agent.Flush()
}

// the public api of the Agent package.
func (agent *Agent) Run() {
	err := termbox.Init()
	if err != nil {
		panic(err)
	}

	defer termbox.Close()
	agent.Clear()
	agent.DrawPromp("")

	agent.ListNew()
	agent.Flush()

loop:

	for {
		switch ev := termbox.PollEvent(); ev.Type {
		case termbox.EventKey:
			switch ev.Key {

			// quit the program.
			case termbox.KeyCtrlC:
				break loop

			// puts " " to the screen, when we press space.
			case termbox.KeySpace:
				str := " "
				agent.DrawPromp(str)

			// support back space key
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				agent.DeletePromp()

			// when we press Enter, parse the command.
			case termbox.KeyEnter:
				agent.ParseCmd()

			// convert rune to string, then draw the promp.
			default:
				str := string(ev.Ch)
				agent.DrawPromp(str)
			}
		}
	}
}
<commit_msg>Bump to v0.4.0<commit_after>package agent

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"unicode/utf8"

	"github.com/mitchellh/go-homedir"
	"github.com/nsf/termbox-go"
	"github.com/weekface/togo/ui"
	"github.com/weekface/togo/util"
)

var help = `
Usage:

  list        List all the todo s.
  add <task>  Add a new todo task.
  help|h      Show help info.
	quit|exit|q Quit TOGO.

Press Ctr+C to quit.`

// type Agent present the togo terminate's all attributes.
type Agent struct {
	Ui *ui.DefaultUi

	// Now, it contain togo's prompt strings.
	Chars string

	// version
	Version string

	// new files path
	NewPath string

	// old files path
	OldPath string
}

// return a default Agent.
func New() *Agent {
	dir, _ := homedir.Dir()
	return &Agent{
		Ui: &ui.DefaultUi{
			Fg:          termbox.ColorWhite,
			Bg:          termbox.ColorBlack,
			AlertString: "Press what you want. Press Ctr+C to quit.",
		},
		Version: "0.4.0",
		NewPath: filepath.Join(dir, ".togo/new"),
		OldPath: filepath.Join(dir, ".togo/old"),
	}
}

// backspace key support
func (agent *Agent) DeletePromp() {
	_, size := utf8.DecodeLastRuneInString(agent.Chars)
	agent.Chars = agent.Chars[:len(agent.Chars)-size]
	agent.DrawPromp("")
}

// sugar
func (agent *Agent) Clear() {
	agent.Ui.Clear()
}

// sugar
func (agent *Agent) Flush() {
	agent.Ui.Flush()
}

// display help info
func (agent *Agent) ShowHelp() {
	agent.Clear()

	agent.Chars = ""
	agent.DrawPromp("")

	lines := strings.Split(help, "\n")

	x := 1
	y := 1

	for _, line := range lines {
		agent.Ui.PrintLine(line, x, y)
		x = 0
		y++
	}

	agent.Flush()
}

// list all new todos.
func (agent *Agent) ListNew() {
	files, err := filepath.Glob(filepath.Join(agent.NewPath, "*"))
	list := make([]string, 0)

	if err != nil {
		panic("Read todos fail!")
	}

	for _, file := range files {
		data, err := ioutil.ReadFile(file)
		if err != nil {
			panic(err)
		}

		list = append(list, string(data))
	}

	agent.Clear()
	agent.Ui.PrintLines(list, 0, 2)
	agent.Chars = ""
	agent.DrawPromp("")
	agent.Flush()
}

// add a todo.
func (agent *Agent) Add() {
	slice := strings.SplitN(agent.Chars, " ", 2)

	if len(slice) == 2 {
		filename := util.Hash(slice[1])
		ioutil.WriteFile(filepath.Join(agent.NewPath, filename), []byte(slice[1]), 0644)
		agent.ListNew()
	}
}

// quit.
func (agent *Agent) Quit() {
	termbox.Close()
	os.Exit(0)
}

// parse command. Now, just support help command to show the help info.
func (agent *Agent) ParseCmd() {
	command := strings.Split(agent.Chars, " ")[0]

	switch command {

	// help command.
	case "help", "h":
		agent.ShowHelp()

	// quit command.
	case "quit", "Quit", "Exit", "exit", "q":
		agent.Quit()

	// list command.
	case "list":
		agent.ListNew()

	// add command.
	case "add", "Add":
		agent.Add()

	// default, print what you press.
	default:
		agent.ShowHelp()
	}
}

// draw the promp line.
func (agent *Agent) DrawPromp(str string) {
	agent.Chars = agent.Chars + string(str)

	prompStr := "TOGO> " + agent.Chars

	x := 0
	y := 0

	width := agent.Ui.PrintLine(prompStr, x, y)
	agent.Ui.Cursor(width, 0)
	agent.Flush()
}

// the public api of the Agent package.
func (agent *Agent) Run() {
	err := termbox.Init()
	if err != nil {
		panic(err)
	}

	defer termbox.Close()
	agent.Clear()
	agent.DrawPromp("")

	agent.ListNew()
	agent.Flush()

loop:

	for {
		switch ev := termbox.PollEvent(); ev.Type {
		case termbox.EventKey:
			switch ev.Key {

			// quit the program.
			case termbox.KeyCtrlC:
				break loop

			// puts " " to the screen, when we press space.
			case termbox.KeySpace:
				str := " "
				agent.DrawPromp(str)

			// support back space key
			case termbox.KeyBackspace, termbox.KeyBackspace2:
				agent.DeletePromp()

			// when we press Enter, parse the command.
			case termbox.KeyEnter:
				agent.ParseCmd()

			// convert rune to string, then draw the promp.
			default:
				str := string(ev.Ch)
				agent.DrawPromp(str)
			}
		}
	}
}
<|endoftext|>