<commit_before>// Copyright 2012 Brian "bojo" Jones. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package sessions

import (
	"bytes"
	"encoding/base32"
	"encoding/gob"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
)

// Amount of time for cookies/redis keys to expire.
var sessionExpire = 86400 * 30

// RediStore stores sessions in a redis backend.
type RediStore struct {
	Pool          *redis.Pool
	Codecs        []securecookie.Codec
	Options       *sessions.Options // default configuration
	DefaultMaxAge int               // default Redis TTL for a MaxAge == 0 session
	maxLength     int
}

// SetMaxLength sets RediStore.maxLength if the `l` argument is greater or equal 0
// maxLength restricts the maximum length of new sessions to l.
// If l is 0 there is no limit to the size of a session, use with caution.
// The default for a new RediStore is 4096. Redis allows for max.
// value sizes of up to 512MB (http://redis.io/topics/data-types)
// Default: 4096,
func (s *RediStore) SetMaxLength(l int) {
	if l >= 0 {
		s.maxLength = l
	}
}

// SetMaxAge restricts the maximum age, in seconds, of the session record
// both in database and a browser. This is to change session storage configuration.
// If you want just to remove session use your session `s` object and change it's
// `Options.MaxAge` to -1, as specified in
//    http://godoc.org/github.com/gorilla/sessions#Options
//
// Default is the one provided by this package value - `sessionExpire`.
// Set it to 0 for no restriction.
// Because we use `MaxAge` also in SecureCookie crypting algorithm you should
// use this function to change `MaxAge` value.
func (s *RediStore) SetMaxAge(v int) {
	var c *securecookie.SecureCookie
	var ok bool
	s.Options.MaxAge = v
	for i := range s.Codecs {
		if c, ok = s.Codecs[i].(*securecookie.SecureCookie); ok {
			c.MaxAge(v)
		} else {
			fmt.Printf("Can't change MaxAge on codec %t\n", s.Codecs[i])
		}
	}
}

func dial(network, server, password string) (redis.Conn, error) {

	c, err := redis.Dial(network, server)
	if err != nil {
		return nil, err
	}
	if password != "" {
		if _, err := c.Do("AUTH", password); err != nil {
			c.Close()
			return nil, err
		}
	}
	return c, err
}

// NewRediStore returns a new RediStore.
// size: maximum number of idle connections.
func NewRediStore(size int, network, address, password string, keyPairs ...[]byte) (*RediStore, error) {
	return NewRediStoreWithPool(&redis.Pool{
		MaxIdle:     size,
		IdleTimeout: 240 * time.Second,
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
		Dial: func() (redis.Conn, error) {
			return dial(network, address, password)
		},
	}, keyPairs...)
}

// NewRedisStoreWithDB - like NewRedisStore but accepts `DB` parameter to select
// redis DB instead of using the default one ("0")
func NewRediStoreWithDB(size int, network, address, password, DB string, keyPairs ...[]byte) (*RediStore, error) {
	rs, _ := NewRediStore(size, network, address, password, keyPairs...)
	rs.Pool.Dial = func() (redis.Conn, error) {
		c, err := dial(network, address, password)
		if err != nil {
			return c, err
		}
		if _, err := c.Do("SELECT", DB); err != nil {
			c.Close()
			return nil, err
		}
		return c, err
	}
	_, err := rs.ping()
	return rs, err
}

// NewRediStoreWithPool instantiates a RediStore with a *redis.Pool passed in.
func NewRediStoreWithPool(pool *redis.Pool, keyPairs ...[]byte) (*RediStore, error) {
	rs := &RediStore{
		// http://godoc.org/github.com/garyburd/redigo/redis#Pool
		Pool:   pool,
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: sessionExpire,
		},
		DefaultMaxAge: 60 * 20, // 20 minutes seems like a reasonable default
		maxLength:     4096,
	}
	_, err := rs.ping()
	return rs, err
}

// Close closes the underlying *redis.Pool
func (s *RediStore) Close() error {
	return s.Pool.Close()
}

// Get returns a session for the given name after adding it to the registry.
//
// See gorilla/sessions FilesystemStore.Get().
func (s *RediStore) Get(r *http.Request, name string) (*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See gorilla/sessions FilesystemStore.New().
func (s *RediStore) New(r *http.Request, name string) (*sessions.Session, error) {
	var err error
	session := sessions.NewSession(s, name)
	session.Options = s.Options
	session.IsNew = true
	if c, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, c.Value, &session.ID, s.Codecs...)
		if err == nil {
			ok, err := s.load(session)
			session.IsNew = !(err == nil && ok) // not new if no error and data available
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *RediStore) Save(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
	// Marked for deletion.
	if session.Options.MaxAge < 0 {
		if err := s.delete(session); err != nil {
			return err
		}
		http.SetCookie(w, sessions.NewCookie(session.Name(), "", session.Options))
	} else {
		// Build an alphanumeric key for the redis store.
		if session.ID == "" {
			session.ID = strings.TrimRight(base32.StdEncoding.EncodeToString(securecookie.GenerateRandomKey(32)), "=")
		}
		if err := s.save(session); err != nil {
			return err
		}
		encoded, err := securecookie.EncodeMulti(session.Name(), session.ID, s.Codecs...)
		if err != nil {
			return err
		}
		http.SetCookie(w, sessions.NewCookie(session.Name(), encoded, session.Options))
	}
	return nil
}

// Delete removes the session from redis, and sets the cookie to expire.
//
// WARNING: This method should be considered deprecated since it is not exposed via the gorilla/sessions interface.
// Set session.Options.MaxAge = -1 and call Save instead. - July 18th, 2013
func (s *RediStore) Delete(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
	conn := s.Pool.Get()
	defer conn.Close()
	if _, err := conn.Do("DEL", "session_"+session.ID); err != nil {
		return err
	}
	// Set cookie to expire.
	options := *session.Options
	options.MaxAge = -1
	http.SetCookie(w, sessions.NewCookie(session.Name(), "", &options))
	// Clear session values.
	for k := range session.Values {
		delete(session.Values, k)
	}
	return nil
}

// ping does an internal ping against a server to check if it is alive.
func (s *RediStore) ping() (bool, error) {
	conn := s.Pool.Get()
	defer conn.Close()
	data, err := conn.Do("PING")
	if err != nil || data == nil {
		return false, err
	}
	return (data == "PONG"), nil
}

// save stores the session in redis.
func (s *RediStore) save(session *sessions.Session) error {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(session.Values)
	if err != nil {
		return err
	}
	b := buf.Bytes()
	if s.maxLength != 0 && len(b) > s.maxLength {
		return errors.New("SessionStore: the value to store is too big")
	}
	conn := s.Pool.Get()
	defer conn.Close()
	if err = conn.Err(); err != nil {
		return err
	}
	age := session.Options.MaxAge
	if age == 0 {
		age = s.DefaultMaxAge
	}
	_, err = conn.Do("SETEX", "session_"+session.ID, age, b)
	return err
}

// load reads the session from redis.
// returns true if there is a sessoin data in DB
func (s *RediStore) load(session *sessions.Session) (bool, error) {
	conn := s.Pool.Get()
	defer conn.Close()
	if err := conn.Err(); err != nil {
		return false, err
	}
	data, err := conn.Do("GET", "session_"+session.ID)
	if err != nil {
		return false, err
	}
	if data == nil {
		return false, nil // no data was associated with this key
	}
	b, err := redis.Bytes(data, err)
	if err != nil {
		return false, err
	}
	dec := gob.NewDecoder(bytes.NewBuffer(b))
	return true, dec.Decode(&session.Values)
}

// delete removes keys from redis if MaxAge<0
func (s *RediStore) delete(session *sessions.Session) error {
	conn := s.Pool.Get()
	defer conn.Close()
	if _, err := conn.Do("DEL", "session_"+session.ID); err != nil {
		return err
	}
	return nil
}
<commit_msg>Session wrapper around boj/redistore for Martini sessions.<commit_after>package sessions

import (
	"github.com/boj/redistore"
	"github.com/gorilla/sessions"
)

// RedisStore is an interface that represents a Cookie based storage
// for Sessions.
type RediStore interface {
	// Store is an embedded interface so that RedisStore can be used
	// as a session store.
	Store
	// Options sets the default options for each session stored in this
	// CookieStore.
	Options(Options)
}

// NewCookieStore returns a new CookieStore.
//
// Keys are defined in pairs to allow key rotation, but the common case is to set a single
// authentication key and optionally an encryption key.
//
// The first key in a pair is used for authentication and the second for encryption. The
// encryption key can be set to nil or omitted in the last pair, but the authentication key
// is required in all pairs.
//
// It is recommended to use an authentication key with 32 or 64 bytes. The encryption key,
// if set, must be either 16, 24, or 32 bytes to select AES-128, AES-192, or AES-256 modes.
func NewRediStore(size int, network, address, password string, keyPairs ...[]byte) (RediStore, error) {
	store, err := redistore.NewRediStore(size, network, address, password, keyPairs...)
	if err != nil {
		return nil, err
	}
	return &rediStore{store}, nil
}

type rediStore struct {
	*redistore.RediStore
}

func (c *rediStore) Options(options Options) {
	c.RediStore.Options = &sessions.Options{
		Path:     options.Path,
		Domain:   options.Domain,
		MaxAge:   options.MaxAge,
		Secure:   options.Secure,
		HttpOnly: options.HttpOnly,
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Théo Crevon
//
// See the file LICENSE for copying permission.

/*
Package reflections provides high level abstractions above the
reflect library.

Reflect library is very low-level and as can be quite complex when it comes to do simple things like accessing a structure field value, a field tag...

The purpose of reflections package is to make developers life easier when it comes to introspect structures at runtime.
It's API is freely inspired from python language (getattr, setattr, hasattr...) and provides a simplified access to structure fields and tags.
*/
package reflections

import (
	"errors"
	"fmt"
	"reflect"
)

// GetField returns the value of the provided obj field. obj can whether
// be a structure or pointer to structure.
func GetField(obj interface{}, name string) (interface{}, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	field := objValue.FieldByName(name)
	if !field.IsValid() {
		return nil, fmt.Errorf("No such field: %s in obj", name)
	}

	return field.Interface(), nil
}

// GetFieldKind returns the kind of the provided obj field. obj can whether
// be a structure or pointer to structure.
func GetFieldKind(obj interface{}, name string) (reflect.Kind, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return reflect.Invalid, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	field := objValue.FieldByName(name)

	if !field.IsValid() {
		return reflect.Invalid, fmt.Errorf("No such field: %s in obj", name)
	}

	return field.Type().Kind(), nil
}

// GetFieldTag returns the provided obj field tag value. obj can whether
// be a structure or pointer to structure.
func GetFieldTag(obj interface{}, fieldName, tagKey string) (string, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return "", errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()

	field, ok := objType.FieldByName(fieldName)
	if !ok {
		return "", fmt.Errorf("No such field: %s in obj", fieldName)
	}

	if !isExportableField(field) {
		return "", errors.New("Cannot GetFieldTag on a non-exported struct field")
	}

	return field.Tag.Get(tagKey), nil
}

// SetField sets the provided obj field with provided value. obj param has
// to be a pointer to a struct, otherwise it will soundly fail. Provided
// value type should match with the struct field you're trying to set.
func SetField(obj interface{}, name string, value interface{}) error {
	// Fetch the field reflect.Value
	structValue := reflect.ValueOf(obj).Elem()
	structFieldValue := structValue.FieldByName(name)

	if !structFieldValue.IsValid() {
		return fmt.Errorf("No such field: %s in obj", name)
	}

	// If obj field value is not settable an error is thrown
	if !structFieldValue.CanSet() {
		return fmt.Errorf("Cannot set %s field value", name)
	}

	structFieldType := structFieldValue.Type()
	val := reflect.ValueOf(value)
	if structFieldType != val.Type() {
		invalidTypeError := errors.New("Provided value type didn't match obj field type")
		return invalidTypeError
	}

	structFieldValue.Set(val)
	return nil
}

// HasField checks if the provided field name is part of a struct. obj can whether
// be a structure or pointer to structure.
func HasField(obj interface{}, name string) (bool, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return false, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	field, ok := objType.FieldByName(name)
	if !ok || !isExportableField(field) {
		return false, nil
	}

	return true, nil
}

// Fields returns the struct fields names list. obj can whether
// be a structure or pointer to structure.
func Fields(obj interface{}) ([]string, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	fieldsCount := objType.NumField()

	var fields []string
	for i := 0; i < fieldsCount; i++ {
		field := objType.Field(i)
		if isExportableField(field) {
			fields = append(fields, field.Name)
		}
	}

	return fields, nil
}

// Items returns the field - value struct pairs as a map. obj can whether
// be a structure or pointer to structure.
func Items(obj interface{}) (map[string]interface{}, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	fieldsCount := objType.NumField()

	items := make(map[string]interface{})

	for i := 0; i < fieldsCount; i++ {
		field := objType.Field(i)
		fieldValue := objValue.Field(i)

		// Make sure only exportable and addressable fields are
		// returned by Items
		if isExportableField(field) {
			items[field.Name] = fieldValue.Interface()
		}
	}

	return items, nil
}

// Tags lists the struct tag fields. obj can whether
// be a structure or pointer to structure.
func Tags(obj interface{}, key string) (map[string]string, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	fieldsCount := objType.NumField()

	tags := make(map[string]string)

	for i := 0; i < fieldsCount; i++ {
		structField := objType.Field(i)

		if isExportableField(structField) {
			tags[structField.Name] = structField.Tag.Get(key)
		}
	}

	return tags, nil
}

func reflectValue(obj interface{}) reflect.Value {
    var val reflect.Value

    if reflect.TypeOf(obj).Kind() == reflect.Ptr {
        val = reflect.ValueOf(obj).Elem()
    } else {
        val = reflect.ValueOf(obj)
    }

    return val
}

func isExportableField(field reflect.StructField) bool {
	// PkgPath is empty for exported fields.
	return field.PkgPath == ""
}

func hasValidType(obj interface{}, types []reflect.Kind) bool{
	for _, t := range types {
		if reflect.TypeOf(obj).Kind() == t {
			return true
		}
	}

	return false
}

func isStruct(obj interface{}) bool {
	return reflect.TypeOf(obj).Kind() == reflect.Struct
}

func isPointer(obj interface{}) bool {
	return reflect.TypeOf(obj).Kind() == reflect.Ptr
}
<commit_msg>only return tags we are interested in, not all tags<commit_after>// Copyright (c) 2013 Théo Crevon
//
// See the file LICENSE for copying permission.

/*
Package reflections provides high level abstractions above the
reflect library.

Reflect library is very low-level and as can be quite complex when it comes to do simple things like accessing a structure field value, a field tag...

The purpose of reflections package is to make developers life easier when it comes to introspect structures at runtime.
It's API is freely inspired from python language (getattr, setattr, hasattr...) and provides a simplified access to structure fields and tags.
*/
package reflections

import (
	"errors"
	"fmt"
	"reflect"
)

// GetField returns the value of the provided obj field. obj can whether
// be a structure or pointer to structure.
func GetField(obj interface{}, name string) (interface{}, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	field := objValue.FieldByName(name)
	if !field.IsValid() {
		return nil, fmt.Errorf("No such field: %s in obj", name)
	}

	return field.Interface(), nil
}

// GetFieldKind returns the kind of the provided obj field. obj can whether
// be a structure or pointer to structure.
func GetFieldKind(obj interface{}, name string) (reflect.Kind, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return reflect.Invalid, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	field := objValue.FieldByName(name)

	if !field.IsValid() {
		return reflect.Invalid, fmt.Errorf("No such field: %s in obj", name)
	}

	return field.Type().Kind(), nil
}

// GetFieldTag returns the provided obj field tag value. obj can whether
// be a structure or pointer to structure.
func GetFieldTag(obj interface{}, fieldName, tagKey string) (string, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return "", errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()

	field, ok := objType.FieldByName(fieldName)
	if !ok {
		return "", fmt.Errorf("No such field: %s in obj", fieldName)
	}

	if !isExportableField(field) {
		return "", errors.New("Cannot GetFieldTag on a non-exported struct field")
	}

	return field.Tag.Get(tagKey), nil
}

// SetField sets the provided obj field with provided value. obj param has
// to be a pointer to a struct, otherwise it will soundly fail. Provided
// value type should match with the struct field you're trying to set.
func SetField(obj interface{}, name string, value interface{}) error {
	// Fetch the field reflect.Value
	structValue := reflect.ValueOf(obj).Elem()
	structFieldValue := structValue.FieldByName(name)

	if !structFieldValue.IsValid() {
		return fmt.Errorf("No such field: %s in obj", name)
	}

	// If obj field value is not settable an error is thrown
	if !structFieldValue.CanSet() {
		return fmt.Errorf("Cannot set %s field value", name)
	}

	structFieldType := structFieldValue.Type()
	val := reflect.ValueOf(value)
	if structFieldType != val.Type() {
		invalidTypeError := errors.New("Provided value type didn't match obj field type")
		return invalidTypeError
	}

	structFieldValue.Set(val)
	return nil
}

// HasField checks if the provided field name is part of a struct. obj can whether
// be a structure or pointer to structure.
func HasField(obj interface{}, name string) (bool, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return false, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	field, ok := objType.FieldByName(name)
	if !ok || !isExportableField(field) {
		return false, nil
	}

	return true, nil
}

// Fields returns the struct fields names list. obj can whether
// be a structure or pointer to structure.
func Fields(obj interface{}) ([]string, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	fieldsCount := objType.NumField()

	var fields []string
	for i := 0; i < fieldsCount; i++ {
		field := objType.Field(i)
		if isExportableField(field) {
			fields = append(fields, field.Name)
		}
	}

	return fields, nil
}

// Items returns the field - value struct pairs as a map. obj can whether
// be a structure or pointer to structure.
func Items(obj interface{}) (map[string]interface{}, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	fieldsCount := objType.NumField()

	items := make(map[string]interface{})

	for i := 0; i < fieldsCount; i++ {
		field := objType.Field(i)
		fieldValue := objValue.Field(i)

		// Make sure only exportable and addressable fields are
		// returned by Items
		if isExportableField(field) {
			items[field.Name] = fieldValue.Interface()
		}
	}

	return items, nil
}

// Tags lists the struct tag fields. obj can whether
// be a structure or pointer to structure.
func Tags(obj interface{}, key string) (map[string]string, error) {
	if !hasValidType(obj, []reflect.Kind{reflect.Struct, reflect.Ptr}) {
		return nil, errors.New("Cannot use GetField on a non-struct interface")
	}

	objValue := reflectValue(obj)
	objType := objValue.Type()
	fieldsCount := objType.NumField()

	tags := make(map[string]string)

	for i := 0; i < fieldsCount; i++ {
		structField := objType.Field(i)

		if isExportableField(structField) {
			val := structField.Tag.Get(key)
			if val != "" {
				tags[structField.Name] = structField.Tag.Get(key)
			}
		}
	}

	return tags, nil
}

func reflectValue(obj interface{}) reflect.Value {
	var val reflect.Value

	if reflect.TypeOf(obj).Kind() == reflect.Ptr {
		val = reflect.ValueOf(obj).Elem()
	} else {
		val = reflect.ValueOf(obj)
	}

	return val
}

func isExportableField(field reflect.StructField) bool {
	// PkgPath is empty for exported fields.
	return field.PkgPath == ""
}

func hasValidType(obj interface{}, types []reflect.Kind) bool {
	for _, t := range types {
		if reflect.TypeOf(obj).Kind() == t {
			return true
		}
	}

	return false
}

func isStruct(obj interface{}) bool {
	return reflect.TypeOf(obj).Kind() == reflect.Struct
}

func isPointer(obj interface{}) bool {
	return reflect.TypeOf(obj).Kind() == reflect.Ptr
}
<|endoftext|>


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

package ec2

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"

	"github.com/aws/aws-sdk-go/service/ec2"
	config_util "github.com/prometheus/common/config"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	"github.com/prometheus/prometheus/util/strutil"
	yaml_util "github.com/prometheus/prometheus/util/yaml"
)

const (
	ec2Label              = model.MetaLabelPrefix + "ec2_"
	ec2LabelAZ            = ec2Label + "availability_zone"
	ec2LabelInstanceID    = ec2Label + "instance_id"
	ec2LabelInstanceState = ec2Label + "instance_state"
	ec2LabelInstanceType  = ec2Label + "instance_type"
	ec2LabelPublicDNS     = ec2Label + "public_dns_name"
	ec2LabelPublicIP      = ec2Label + "public_ip"
	ec2LabelPrivateIP     = ec2Label + "private_ip"
	ec2LabelSubnetID      = ec2Label + "subnet_id"
	ec2LabelTag           = ec2Label + "tag_"
	ec2LabelVPCID         = ec2Label + "vpc_id"
	subnetSeparator       = ","
)

var (
	ec2SDRefreshFailuresCount = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "prometheus_sd_ec2_refresh_failures_total",
			Help: "The number of EC2-SD scrape failures.",
		})
	ec2SDRefreshDuration = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Name: "prometheus_sd_ec2_refresh_duration_seconds",
			Help: "The duration of a EC2-SD refresh in seconds.",
		})
	// DefaultSDConfig is the default EC2 SD configuration.
	DefaultSDConfig = SDConfig{
		Port:            80,
		RefreshInterval: model.Duration(60 * time.Second),
	}
)

// SDConfig is the configuration for EC2 based service discovery.
type SDConfig struct {
	Region          string             `yaml:"region"`
	AccessKey       string             `yaml:"access_key,omitempty"`
	SecretKey       config_util.Secret `yaml:"secret_key,omitempty"`
	Profile         string             `yaml:"profile,omitempty"`
	RoleARN         string             `yaml:"role_arn,omitempty"`
	RefreshInterval model.Duration     `yaml:"refresh_interval,omitempty"`
	Port            int                `yaml:"port"`

	// Catches all undefined fields and must be empty after parsing.
	XXX map[string]interface{} `yaml:",inline"`
}

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (c *SDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
	*c = DefaultSDConfig
	type plain SDConfig
	err := unmarshal((*plain)(c))
	if err != nil {
		return err
	}
	if err := yaml_util.CheckOverflow(c.XXX, "ec2_sd_config"); err != nil {
		return err
	}
	if c.Region == "" {
		sess, err := session.NewSession()
		if err != nil {
			return err
		}
		metadata := ec2metadata.New(sess)
		region, err := metadata.Region()
		if err != nil {
			return fmt.Errorf("EC2 SD configuration requires a region")
		}
		c.Region = region
	}
	return nil
}

func init() {
	prometheus.MustRegister(ec2SDRefreshFailuresCount)
	prometheus.MustRegister(ec2SDRefreshDuration)
}

// Discovery periodically performs EC2-SD requests. It implements
// the Discoverer interface.
type Discovery struct {
	aws      *aws.Config
	interval time.Duration
	profile  string
	roleARN  string
	port     int
	logger   log.Logger
}

// NewDiscovery returns a new EC2Discovery which periodically refreshes its targets.
func NewDiscovery(conf *SDConfig, logger log.Logger) *Discovery {
	creds := credentials.NewStaticCredentials(conf.AccessKey, string(conf.SecretKey), "")
	if conf.AccessKey == "" && conf.SecretKey == "" {
		creds = nil
	}
	if logger == nil {
		logger = log.NewNopLogger()
	}
	return &Discovery{
		aws: &aws.Config{
			Region:      &conf.Region,
			Credentials: creds,
		},
		profile:  conf.Profile,
		roleARN:  conf.RoleARN,
		interval: time.Duration(conf.RefreshInterval),
		port:     conf.Port,
		logger:   logger,
	}
}

// Run implements the Discoverer interface.
func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) {
	ticker := time.NewTicker(d.interval)
	defer ticker.Stop()

	// Get an initial set right away.
	tg, err := d.refresh()
	if err != nil {
		level.Error(d.logger).Log("msg", "Refresh failed", "err", err)
	} else {
		select {
		case ch <- []*targetgroup.Group{tg}:
		case <-ctx.Done():
			return
		}
	}

	for {
		select {
		case <-ticker.C:
			tg, err := d.refresh()
			if err != nil {
				level.Error(d.logger).Log("msg", "Refresh failed", "err", err)
				continue
			}

			select {
			case ch <- []*targetgroup.Group{tg}:
			case <-ctx.Done():
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (d *Discovery) ec2MetadataAvailable(sess *session.Session) (isAvailable bool) {
	svc := ec2metadata.New(sess, &aws.Config{
		MaxRetries: aws.Int(0),
	})

	isAvailable = svc.Available()

	return isAvailable
}

func (d *Discovery) refresh() (tg *targetgroup.Group, err error) {
	t0 := time.Now()
	defer func() {
		ec2SDRefreshDuration.Observe(time.Since(t0).Seconds())
		if err != nil {
			ec2SDRefreshFailuresCount.Inc()
		}
	}()

	sess, err := session.NewSessionWithOptions(session.Options{
		Config:  *d.aws,
		Profile: d.profile,
	})
	if err != nil {
		return nil, fmt.Errorf("could not create aws session: %s", err)
	}

	var ec2s *ec2.EC2
	if d.roleARN != "" {
		creds := stscreds.NewCredentials(sess, d.roleARN)
		ec2s = ec2.New(sess, &aws.Config{Credentials: creds})
	} else {
		if d.aws.Credentials == nil && d.ec2MetadataAvailable(sess) {
			creds := ec2rolecreds.NewCredentials(sess)
			ec2s = ec2.New(sess, &aws.Config{Credentials: creds})
		} else {
			ec2s = ec2.New(sess)
		}
	}
	tg = &targetgroup.Group{
		Source: *d.aws.Region,
	}
	if err = ec2s.DescribeInstancesPages(nil, func(p *ec2.DescribeInstancesOutput, lastPage bool) bool {
		for _, r := range p.Reservations {
			for _, inst := range r.Instances {
				if inst.PrivateIpAddress == nil {
					continue
				}
				labels := model.LabelSet{
					ec2LabelInstanceID: model.LabelValue(*inst.InstanceId),
				}
				labels[ec2LabelPrivateIP] = model.LabelValue(*inst.PrivateIpAddress)
				addr := net.JoinHostPort(*inst.PrivateIpAddress, fmt.Sprintf("%d", d.port))
				labels[model.AddressLabel] = model.LabelValue(addr)

				if inst.PublicIpAddress != nil {
					labels[ec2LabelPublicIP] = model.LabelValue(*inst.PublicIpAddress)
					labels[ec2LabelPublicDNS] = model.LabelValue(*inst.PublicDnsName)
				}

				labels[ec2LabelAZ] = model.LabelValue(*inst.Placement.AvailabilityZone)
				labels[ec2LabelInstanceState] = model.LabelValue(*inst.State.Name)
				labels[ec2LabelInstanceType] = model.LabelValue(*inst.InstanceType)

				if inst.VpcId != nil {
					labels[ec2LabelVPCID] = model.LabelValue(*inst.VpcId)

					subnetsMap := make(map[string]struct{})
					for _, eni := range inst.NetworkInterfaces {
						subnetsMap[*eni.SubnetId] = struct{}{}
					}
					subnets := []string{}
					for k := range subnetsMap {
						subnets = append(subnets, k)
					}
					labels[ec2LabelSubnetID] = model.LabelValue(
						subnetSeparator +
							strings.Join(subnets, subnetSeparator) +
							subnetSeparator)
				}

				for _, t := range inst.Tags {
					if t == nil || t.Key == nil || t.Value == nil {
						continue
					}
					name := strutil.SanitizeLabelName(*t.Key)
					labels[ec2LabelTag+model.LabelName(name)] = model.LabelValue(*t.Value)
				}
				tg.Targets = append(tg.Targets, labels)
			}
		}
		return true
	}); err != nil {
		return nil, fmt.Errorf("could not describe instances: %s", err)
	}
	return tg, nil
}
<commit_msg>Revert "Feature: Allow getting credentials via EC2 role (#3343)" (#3985)<commit_after>// Copyright 2015 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ec2

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"

	"github.com/aws/aws-sdk-go/service/ec2"
	config_util "github.com/prometheus/common/config"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	"github.com/prometheus/prometheus/util/strutil"
	yaml_util "github.com/prometheus/prometheus/util/yaml"
)

const (
	ec2Label              = model.MetaLabelPrefix + "ec2_"
	ec2LabelAZ            = ec2Label + "availability_zone"
	ec2LabelInstanceID    = ec2Label + "instance_id"
	ec2LabelInstanceState = ec2Label + "instance_state"
	ec2LabelInstanceType  = ec2Label + "instance_type"
	ec2LabelPublicDNS     = ec2Label + "public_dns_name"
	ec2LabelPublicIP      = ec2Label + "public_ip"
	ec2LabelPrivateIP     = ec2Label + "private_ip"
	ec2LabelSubnetID      = ec2Label + "subnet_id"
	ec2LabelTag           = ec2Label + "tag_"
	ec2LabelVPCID         = ec2Label + "vpc_id"
	subnetSeparator       = ","
)

var (
	ec2SDRefreshFailuresCount = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "prometheus_sd_ec2_refresh_failures_total",
			Help: "The number of EC2-SD scrape failures.",
		})
	ec2SDRefreshDuration = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Name: "prometheus_sd_ec2_refresh_duration_seconds",
			Help: "The duration of a EC2-SD refresh in seconds.",
		})
	// DefaultSDConfig is the default EC2 SD configuration.
	DefaultSDConfig = SDConfig{
		Port:            80,
		RefreshInterval: model.Duration(60 * time.Second),
	}
)

// SDConfig is the configuration for EC2 based service discovery.
type SDConfig struct {
	Region          string             `yaml:"region"`
	AccessKey       string             `yaml:"access_key,omitempty"`
	SecretKey       config_util.Secret `yaml:"secret_key,omitempty"`
	Profile         string             `yaml:"profile,omitempty"`
	RoleARN         string             `yaml:"role_arn,omitempty"`
	RefreshInterval model.Duration     `yaml:"refresh_interval,omitempty"`
	Port            int                `yaml:"port"`

	// Catches all undefined fields and must be empty after parsing.
	XXX map[string]interface{} `yaml:",inline"`
}

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (c *SDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
	*c = DefaultSDConfig
	type plain SDConfig
	err := unmarshal((*plain)(c))
	if err != nil {
		return err
	}
	if err := yaml_util.CheckOverflow(c.XXX, "ec2_sd_config"); err != nil {
		return err
	}
	if c.Region == "" {
		sess, err := session.NewSession()
		if err != nil {
			return err
		}
		metadata := ec2metadata.New(sess)
		region, err := metadata.Region()
		if err != nil {
			return fmt.Errorf("EC2 SD configuration requires a region")
		}
		c.Region = region
	}
	return nil
}

func init() {
	prometheus.MustRegister(ec2SDRefreshFailuresCount)
	prometheus.MustRegister(ec2SDRefreshDuration)
}

// Discovery periodically performs EC2-SD requests. It implements
// the Discoverer interface.
type Discovery struct {
	aws      *aws.Config
	interval time.Duration
	profile  string
	roleARN  string
	port     int
	logger   log.Logger
}

// NewDiscovery returns a new EC2Discovery which periodically refreshes its targets.
func NewDiscovery(conf *SDConfig, logger log.Logger) *Discovery {
	creds := credentials.NewStaticCredentials(conf.AccessKey, string(conf.SecretKey), "")
	if conf.AccessKey == "" && conf.SecretKey == "" {
		creds = nil
	}
	if logger == nil {
		logger = log.NewNopLogger()
	}
	return &Discovery{
		aws: &aws.Config{
			Region:      &conf.Region,
			Credentials: creds,
		},
		profile:  conf.Profile,
		roleARN:  conf.RoleARN,
		interval: time.Duration(conf.RefreshInterval),
		port:     conf.Port,
		logger:   logger,
	}
}

// Run implements the Discoverer interface.
func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) {
	ticker := time.NewTicker(d.interval)
	defer ticker.Stop()

	// Get an initial set right away.
	tg, err := d.refresh()
	if err != nil {
		level.Error(d.logger).Log("msg", "Refresh failed", "err", err)
	} else {
		select {
		case ch <- []*targetgroup.Group{tg}:
		case <-ctx.Done():
			return
		}
	}

	for {
		select {
		case <-ticker.C:
			tg, err := d.refresh()
			if err != nil {
				level.Error(d.logger).Log("msg", "Refresh failed", "err", err)
				continue
			}

			select {
			case ch <- []*targetgroup.Group{tg}:
			case <-ctx.Done():
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (d *Discovery) refresh() (tg *targetgroup.Group, err error) {
	t0 := time.Now()
	defer func() {
		ec2SDRefreshDuration.Observe(time.Since(t0).Seconds())
		if err != nil {
			ec2SDRefreshFailuresCount.Inc()
		}
	}()

	sess, err := session.NewSessionWithOptions(session.Options{
		Config:  *d.aws,
		Profile: d.profile,
	})
	if err != nil {
		return nil, fmt.Errorf("could not create aws session: %s", err)
	}

	var ec2s *ec2.EC2
	if d.roleARN != "" {
		creds := stscreds.NewCredentials(sess, d.roleARN)
		ec2s = ec2.New(sess, &aws.Config{Credentials: creds})
	} else {
		ec2s = ec2.New(sess)
	}
	tg = &targetgroup.Group{
		Source: *d.aws.Region,
	}
	if err = ec2s.DescribeInstancesPages(nil, func(p *ec2.DescribeInstancesOutput, lastPage bool) bool {
		for _, r := range p.Reservations {
			for _, inst := range r.Instances {
				if inst.PrivateIpAddress == nil {
					continue
				}
				labels := model.LabelSet{
					ec2LabelInstanceID: model.LabelValue(*inst.InstanceId),
				}
				labels[ec2LabelPrivateIP] = model.LabelValue(*inst.PrivateIpAddress)
				addr := net.JoinHostPort(*inst.PrivateIpAddress, fmt.Sprintf("%d", d.port))
				labels[model.AddressLabel] = model.LabelValue(addr)

				if inst.PublicIpAddress != nil {
					labels[ec2LabelPublicIP] = model.LabelValue(*inst.PublicIpAddress)
					labels[ec2LabelPublicDNS] = model.LabelValue(*inst.PublicDnsName)
				}

				labels[ec2LabelAZ] = model.LabelValue(*inst.Placement.AvailabilityZone)
				labels[ec2LabelInstanceState] = model.LabelValue(*inst.State.Name)
				labels[ec2LabelInstanceType] = model.LabelValue(*inst.InstanceType)

				if inst.VpcId != nil {
					labels[ec2LabelVPCID] = model.LabelValue(*inst.VpcId)

					subnetsMap := make(map[string]struct{})
					for _, eni := range inst.NetworkInterfaces {
						subnetsMap[*eni.SubnetId] = struct{}{}
					}
					subnets := []string{}
					for k := range subnetsMap {
						subnets = append(subnets, k)
					}
					labels[ec2LabelSubnetID] = model.LabelValue(
						subnetSeparator +
							strings.Join(subnets, subnetSeparator) +
							subnetSeparator)
				}

				for _, t := range inst.Tags {
					if t == nil || t.Key == nil || t.Value == nil {
						continue
					}
					name := strutil.SanitizeLabelName(*t.Key)
					labels[ec2LabelTag+model.LabelName(name)] = model.LabelValue(*t.Value)
				}
				tg.Targets = append(tg.Targets, labels)
			}
		}
		return true
	}); err != nil {
		return nil, fmt.Errorf("could not describe instances: %s", err)
	}
	return tg, nil
}
<|endoftext|>


<commit_before>// +build windows

package disk

import (
	"bytes"
	"context"
	"unsafe"

	"github.com/shirou/gopsutil/internal/common"
	"golang.org/x/sys/windows"
)

var (
	procGetDiskFreeSpaceExW     = common.Modkernel32.NewProc("GetDiskFreeSpaceExW")
	procGetLogicalDriveStringsW = common.Modkernel32.NewProc("GetLogicalDriveStringsW")
	procGetDriveType            = common.Modkernel32.NewProc("GetDriveTypeW")
	provGetVolumeInformation    = common.Modkernel32.NewProc("GetVolumeInformationW")
)

var (
	FileFileCompression = int64(16)     // 0x00000010
	FileReadOnlyVolume  = int64(524288) // 0x00080000
)

type Win32_PerfFormattedData struct {
	Name                    string
	AvgDiskBytesPerRead     uint64
	AvgDiskBytesPerWrite    uint64
	AvgDiskReadQueueLength  uint64
	AvgDiskWriteQueueLength uint64
	AvgDisksecPerRead       uint64
	AvgDisksecPerWrite      uint64
}
type Win32_DiskDrive struct {
	DeviceID         string
	FirmwareRevision string
	InterfaceType    string
	SerialNumber     string
}
type Win32_DiskPartition struct {
	DeviceID string
}

const WaitMSec = 500

func Usage(path string) (*UsageStat, error) {
	return UsageWithContext(context.Background(), path)
}

func UsageWithContext(ctx context.Context, path string) (*UsageStat, error) {
	ret := &UsageStat{}

	lpFreeBytesAvailable := int64(0)
	lpTotalNumberOfBytes := int64(0)
	lpTotalNumberOfFreeBytes := int64(0)
	diskret, _, err := procGetDiskFreeSpaceExW.Call(
		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(path))),
		uintptr(unsafe.Pointer(&lpFreeBytesAvailable)),
		uintptr(unsafe.Pointer(&lpTotalNumberOfBytes)),
		uintptr(unsafe.Pointer(&lpTotalNumberOfFreeBytes)))
	if diskret == 0 {
		return nil, err
	}
	ret = &UsageStat{
		Path:        path,
		Total:       uint64(lpTotalNumberOfBytes),
		Free:        uint64(lpTotalNumberOfFreeBytes),
		Used:        uint64(lpTotalNumberOfBytes) - uint64(lpTotalNumberOfFreeBytes),
		UsedPercent: (float64(lpTotalNumberOfBytes) - float64(lpTotalNumberOfFreeBytes)) / float64(lpTotalNumberOfBytes) * 100,
		// InodesTotal: 0,
		// InodesFree: 0,
		// InodesUsed: 0,
		// InodesUsedPercent: 0,
	}
	return ret, nil
}

func Partitions(all bool) ([]PartitionStat, error) {
	return PartitionsWithContext(context.Background(), all)
}

func PartitionsWithContext(ctx context.Context, all bool) ([]PartitionStat, error) {
	var ret []PartitionStat
	lpBuffer := make([]byte, 254)
	diskret, _, err := procGetLogicalDriveStringsW.Call(
		uintptr(len(lpBuffer)),
		uintptr(unsafe.Pointer(&lpBuffer[0])))
	if diskret == 0 {
		return ret, err
	}
	for _, v := range lpBuffer {
		if v >= 65 && v <= 90 {
			path := string(v) + ":"
			if path == "A:" || path == "B:" { // skip floppy drives
				continue
			}
			typepath, _ := windows.UTF16PtrFromString(path)
			typeret, _, _ := procGetDriveType.Call(uintptr(unsafe.Pointer(typepath)))
			if typeret == 0 {
				return ret, windows.GetLastError()
			}
			// 2: DRIVE_REMOVABLE 3: DRIVE_FIXED 4: DRIVE_REMOTE 5: DRIVE_CDROM

			if typeret == 2 || typeret == 3 || typeret == 4 || typeret == 5 {
				lpVolumeNameBuffer := make([]byte, 256)
				lpVolumeSerialNumber := int64(0)
				lpMaximumComponentLength := int64(0)
				lpFileSystemFlags := int64(0)
				lpFileSystemNameBuffer := make([]byte, 256)
				volpath, _ := windows.UTF16PtrFromString(string(v) + ":/")
				driveret, _, err := provGetVolumeInformation.Call(
					uintptr(unsafe.Pointer(volpath)),
					uintptr(unsafe.Pointer(&lpVolumeNameBuffer[0])),
					uintptr(len(lpVolumeNameBuffer)),
					uintptr(unsafe.Pointer(&lpVolumeSerialNumber)),
					uintptr(unsafe.Pointer(&lpMaximumComponentLength)),
					uintptr(unsafe.Pointer(&lpFileSystemFlags)),
					uintptr(unsafe.Pointer(&lpFileSystemNameBuffer[0])),
					uintptr(len(lpFileSystemNameBuffer)))
				if driveret == 0 {
					if typeret == 5 || typeret == 2 {
						continue //device is not ready will happen if there is no disk in the drive
					}
					return ret, err
				}
				opts := "rw"
				if lpFileSystemFlags&FileReadOnlyVolume != 0 {
					opts = "ro"
				}
				if lpFileSystemFlags&FileFileCompression != 0 {
					opts += ".compress"
				}

				d := PartitionStat{
					Mountpoint: path,
					Device:     path,
					Fstype:     string(bytes.Replace(lpFileSystemNameBuffer, []byte("\x00"), []byte(""), -1)),
					Opts:       opts,
				}
				ret = append(ret, d)
			}
		}
	}
	return ret, nil
}

func IOCounters(names ...string) (map[string]IOCountersStat, error) {
	return IOCountersWithContext(context.Background(), names...)
}

func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOCountersStat, error) {
	ret := make(map[string]IOCountersStat, 0)
	var dst []Win32_PerfFormattedData

	err := common.WMIQueryWithContext(ctx, "SELECT * FROM Win32_PerfFormattedData_PerfDisk_LogicalDisk", &dst)
	if err != nil {
		return ret, err
	}
	for _, d := range dst {
		if len(d.Name) > 3 { // not get _Total or Harddrive
			continue
		}

		if len(names) > 0 && !common.StringsHas(names, d.Name) {
			continue
		}

		tmpIO := IOCountersStat{
			Name:       d.Name,
			ReadCount:  uint64(d.AvgDiskReadQueueLength),
			WriteCount: d.AvgDiskWriteQueueLength,
			ReadBytes:  uint64(d.AvgDiskBytesPerRead),
			WriteBytes: uint64(d.AvgDiskBytesPerWrite),
			ReadTime:   d.AvgDisksecPerRead,
			WriteTime:  d.AvgDisksecPerWrite,
		}
		tmpIO.SerialNumber = GetDiskSerialNumber(d.Name)
		ret[d.Name] = d
	}
	return ret, nil
}

// return disk serial number(not volume serial number) of given device or empty string on error. Name of device is drive letter, eg. C:
func GetDiskSerialNumber(name string) string {
	return GetDiskSerialNumberWithContext(context.Background(), name)
}

func GetDiskSerialNumberWithContext(ctx context.Context, name string) string {
	var diskPart []Win32_DiskPartition
	var diskDrive []Win32_DiskDrive
	err := common.WMIQueryWithContext(ctx, "Associators of {Win32_LogicalDisk.DeviceID='"+name+"'} where AssocClass=Win32_LogicalDiskToPartition", &diskPart)
	if err != nil || len(diskPart) <= 0 {
		return ""
	}
	err = common.WMIQueryWithContext(ctx, "Associators of {Win32_DiskPartition.DeviceID='"+diskPart[0].DeviceID+"'} where AssocClass=Win32_DiskDriveToDiskPartition", &diskDrive)
	if err != nil || len(diskDrive) <= 0 {
		return ""
	}
	return diskDrive[0].SerialNumber
}
<commit_msg>wrong variable<commit_after>// +build windows

package disk

import (
	"bytes"
	"context"
	"unsafe"

	"github.com/shirou/gopsutil/internal/common"
	"golang.org/x/sys/windows"
)

var (
	procGetDiskFreeSpaceExW     = common.Modkernel32.NewProc("GetDiskFreeSpaceExW")
	procGetLogicalDriveStringsW = common.Modkernel32.NewProc("GetLogicalDriveStringsW")
	procGetDriveType            = common.Modkernel32.NewProc("GetDriveTypeW")
	provGetVolumeInformation    = common.Modkernel32.NewProc("GetVolumeInformationW")
)

var (
	FileFileCompression = int64(16)     // 0x00000010
	FileReadOnlyVolume  = int64(524288) // 0x00080000
)

type Win32_PerfFormattedData struct {
	Name                    string
	AvgDiskBytesPerRead     uint64
	AvgDiskBytesPerWrite    uint64
	AvgDiskReadQueueLength  uint64
	AvgDiskWriteQueueLength uint64
	AvgDisksecPerRead       uint64
	AvgDisksecPerWrite      uint64
}
type Win32_DiskDrive struct {
	DeviceID         string
	FirmwareRevision string
	InterfaceType    string
	SerialNumber     string
}
type Win32_DiskPartition struct {
	DeviceID string
}

const WaitMSec = 500

func Usage(path string) (*UsageStat, error) {
	return UsageWithContext(context.Background(), path)
}

func UsageWithContext(ctx context.Context, path string) (*UsageStat, error) {
	ret := &UsageStat{}

	lpFreeBytesAvailable := int64(0)
	lpTotalNumberOfBytes := int64(0)
	lpTotalNumberOfFreeBytes := int64(0)
	diskret, _, err := procGetDiskFreeSpaceExW.Call(
		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(path))),
		uintptr(unsafe.Pointer(&lpFreeBytesAvailable)),
		uintptr(unsafe.Pointer(&lpTotalNumberOfBytes)),
		uintptr(unsafe.Pointer(&lpTotalNumberOfFreeBytes)))
	if diskret == 0 {
		return nil, err
	}
	ret = &UsageStat{
		Path:        path,
		Total:       uint64(lpTotalNumberOfBytes),
		Free:        uint64(lpTotalNumberOfFreeBytes),
		Used:        uint64(lpTotalNumberOfBytes) - uint64(lpTotalNumberOfFreeBytes),
		UsedPercent: (float64(lpTotalNumberOfBytes) - float64(lpTotalNumberOfFreeBytes)) / float64(lpTotalNumberOfBytes) * 100,
		// InodesTotal: 0,
		// InodesFree: 0,
		// InodesUsed: 0,
		// InodesUsedPercent: 0,
	}
	return ret, nil
}

func Partitions(all bool) ([]PartitionStat, error) {
	return PartitionsWithContext(context.Background(), all)
}

func PartitionsWithContext(ctx context.Context, all bool) ([]PartitionStat, error) {
	var ret []PartitionStat
	lpBuffer := make([]byte, 254)
	diskret, _, err := procGetLogicalDriveStringsW.Call(
		uintptr(len(lpBuffer)),
		uintptr(unsafe.Pointer(&lpBuffer[0])))
	if diskret == 0 {
		return ret, err
	}
	for _, v := range lpBuffer {
		if v >= 65 && v <= 90 {
			path := string(v) + ":"
			if path == "A:" || path == "B:" { // skip floppy drives
				continue
			}
			typepath, _ := windows.UTF16PtrFromString(path)
			typeret, _, _ := procGetDriveType.Call(uintptr(unsafe.Pointer(typepath)))
			if typeret == 0 {
				return ret, windows.GetLastError()
			}
			// 2: DRIVE_REMOVABLE 3: DRIVE_FIXED 4: DRIVE_REMOTE 5: DRIVE_CDROM

			if typeret == 2 || typeret == 3 || typeret == 4 || typeret == 5 {
				lpVolumeNameBuffer := make([]byte, 256)
				lpVolumeSerialNumber := int64(0)
				lpMaximumComponentLength := int64(0)
				lpFileSystemFlags := int64(0)
				lpFileSystemNameBuffer := make([]byte, 256)
				volpath, _ := windows.UTF16PtrFromString(string(v) + ":/")
				driveret, _, err := provGetVolumeInformation.Call(
					uintptr(unsafe.Pointer(volpath)),
					uintptr(unsafe.Pointer(&lpVolumeNameBuffer[0])),
					uintptr(len(lpVolumeNameBuffer)),
					uintptr(unsafe.Pointer(&lpVolumeSerialNumber)),
					uintptr(unsafe.Pointer(&lpMaximumComponentLength)),
					uintptr(unsafe.Pointer(&lpFileSystemFlags)),
					uintptr(unsafe.Pointer(&lpFileSystemNameBuffer[0])),
					uintptr(len(lpFileSystemNameBuffer)))
				if driveret == 0 {
					if typeret == 5 || typeret == 2 {
						continue //device is not ready will happen if there is no disk in the drive
					}
					return ret, err
				}
				opts := "rw"
				if lpFileSystemFlags&FileReadOnlyVolume != 0 {
					opts = "ro"
				}
				if lpFileSystemFlags&FileFileCompression != 0 {
					opts += ".compress"
				}

				d := PartitionStat{
					Mountpoint: path,
					Device:     path,
					Fstype:     string(bytes.Replace(lpFileSystemNameBuffer, []byte("\x00"), []byte(""), -1)),
					Opts:       opts,
				}
				ret = append(ret, d)
			}
		}
	}
	return ret, nil
}

func IOCounters(names ...string) (map[string]IOCountersStat, error) {
	return IOCountersWithContext(context.Background(), names...)
}

func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOCountersStat, error) {
	ret := make(map[string]IOCountersStat, 0)
	var dst []Win32_PerfFormattedData

	err := common.WMIQueryWithContext(ctx, "SELECT * FROM Win32_PerfFormattedData_PerfDisk_LogicalDisk", &dst)
	if err != nil {
		return ret, err
	}
	for _, d := range dst {
		if len(d.Name) > 3 { // not get _Total or Harddrive
			continue
		}

		if len(names) > 0 && !common.StringsHas(names, d.Name) {
			continue
		}

		tmpIO := IOCountersStat{
			Name:       d.Name,
			ReadCount:  uint64(d.AvgDiskReadQueueLength),
			WriteCount: d.AvgDiskWriteQueueLength,
			ReadBytes:  uint64(d.AvgDiskBytesPerRead),
			WriteBytes: uint64(d.AvgDiskBytesPerWrite),
			ReadTime:   d.AvgDisksecPerRead,
			WriteTime:  d.AvgDisksecPerWrite,
		}
		tmpIO.SerialNumber = GetDiskSerialNumber(d.Name)
		ret[d.Name] = tmpIO
	}
	return ret, nil
}

// return disk serial number(not volume serial number) of given device or empty string on error. Name of device is drive letter, eg. C:
func GetDiskSerialNumber(name string) string {
	return GetDiskSerialNumberWithContext(context.Background(), name)
}

func GetDiskSerialNumberWithContext(ctx context.Context, name string) string {
	var diskPart []Win32_DiskPartition
	var diskDrive []Win32_DiskDrive
	err := common.WMIQueryWithContext(ctx, "Associators of {Win32_LogicalDisk.DeviceID='"+name+"'} where AssocClass=Win32_LogicalDiskToPartition", &diskPart)
	if err != nil || len(diskPart) <= 0 {
		return ""
	}
	err = common.WMIQueryWithContext(ctx, "Associators of {Win32_DiskPartition.DeviceID='"+diskPart[0].DeviceID+"'} where AssocClass=Win32_DiskDriveToDiskPartition", &diskDrive)
	if err != nil || len(diskDrive) <= 0 {
		return ""
	}
	return diskDrive[0].SerialNumber
}
<|endoftext|>


<commit_before>// dispatch package is used to dispatch irc messages to event handlers in an
// asynchronous fashion. It supports various event handling types to easily
// extract information from events, as well as define more succint handlers.
package dispatch

import (
	"errors"
	"github.com/aarondl/ultimateq/irc"
	"math/rand"
	"strings"
)

var (
	// errProtoCapsMissing is returned by CreateRichDispatch if nil is provided
	// instead of a irc.ProtoCaps pointer.
	errProtoCapsMissing = errors.New(
		"dispatch: Protocaps missing in CreateRichDispatch")
)

// EventHandler is the basic interface that will deal with handling any message
// as a raw IrcMessage event. However there are other message types are specific
// to very common irc events that are more helpful than this interface.
type EventHandler interface {
	HandleRaw(event *irc.IrcMessage)
}

// eventTable is the storage used to keep id -> interface{} mappings in the
// eventTableStore map.
type eventTable map[int]interface{}

// eventTableStore is the map used to hold the event handlers for an event
type eventTableStore map[string]eventTable

// Dispatcher is made for handling bot-local dispatching of irc
// events.
type Dispatcher struct {
	events eventTableStore
	caps   *irc.ProtoCaps
	finder *irc.ChannelFinder
}

// CreateDispatcher initializes an empty dispatcher ready to register events.
func CreateDispatcher() *Dispatcher {
	return &Dispatcher{
		events: make(eventTableStore),
	}
}

// CreateRichDispatcher initializes empty dispatcher ready to register events
// and additionally creates a channelfinder from a set of irc.ProtoCaps in order to
// properly send Privmsg(User|Channel)/Notice(User|Channel) events.
func CreateRichDispatcher(caps *irc.ProtoCaps) (*Dispatcher, error) {
	if caps == nil {
		return nil, errProtoCapsMissing
	}
	f, err := irc.CreateChannelFinder(caps.Chantypes)
	if err != nil {
		return nil, err
	}
	return &Dispatcher{
		events: make(eventTableStore),
		caps:   caps,
		finder: f,
	}, nil
}

// Register registers an event handler to a particular event. In return a
// unique identifer is given to later pass into Unregister in case of a need
// to unregister the event handler.
func (d *Dispatcher) Register(event string, handler interface{}) int {
	event = strings.ToUpper(event)
	id := rand.Int()
	if ev, ok := d.events[event]; !ok {
		d.events[event] = make(eventTable)
	} else {
		for _, has := ev[id]; has; id = rand.Int() {
		}
	}

	d.events[event][id] = handler
	return id
}

// Unregister uses the event name, and the identifier returned by Register to
// unregister a callback from the Dispatcher. If the callback was removed it
// returns true, false if it could not be found.
func (d *Dispatcher) Unregister(event string, id int) bool {
	event = strings.ToUpper(event)
	if ev, ok := d.events[event]; ok {
		if _, ok := ev[id]; ok {
			delete(ev, id)
			return true
		}
	}
	return false
}

// Dispatch an IrcMessage to event handlers handling event also ensures all raw
// handlers receive all messages. Returns false if no eventtable was found for
// the primary sent event.
func (d *Dispatcher) Dispatch(event string, msg *irc.IrcMessage) bool {
	event = strings.ToUpper(event)
	handled := d.dispatchHelper(event, msg)
	d.dispatchHelper(irc.RAW, msg)

	return handled
}

// dispatchHelper locates a handler and attempts to resolve it with
// resolveHandler. It returns true if it was able to find an event table.
func (d *Dispatcher) dispatchHelper(event string, msg *irc.IrcMessage) bool {
	if evtable, ok := d.events[event]; ok {
		for _, handler := range evtable {
			go d.resolveHandler(handler, event, msg)
		}
		return true
	}
	return false
}

// resolveHandler checks the type of the handler passed in, resolves it to a
// real type, coerces the IrcMessage in whatever way necessary and then
// calls that handlers primary dispatch method with the coerced message.
func (d *Dispatcher) resolveHandler(
	handler interface{}, event string, msg *irc.IrcMessage) {

	switch t := handler.(type) {
	case PrivmsgUserHandler:
		if d.finder != nil && !d.finder.IsChannel(msg.Args[0]) {
			t.PrivmsgUser(&Message{msg})
		}
	case PrivmsgChannelHandler:
		if d.finder != nil && d.finder.IsChannel(msg.Args[0]) {
			t.PrivmsgChannel(&Message{msg})
		}
	case PrivmsgHandler:
		t.Privmsg(&Message{msg})
	case NoticeUserHandler:
		if d.finder != nil && !d.finder.IsChannel(msg.Args[0]) {
			t.NoticeUser(&Message{msg})
		}
	case NoticeChannelHandler:
		if d.finder != nil && d.finder.IsChannel(msg.Args[0]) {
			t.NoticeChannel(&Message{msg})
		}
	case NoticeHandler:
		t.Notice(&Message{msg})
	case EventHandler:
		t.HandleRaw(msg)
	}
}
<commit_msg>Fixes a comment over 80 chars.<commit_after>// dispatch package is used to dispatch irc messages to event handlers in an
// asynchronous fashion. It supports various event handling types to easily
// extract information from events, as well as define more succint handlers.
package dispatch

import (
	"errors"
	"github.com/aarondl/ultimateq/irc"
	"math/rand"
	"strings"
)

var (
	// errProtoCapsMissing is returned by CreateRichDispatch if nil is provided
	// instead of a irc.ProtoCaps pointer.
	errProtoCapsMissing = errors.New(
		"dispatch: Protocaps missing in CreateRichDispatch")
)

// EventHandler is the basic interface that will deal with handling any message
// as a raw IrcMessage event. However there are other message types are specific
// to very common irc events that are more helpful than this interface.
type EventHandler interface {
	HandleRaw(event *irc.IrcMessage)
}

// eventTable is the storage used to keep id -> interface{} mappings in the
// eventTableStore map.
type eventTable map[int]interface{}

// eventTableStore is the map used to hold the event handlers for an event
type eventTableStore map[string]eventTable

// Dispatcher is made for handling bot-local dispatching of irc
// events.
type Dispatcher struct {
	events eventTableStore
	caps   *irc.ProtoCaps
	finder *irc.ChannelFinder
}

// CreateDispatcher initializes an empty dispatcher ready to register events.
func CreateDispatcher() *Dispatcher {
	return &Dispatcher{
		events: make(eventTableStore),
	}
}

// CreateRichDispatcher initializes empty dispatcher ready to register events
// and additionally creates a channelfinder from a set of irc.ProtoCaps in order
// to properly send Privmsg(User|Channel)/Notice(User|Channel) events.
func CreateRichDispatcher(caps *irc.ProtoCaps) (*Dispatcher, error) {
	if caps == nil {
		return nil, errProtoCapsMissing
	}
	f, err := irc.CreateChannelFinder(caps.Chantypes)
	if err != nil {
		return nil, err
	}
	return &Dispatcher{
		events: make(eventTableStore),
		caps:   caps,
		finder: f,
	}, nil
}

// Register registers an event handler to a particular event. In return a
// unique identifer is given to later pass into Unregister in case of a need
// to unregister the event handler.
func (d *Dispatcher) Register(event string, handler interface{}) int {
	event = strings.ToUpper(event)
	id := rand.Int()
	if ev, ok := d.events[event]; !ok {
		d.events[event] = make(eventTable)
	} else {
		for _, has := ev[id]; has; id = rand.Int() {
		}
	}

	d.events[event][id] = handler
	return id
}

// Unregister uses the event name, and the identifier returned by Register to
// unregister a callback from the Dispatcher. If the callback was removed it
// returns true, false if it could not be found.
func (d *Dispatcher) Unregister(event string, id int) bool {
	event = strings.ToUpper(event)
	if ev, ok := d.events[event]; ok {
		if _, ok := ev[id]; ok {
			delete(ev, id)
			return true
		}
	}
	return false
}

// Dispatch an IrcMessage to event handlers handling event also ensures all raw
// handlers receive all messages. Returns false if no eventtable was found for
// the primary sent event.
func (d *Dispatcher) Dispatch(event string, msg *irc.IrcMessage) bool {
	event = strings.ToUpper(event)
	handled := d.dispatchHelper(event, msg)
	d.dispatchHelper(irc.RAW, msg)

	return handled
}

// dispatchHelper locates a handler and attempts to resolve it with
// resolveHandler. It returns true if it was able to find an event table.
func (d *Dispatcher) dispatchHelper(event string, msg *irc.IrcMessage) bool {
	if evtable, ok := d.events[event]; ok {
		for _, handler := range evtable {
			go d.resolveHandler(handler, event, msg)
		}
		return true
	}
	return false
}

// resolveHandler checks the type of the handler passed in, resolves it to a
// real type, coerces the IrcMessage in whatever way necessary and then
// calls that handlers primary dispatch method with the coerced message.
func (d *Dispatcher) resolveHandler(
	handler interface{}, event string, msg *irc.IrcMessage) {

	switch t := handler.(type) {
	case PrivmsgUserHandler:
		if d.finder != nil && !d.finder.IsChannel(msg.Args[0]) {
			t.PrivmsgUser(&Message{msg})
		}
	case PrivmsgChannelHandler:
		if d.finder != nil && d.finder.IsChannel(msg.Args[0]) {
			t.PrivmsgChannel(&Message{msg})
		}
	case PrivmsgHandler:
		t.Privmsg(&Message{msg})
	case NoticeUserHandler:
		if d.finder != nil && !d.finder.IsChannel(msg.Args[0]) {
			t.NoticeUser(&Message{msg})
		}
	case NoticeChannelHandler:
		if d.finder != nil && d.finder.IsChannel(msg.Args[0]) {
			t.NoticeChannel(&Message{msg})
		}
	case NoticeHandler:
		t.Notice(&Message{msg})
	case EventHandler:
		t.HandleRaw(msg)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"io"
	"log"
	"strings"
	"launchpad.net/goyaml"

	"code.google.com/p/go.crypto/ssh"
	_ "github.com/go-sql-driver/mysql"
	"github.com/thinxer/tagbbs"
)

var (
	flagDB = flag.String("db", "mysql://bbs:bbs@/bbs?parseTime=true", "connection string")
	bbs    *tagbbs.BBS
)

func bbsinit() {
	var (
		store tagbbs.Storage
		err   error
	)
	parts := strings.SplitN(*flagDB, "://", 2)
	driver := parts[0]
	if driver == "mysql" {
		store, err = tagbbs.NewSQLStore(driver, parts[1], "kvs")
		if err != nil {
			panic(err)
		}
	} else if driver == "mem" {
		store = tagbbs.MemStore{}
	} else {
		panic("unknown driver: " + driver)
	}
	bbs = tagbbs.NewBBS(store)
}

func userauth(user string, password string) bool {
	return bbs.Auth(user, password)
}

func userpubkey(user string, algo string, pubkey []byte) bool {
	type Profile struct {
		AuthorizedKeys string `yaml:"authorized_keys"`
	}
	post, err := bbs.Get("user:"+user, user)
	if err != nil {
		return false
	}
	profile := Profile{}
	err = goyaml.Unmarshal(post.Content, &profile)
	if err != nil {
		return false
	}
	keys := []byte(profile.AuthorizedKeys)
	for {
		var (
			pkey ssh.PublicKey
			ok   bool
		)
		pkey, _, _, keys, ok = ssh.ParseAuthorizedKey(keys)
		if !ok {
			return false
		}
		if pkey.PublicKeyAlgo() == algo && bytes.Compare(ssh.MarshalPublicKey(pkey), pubkey) == 0 {
			return true
		}
	}
}

// Replace with custom terminal implementation.
func usermain(user string, ch ssh.Channel) {
	// catch all errors so that the main server won't crash
	defer func() {
		if err := recover(); err != nil {
			log.Println("!!! Err:", err)
		}
	}()

	// create a term on top of the connection
	term := NewTerminal(ch, "> ")
	serverTerm := &ssh.ServerTerminal{
		Term:    term,
		Channel: ch,
	}

	// ready to accept the connection
	ch.Accept()
	defer ch.Close()

	// real logic here
	for {
		line, err := serverTerm.ReadLine()
		if err == io.EOF {
			return
		} else if err != nil {
			log.Println("readLine err:", err)
			continue
		}
		// deal with the command
		cmds := strings.Split(line, " ")
		switch cmds[0] {
		case "help":
			term.Println("Available Commands:")
			term.Println("help\r\nregister\r\npasswd\r\nwho\r\nlist [tag]\r\nget key\r\nput [key]")
		case "register":
			err := bbs.NewUser(user)
			term.PerrorIf(err)
		case "passwd":
			pass1, _ := term.ReadPassword("New Password:")
			pass2, _ := term.ReadPassword("New Password Again:")
			if pass1 != pass2 {
				term.Perror("Password Mismatch")
			} else {
				term.PerrorIf(bbs.SetUserPass(user, pass1))
			}
		case "who":
			term.Println(user)
			term.PerrorIf(nil)
		case "list":
			name := ""
			if len(cmds) > 1 {
				name = cmds[1]
			}
			ids, err := bbs.Query(name)
			for _, id := range ids {
				p, _ := bbs.Get(id, user)
				fm := p.FrontMatter()
				if fm == nil {
					continue
				}
				paddedTitle := fm.Title
				width := StringWidth(paddedTitle)
				if width < 40 {
					paddedTitle += strings.Repeat(" ", 40-width)
				}
				term.Printf("%-12s %s %v %v\r\n", id, paddedTitle, fm.Authors, fm.Tags)
			}
			term.PerrorIf(err)
		case "get":
			key := ""
			if len(cmds) > 1 {
				key = cmds[1]
			}
			post, err := bbs.Get(key, user)

			if err != nil {
				term.Perror(err)
			} else {
				term.Pokay(post.Rev, post.Timestamp)
				term.WriteUnix(post.Content)
			}
		case "put":
			var (
				key  string
				post tagbbs.Post
			)
			if len(cmds) > 1 {
				key = cmds[1]
			}
			if len(key) > 0 {
				var err error
				post, err = bbs.Get(key, user)
				if err != nil {
					term.Perror(err)
				}
			}
			term.SetPrompt("")
			term.Pokay("Type EOF to end the file.")
			content, err := readutil(serverTerm, "EOF")
			if err == nil {
				post.Rev++
				post.Content = []byte(content)
				if len(key) == 0 {
					key = bbs.NewPostKey()
				}
				term.PerrorIf(bbs.Put(key, post, user))
			}
			term.SetPrompt("> ")
		}
	}
}

func readutil(term *ssh.ServerTerminal, until string) (content string, err error) {
	var line string
	for {
		line, err = term.ReadLine()
		if err != nil {
			return
		}
		if line == "EOF" {
			break
		}
		content += line + "\n"
	}
	return
}
<commit_msg>term: allow discard<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"io"
	"log"
	"strings"
	"launchpad.net/goyaml"

	"code.google.com/p/go.crypto/ssh"
	_ "github.com/go-sql-driver/mysql"
	"github.com/thinxer/tagbbs"
)

var (
	flagDB = flag.String("db", "mysql://bbs:bbs@/bbs?parseTime=true", "connection string")
	bbs    *tagbbs.BBS
)

var (
	ErrDiscarded = errors.New("Discarded")
)

func bbsinit() {
	var (
		store tagbbs.Storage
		err   error
	)
	parts := strings.SplitN(*flagDB, "://", 2)
	driver := parts[0]
	if driver == "mysql" {
		store, err = tagbbs.NewSQLStore(driver, parts[1], "kvs")
		if err != nil {
			panic(err)
		}
	} else if driver == "mem" {
		store = tagbbs.MemStore{}
	} else {
		panic("unknown driver: " + driver)
	}
	bbs = tagbbs.NewBBS(store)
}

func userauth(user string, password string) bool {
	return bbs.Auth(user, password)
}

func userpubkey(user string, algo string, pubkey []byte) bool {
	type Profile struct {
		AuthorizedKeys string `yaml:"authorized_keys"`
	}
	post, err := bbs.Get("user:"+user, user)
	if err != nil {
		return false
	}
	profile := Profile{}
	err = goyaml.Unmarshal(post.Content, &profile)
	if err != nil {
		return false
	}
	keys := []byte(profile.AuthorizedKeys)
	for {
		var (
			pkey ssh.PublicKey
			ok   bool
		)
		pkey, _, _, keys, ok = ssh.ParseAuthorizedKey(keys)
		if !ok {
			return false
		}
		if pkey.PublicKeyAlgo() == algo && bytes.Compare(ssh.MarshalPublicKey(pkey), pubkey) == 0 {
			return true
		}
	}
}

// Replace with custom terminal implementation.
func usermain(user string, ch ssh.Channel) {
	// catch all errors so that the main server won't crash
	defer func() {
		if err := recover(); err != nil {
			log.Println("!!! Err:", err)
		}
	}()

	// create a term on top of the connection
	term := NewTerminal(ch, "> ")
	serverTerm := &ssh.ServerTerminal{
		Term:    term,
		Channel: ch,
	}

	// ready to accept the connection
	ch.Accept()
	defer ch.Close()

	// real logic here
	for {
		line, err := serverTerm.ReadLine()
		if err == io.EOF {
			return
		} else if err != nil {
			log.Println("readLine err:", err)
			continue
		}
		// deal with the command
		cmds := strings.Split(line, " ")
		switch cmds[0] {
		case "help":
			term.Println("Available Commands:")
			term.Println("help\r\nregister\r\npasswd\r\nwho\r\nlist [tag]\r\nget key\r\nput [key]")
		case "register":
			err := bbs.NewUser(user)
			term.PerrorIf(err)
		case "passwd":
			pass1, _ := term.ReadPassword("New Password:")
			pass2, _ := term.ReadPassword("New Password Again:")
			if pass1 != pass2 {
				term.Perror("Password Mismatch")
			} else {
				term.PerrorIf(bbs.SetUserPass(user, pass1))
			}
		case "who":
			term.Println(user)
			term.PerrorIf(nil)
		case "list":
			name := ""
			if len(cmds) > 1 {
				name = cmds[1]
			}
			ids, err := bbs.Query(name)
			for _, id := range ids {
				p, _ := bbs.Get(id, user)
				fm := p.FrontMatter()
				if fm == nil {
					continue
				}
				paddedTitle := fm.Title
				width := StringWidth(paddedTitle)
				if width < 40 {
					paddedTitle += strings.Repeat(" ", 40-width)
				}
				term.Printf("%-12s %s %v %v\r\n", id, paddedTitle, fm.Authors, fm.Tags)
			}
			term.PerrorIf(err)
		case "get":
			key := ""
			if len(cmds) > 1 {
				key = cmds[1]
			}
			post, err := bbs.Get(key, user)

			if err != nil {
				term.Perror(err)
			} else {
				term.Pokay(post.Rev, post.Timestamp)
				term.WriteUnix(post.Content)
			}
		case "put":
			var (
				key  string
				post tagbbs.Post
			)
			if len(cmds) > 1 {
				key = cmds[1]
			}
			if len(key) > 0 {
				var err error
				post, err = bbs.Get(key, user)
				if err != nil {
					term.Perror(err)
				}
			}
			term.SetPrompt("")
			term.Pokay("Type EOF to end the file. Type DISCARD to cancel.")
			content, err := readutil(serverTerm, "EOF")
			if err == nil {
				post.Rev++
				post.Content = []byte(content)
				if len(key) == 0 {
					key = bbs.NewPostKey()
				}
				term.PerrorIf(bbs.Put(key, post, user))
			} else {
				term.Perror(err)
			}
			term.SetPrompt("> ")
		}
	}
}

func readutil(term *ssh.ServerTerminal, until string) (content string, err error) {
	var line string
	for {
		line, err = term.ReadLine()
		if err != nil {
			return
		}
		if line == "EOF" {
			break
		} else if line == "DISCARD" {
			content = ""
			err = ErrDiscarded
			break
		}
		content += line + "\n"
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	mkr "github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
)

var commandAlerts = cli.Command{
	Name:  "alerts",
	Usage: "Retrieve/Close alerts",
	Description: `
    Retrieve/Close alerts. Without subcommand, show all alerts.
    Request APIs under "/api/v0/alerts". See http://help-ja.mackerel.io/entry/spec/api/v0 .
`,
	Action: doAlertsList,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "format, f", Value: "", Usage: "Output format. (human/json)"},
	},
	Subcommands: []cli.Command{
		{
			Name:        "close",
			Usage:       "close alerts",
			Description: "Pull monitor rules from Mackerel server and save them to a file. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doAlertsClose,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "reason, r", Value: "", Usage: "Reason of closing alert."},
				cli.BoolFlag{Name: "verbose, v", Usage: "Verbose output mode"},
			},
		},
	},
}

type alertSet struct {
	Alert   *mkr.Alert
	Host    *mkr.Host
	Monitor *mkr.Monitor
}

func joinMonitorsAndHosts(client *mkr.Client, alerts []*mkr.Alert) []*alertSet {
	hostsJSON, err := client.FindHosts(&mkr.FindHostsParam{
		Statuses: []string{"working", "standby", "poweroff", "maintenance"},
	})
	logger.DieIf(err)

	hosts := map[string]*mkr.Host{}
	for _, host := range hostsJSON {
		hosts[host.ID] = host
	}

	monitorsJSON, err := client.FindMonitors()
	logger.DieIf(err)

	monitors := map[string]*mkr.Monitor{}
	for _, monitor := range monitorsJSON {
		monitors[monitor.ID] = monitor
	}

	alertSets := []*alertSet{}
	for _, alert := range alerts {
		alertSets = append(
			alertSets,
			&alertSet{Alert: alert, Host: hosts[alert.HostID], Monitor: monitors[alert.MonitorID]},
		)
	}
	return alertSets
}

func formatJoinedAlert(alertSet *alertSet) string {
	const layout = "2006-01-02 15:04:05"

	host := alertSet.Host
	monitor := alertSet.Monitor
	alert := alertSet.Alert

	hostMsg := ""
	if host != nil {
		hostMsg = fmt.Sprintf("%s %s", host.Name, host.Status)
		roleMsgs := []string{}
		for service, roles := range host.Roles {
			roleMsgs = append(roleMsgs, fmt.Sprintf("%s:%s", service, strings.Join(roles, ",")))
		}
		hostMsg += " [" + strings.Join(roleMsgs, ", ") + "]"
	}

	monitorMsg := ""
	if monitor != nil {
		switch monitor.Type {
		case "connectivity":
			monitorMsg = fmt.Sprintf("%s", monitor.Type)
		case "host":
			switch alert.Status {
			case "CRITICAL":
				monitorMsg = fmt.Sprintf("%s %.2f %s %.2f", monitor.Metric, alert.Value, monitor.Operator, monitor.Critical)
			case "WARNING":
				monitorMsg = fmt.Sprintf("%s %.2f %s %.2f", monitor.Metric, alert.Value, monitor.Operator, monitor.Warning)
			default:
				monitorMsg = fmt.Sprintf("%s %.2f %s %.2f", monitor.Metric, alert.Value, monitor.Operator, monitor.Critical)
			}
		default:
			monitorMsg = fmt.Sprintf("%s", monitor.Type)
		}
	}
	return fmt.Sprintf("%s %s %8s %s %s", alert.ID, time.Unix(alert.OpenedAt, 0).Format(layout), alert.Status, monitorMsg, hostMsg)
}

func doAlertsList(c *cli.Context) {
	conffile := c.GlobalString("conf")
	client := newMackerel(conffile)

	alerts, err := client.FindAlerts()
	logger.DieIf(err)
	if c.String("format") == "json" {
		PrettyPrintJSON(alerts)
		return
	}

	joinedAlerts := joinMonitorsAndHosts(client, alerts)

	for _, joinAlert := range joinedAlerts {
		fmt.Println(formatJoinedAlert(joinAlert))
	}
}

func doAlertsClose(c *cli.Context) {
	conffile := c.GlobalString("conf")
	isVerbose := c.Bool("verbose")
	argAlertIDs := c.Args()
	reason := c.String("reason")

	if len(argAlertIDs) < 1 {
		cli.ShowCommandHelp(c, "alerts")
		os.Exit(1)
	}

	client := newMackerel(conffile)
	for _, alertID := range argAlertIDs {
		alert, err := client.CloseAlert(alertID, reason)
		logger.DieIf(err)

		logger.Log("Alert closed", alertID)
		if isVerbose == true {
			PrettyPrintJSON(alert)
		}
	}
}
<commit_msg>support colorize option<commit_after>package main

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"github.com/fatih/color"
	mkr "github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
)

var commandAlerts = cli.Command{
	Name:  "alerts",
	Usage: "Retrieve/Close alerts",
	Description: `
    Retrieve/Close alerts. Without subcommand, show all alerts.
    Request APIs under "/api/v0/alerts". See http://help-ja.mackerel.io/entry/spec/api/v0 .
`,
	Action: doAlertsList,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "format, f", Value: "", Usage: "Output format. (human/json)"},
		cli.BoolFlag{Name: "color, c", Usage: "Colorized output. default: true"},
	},
	Subcommands: []cli.Command{
		{
			Name:        "close",
			Usage:       "close alerts",
			Description: "Pull monitor rules from Mackerel server and save them to a file. The file can be specified by filepath argument <file>. The default is 'monitors.json'.",
			Action:      doAlertsClose,
			Flags: []cli.Flag{
				cli.StringFlag{Name: "reason, r", Value: "", Usage: "Reason of closing alert."},
				cli.BoolFlag{Name: "verbose, v", Usage: "Verbose output mode"},
			},
		},
	},
}

type alertSet struct {
	Alert   *mkr.Alert
	Host    *mkr.Host
	Monitor *mkr.Monitor
}

func joinMonitorsAndHosts(client *mkr.Client, alerts []*mkr.Alert) []*alertSet {
	hostsJSON, err := client.FindHosts(&mkr.FindHostsParam{
		Statuses: []string{"working", "standby", "poweroff", "maintenance"},
	})
	logger.DieIf(err)

	hosts := map[string]*mkr.Host{}
	for _, host := range hostsJSON {
		hosts[host.ID] = host
	}

	monitorsJSON, err := client.FindMonitors()
	logger.DieIf(err)

	monitors := map[string]*mkr.Monitor{}
	for _, monitor := range monitorsJSON {
		monitors[monitor.ID] = monitor
	}

	alertSets := []*alertSet{}
	for _, alert := range alerts {
		alertSets = append(
			alertSets,
			&alertSet{Alert: alert, Host: hosts[alert.HostID], Monitor: monitors[alert.MonitorID]},
		)
	}
	return alertSets
}

func formatJoinedAlert(alertSet *alertSet, colorize bool) string {
	const layout = "2006-01-02 15:04:05"

	host := alertSet.Host
	monitor := alertSet.Monitor
	alert := alertSet.Alert

	hostMsg := ""
	if host != nil {
		statusMsg := host.Status
		if host.IsRetired == true {
			statusMsg = "retired"
		}
		if colorize {
			switch statusMsg {
			case "working":
				statusMsg = color.BlueString("working")
			case "standby":
				statusMsg = color.GreenString("standby")
			case "poweroff":
				statusMsg = "poweroff"
			case "maintenance":
				statusMsg = color.YellowString("maintenance")
			}
		}
		hostMsg = fmt.Sprintf("%s %s", host.Name, statusMsg)
		roleMsgs := []string{}
		for service, roles := range host.Roles {
			roleMsgs = append(roleMsgs, fmt.Sprintf("%s:%s", service, strings.Join(roles, ",")))
		}
		hostMsg += " [" + strings.Join(roleMsgs, ", ") + "]"
	}

	monitorMsg := ""
	if monitor != nil {
		switch monitor.Type {
		case "connectivity":
			monitorMsg = fmt.Sprintf("%s", monitor.Type)
		case "host":
			switch alert.Status {
			case "CRITICAL":
				monitorMsg = fmt.Sprintf("%s %.2f %s %.2f", monitor.Metric, alert.Value, monitor.Operator, monitor.Critical)
			case "WARNING":
				monitorMsg = fmt.Sprintf("%s %.2f %s %.2f", monitor.Metric, alert.Value, monitor.Operator, monitor.Warning)
			default:
				monitorMsg = fmt.Sprintf("%s %.2f %s %.2f", monitor.Metric, alert.Value, monitor.Operator, monitor.Critical)
			}
		default:
			monitorMsg = fmt.Sprintf("%s", monitor.Type)
		}
	}
	statusMsg := alert.Status
	if colorize {
		switch alert.Status {
		case "CRITICAL":
			statusMsg = color.RedString("CRITICAL")
		case "WARNING":
			statusMsg = color.YellowString("WARNING ")
		}
	}
	return fmt.Sprintf("%s %s %s %s %s", alert.ID, time.Unix(alert.OpenedAt, 0).Format(layout), statusMsg, monitorMsg, hostMsg)
}

func doAlertsList(c *cli.Context) {
	conffile := c.GlobalString("conf")
	client := newMackerel(conffile)

	alerts, err := client.FindAlerts()
	logger.DieIf(err)
	if c.String("format") == "json" {
		PrettyPrintJSON(alerts)
		return
	}

	joinedAlerts := joinMonitorsAndHosts(client, alerts)

	for _, joinAlert := range joinedAlerts {
		fmt.Println(formatJoinedAlert(joinAlert, c.Bool("color")))
	}
}

func doAlertsClose(c *cli.Context) {
	conffile := c.GlobalString("conf")
	isVerbose := c.Bool("verbose")
	argAlertIDs := c.Args()
	reason := c.String("reason")

	if len(argAlertIDs) < 1 {
		cli.ShowCommandHelp(c, "alerts")
		os.Exit(1)
	}

	client := newMackerel(conffile)
	for _, alertID := range argAlertIDs {
		alert, err := client.CloseAlert(alertID, reason)
		logger.DieIf(err)

		logger.Log("Alert closed", alertID)
		if isVerbose == true {
			PrettyPrintJSON(alert)
		}
	}
}
<|endoftext|>


<commit_before>// Package cloudfoundry implements the OAuth2 protocol for authenticating users through Cloud Foundry
// This package can be used as a reference implementation of an OAuth2 provider for Goth.
package cloudfoundry

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/markbates/goth"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

// Provider is the implementation of `goth.Provider` for accessing Cloud Foundry.
type Provider struct {
	AuthURL     string
	TokenURL    string
	UserInfoURL string
	ClientKey   string
	Secret      string
	CallbackURL string
	Client      *http.Client
	config      *oauth2.Config
}

// New creates a new Cloud Foundry provider and sets up important connection details.
// You should always call `cloudfoundry.New` to get a new provider.  Never try to
// create one manually.
func New(uaaURL, clientKey, secret, callbackURL string, scopes ...string) *Provider {
	uaaURL = strings.TrimSuffix(uaaURL, "/")
	p := &Provider{
		Client:      new(http.Client),
		ClientKey:   clientKey,
		Secret:      secret,
		CallbackURL: callbackURL,
		AuthURL:     uaaURL + "/oauth/authorize",
		TokenURL:    uaaURL + "/oauth/token",
		UserInfoURL: uaaURL + "/userinfo",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return "cloudfoundry"
}

// Debug is a no-op for the cloudfoundry package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Cloud Foundry for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	return &Session{
		AuthURL: p.config.AuthCodeURL(state),
	}, nil
}

// FetchUser will go to Cloud Foundry and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	s := session.(*Session)
	user := goth.User{
		AccessToken:  s.AccessToken,
		Provider:     p.Name(),
		RefreshToken: s.RefreshToken,
		ExpiresAt:    s.ExpiresAt,
	}
	req, err := http.NewRequest("GET", p.UserInfoURL, nil)
	if err != nil {
		return user, err
	}
	req.Header.Set("Authorization", "Bearer "+s.AccessToken)
	resp, err := p.Client.Do(req)
	if err != nil {
		if resp != nil {
			resp.Body.Close()
		}
		return user, err
	}
	defer resp.Body.Close()

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

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&user.RawData)
	if err != nil {
		return user, err
	}

	err = userFromReader(bytes.NewReader(bits), &user)
	return user, err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  provider.AuthURL,
			TokenURL: provider.TokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	}
	return c
}

func userFromReader(r io.Reader, user *goth.User) error {
	u := struct {
		Name  string `json:"user_name"`
		Email string `json:"email"`
		ID    string `json:"user_id"`
	}{}
	err := json.NewDecoder(r).Decode(&u)
	if err != nil {
		return err
	}
	user.Name = u.Name
	user.NickName = u.Name
	user.UserID = u.ID
	user.Email = u.Email
	return nil
}

//RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ctx := context.WithValue(oauth2.NoContext, oauth2.HTTPClient, p.Client)
	ts := p.config.TokenSource(ctx, token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}
<commit_msg>Fix initial package doc for CloudFoundry<commit_after>// Package cloudfoundry implements the OAuth2 protocol for authenticating users through Cloud Foundry.
// This package can be used as a reference implementation of an OAuth2 provider for Goth.
package cloudfoundry

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/markbates/goth"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

// Provider is the implementation of `goth.Provider` for accessing Cloud Foundry.
type Provider struct {
	AuthURL     string
	TokenURL    string
	UserInfoURL string
	ClientKey   string
	Secret      string
	CallbackURL string
	Client      *http.Client
	config      *oauth2.Config
}

// New creates a new Cloud Foundry provider and sets up important connection details.
// You should always call `cloudfoundry.New` to get a new provider.  Never try to
// create one manually.
func New(uaaURL, clientKey, secret, callbackURL string, scopes ...string) *Provider {
	uaaURL = strings.TrimSuffix(uaaURL, "/")
	p := &Provider{
		Client:      new(http.Client),
		ClientKey:   clientKey,
		Secret:      secret,
		CallbackURL: callbackURL,
		AuthURL:     uaaURL + "/oauth/authorize",
		TokenURL:    uaaURL + "/oauth/token",
		UserInfoURL: uaaURL + "/userinfo",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Name is the name used to retrieve this provider later.
func (p *Provider) Name() string {
	return "cloudfoundry"
}

// Debug is a no-op for the cloudfoundry package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Cloud Foundry for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	return &Session{
		AuthURL: p.config.AuthCodeURL(state),
	}, nil
}

// FetchUser will go to Cloud Foundry and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	s := session.(*Session)
	user := goth.User{
		AccessToken:  s.AccessToken,
		Provider:     p.Name(),
		RefreshToken: s.RefreshToken,
		ExpiresAt:    s.ExpiresAt,
	}
	req, err := http.NewRequest("GET", p.UserInfoURL, nil)
	if err != nil {
		return user, err
	}
	req.Header.Set("Authorization", "Bearer "+s.AccessToken)
	resp, err := p.Client.Do(req)
	if err != nil {
		if resp != nil {
			resp.Body.Close()
		}
		return user, err
	}
	defer resp.Body.Close()

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

	err = json.NewDecoder(bytes.NewReader(bits)).Decode(&user.RawData)
	if err != nil {
		return user, err
	}

	err = userFromReader(bytes.NewReader(bits), &user)
	return user, err
}

func newConfig(provider *Provider, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  provider.AuthURL,
			TokenURL: provider.TokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	}
	return c
}

func userFromReader(r io.Reader, user *goth.User) error {
	u := struct {
		Name  string `json:"user_name"`
		Email string `json:"email"`
		ID    string `json:"user_id"`
	}{}
	err := json.NewDecoder(r).Decode(&u)
	if err != nil {
		return err
	}
	user.Name = u.Name
	user.NickName = u.Name
	user.UserID = u.ID
	user.Email = u.Email
	return nil
}

//RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}

//RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ctx := context.WithValue(oauth2.NoContext, oauth2.HTTPClient, p.Client)
	ts := p.config.TokenSource(ctx, token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}
<|endoftext|>


<commit_before>/*
Aspect is a relational database toolkit for Go:

* Build complete database schemas
* Create reusable and cross-dialect SQL statements
* Allow struct instances and slices to be directly populated by the database
*/
package aspect

const VERSION = "0.3.1"
<commit_msg>Incremented version to 0.4.0<commit_after>/*
Aspect is a relational database toolkit for Go:

* Build complete database schemas
* Create reusable and cross-dialect SQL statements
* Allow struct instances and slices to be directly populated by the database
*/
package aspect

const VERSION = "0.4.0"
<|endoftext|>