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

// +build cgo

package openssl

// #include <openssl/conf.h>
// #include <openssl/ssl.h>
// #include <openssl/x509v3.h>
//
// void OPENSSL_free_not_a_macro(void *ref) { OPENSSL_free(ref); }
//
import "C"

import (
	"errors"
	"io/ioutil"
	"runtime"
	"time"
	"unsafe"
)

type EVP_MD int

const (
	EVP_NULL      EVP_MD = iota
	EVP_MD5       EVP_MD = iota
	EVP_SHA       EVP_MD = iota
	EVP_SHA1      EVP_MD = iota
	EVP_DSS       EVP_MD = iota
	EVP_DSS1      EVP_MD = iota
	EVP_MDC2      EVP_MD = iota
	EVP_RIPEMD160 EVP_MD = iota
	EVP_SHA224    EVP_MD = iota
	EVP_SHA256    EVP_MD = iota
	EVP_SHA384    EVP_MD = iota
	EVP_SHA512    EVP_MD = iota
)

type Certificate struct {
	x      *C.X509
	Issuer *Certificate
	ref    interface{}
	pubKey PublicKey
}

type CertificateInfo struct {
	Serial       int
	Issued       time.Duration
	Expires      time.Duration
	Country      string
	Organization string
	CommonName   string
}

type Name struct {
	name *C.X509_NAME
}

// Allocate and return a new Name object.
func NewName() (*Name, error) {
	n := C.X509_NAME_new()
	if n == nil {
		return nil, errors.New("could not create x509 name")
	}
	name := &Name{name: n}
	runtime.SetFinalizer(name, func(n *Name) {
		C.X509_NAME_free(n.name)
	})
	return name, nil
}

// AddTextEntry appends a text entry to an X509 NAME.
func (n *Name) AddTextEntry(field, value string) error {
	cfield := C.CString(field)
	defer C.free(unsafe.Pointer(cfield))
	cvalue := (*C.uchar)(unsafe.Pointer(C.CString(value)))
	defer C.free(unsafe.Pointer(cvalue))
	ret := C.X509_NAME_add_entry_by_txt(
		n.name, cfield, C.MBSTRING_ASC, cvalue, -1, -1, 0)
	if ret != 1 {
		return errors.New("failed to add x509 name text entry")
	}
	return nil
}

// AddTextEntries allows adding multiple entries to a name in one call.
func (n *Name) AddTextEntries(entries map[string]string) error {
	for f, v := range entries {
		if err := n.AddTextEntry(f, v); err != nil {
			return err
		}
	}
	return nil
}

// NewCertificate generates a basic certificate based
// on the provided CertificateInfo struct
func NewCertificate(info *CertificateInfo, key PublicKey) (*Certificate, error) {
	c := &Certificate{x: C.X509_new()}
	runtime.SetFinalizer(c, func(c *Certificate) {
		C.X509_free(c.x)
	})

	name, err := c.GetSubjectName()
	if err != nil {
		return nil, err
	}
	err = name.AddTextEntries(map[string]string{
		"C":  info.Country,
		"O":  info.Organization,
		"CN": info.CommonName,
	})
	if err != nil {
		return nil, err
	}
	// self-issue for now
	if err := c.SetIssuerName(name); err != nil {
		return nil, err
	}
	if err := c.SetSerial(info.Serial); err != nil {
		return nil, err
	}
	if err := c.SetIssueDate(info.Issued); err != nil {
		return nil, err
	}
	if err := c.SetExpireDate(info.Expires); err != nil {
		return nil, err
	}
	if err := c.SetPubKey(key); err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Certificate) GetSubjectName() (*Name, error) {
	n := C.X509_get_subject_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get subject name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) GetIssuerName() (*Name, error) {
	n := C.X509_get_issuer_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get issuer name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) SetSubjectName(name *Name) error {
	if C.X509_set_subject_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetIssuer updates the stored Issuer cert
// and the internal x509 Issuer Name of a certificate.
// The stored Issuer reference is used when adding extensions.
func (c *Certificate) SetIssuer(issuer *Certificate) error {
	name, err := issuer.GetSubjectName()
	if err != nil {
		return err
	}
	if err = c.SetIssuerName(name); err != nil {
		return err
	}
	c.Issuer = issuer
	return nil
}

// SetIssuerName populates the issuer name of a certificate.
// Use SetIssuer instead, if possible.
func (c *Certificate) SetIssuerName(name *Name) error {
	if C.X509_set_issuer_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetSerial sets the serial of a certificate.
func (c *Certificate) SetSerial(serial int) error {
	if C.ASN1_INTEGER_set(C.X509_get_serialNumber(c.x), C.long(serial)) != 1 {
		return errors.New("failed to set serial")
	}
	return nil
}

// SetIssueDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetIssueDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notBefore, offset)
	if result == nil {
		return errors.New("failed to set issue date")
	}
	return nil
}

// SetExpireDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetExpireDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notAfter, offset)
	if result == nil {
		return errors.New("failed to set expire date")
	}
	return nil
}

// SetPubKey assigns a new public key to a certificate.
func (c *Certificate) SetPubKey(pubKey PublicKey) error {
	c.pubKey = pubKey
	if C.X509_set_pubkey(c.x, pubKey.evpPKey()) != 1 {
		return errors.New("failed to set public key")
	}
	return nil
}

// Sign a certificate using a private key and a digest name.
// Accepted digest names are 'sha256', 'sha384', and 'sha512'.
func (c *Certificate) Sign(privKey PrivateKey, digest EVP_MD) error {
	var md *C.EVP_MD
	switch digest {
	// please don't use these digest functions
	case EVP_NULL:
		md = C.EVP_md_null()
	case EVP_MD5:
		md = C.EVP_md5()
	case EVP_SHA:
		md = C.EVP_sha()
	case EVP_SHA1:
		md = C.EVP_sha1()
	case EVP_DSS:
		md = C.EVP_dss()
	case EVP_DSS1:
		md = C.EVP_dss1()
	case EVP_RIPEMD160:
		md = C.EVP_ripemd160()
	case EVP_SHA224:
		md = C.EVP_sha224()
	// you actually want one of these
	case EVP_SHA256:
		md = C.EVP_sha256()
	case EVP_SHA384:
		md = C.EVP_sha384()
	case EVP_SHA512:
		md = C.EVP_sha512()
	}
	if C.X509_sign(c.x, privKey.evpPKey(), md) <= 0 {
		return errors.New("failed to sign certificate")
	}
	return nil
}

// Add an extension to a certificate.
// Extension constants are NID_* as found in openssl.
func (c *Certificate) AddExtension(nid NID, value string) error {
	issuer := c
	if c.Issuer != nil {
		issuer = c.Issuer
	}
	var ctx C.X509V3_CTX
	C.X509V3_set_ctx(&ctx, c.x, issuer.x, nil, nil, 0)
	ex := C.X509V3_EXT_conf_nid(nil, &ctx, C.int(nid), C.CString(value))
	if ex == nil {
		return errors.New("failed to create x509v3 extension")
	}
	defer C.X509_EXTENSION_free(ex)
	if C.X509_add_ext(c.x, ex, -1) <= 0 {
		return errors.New("failed to add x509v3 extension")
	}
	return nil
}

// Wraps AddExtension using a map of NID to text extension.
// Will return without finishing if it encounters an error.
func (c *Certificate) AddExtensions(extensions map[NID]string) error {
	for nid, value := range extensions {
		if err := c.AddExtension(nid, value); err != nil {
			return err
		}
	}
	return nil
}

// LoadCertificateFromPEM loads an X509 certificate from a PEM-encoded block.
func LoadCertificateFromPEM(pem_block []byte) (*Certificate, error) {
	if len(pem_block) == 0 {
		return nil, errors.New("empty pem block")
	}
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bio := C.BIO_new_mem_buf(unsafe.Pointer(&pem_block[0]),
		C.int(len(pem_block)))
	cert := C.PEM_read_bio_X509(bio, nil, nil, nil)
	C.BIO_free(bio)
	if cert == nil {
		return nil, errorFromErrorQueue()
	}
	x := &Certificate{x: cert}
	runtime.SetFinalizer(x, func(x *Certificate) {
		C.X509_free(x.x)
	})
	return x, nil
}

// MarshalPEM converts the X509 certificate to PEM-encoded format
func (c *Certificate) MarshalPEM() (pem_block []byte, err error) {
	bio := C.BIO_new(C.BIO_s_mem())
	if bio == nil {
		return nil, errors.New("failed to allocate memory BIO")
	}
	defer C.BIO_free(bio)
	if int(C.PEM_write_bio_X509(bio, c.x)) != 1 {
		return nil, errors.New("failed dumping certificate")
	}
	return ioutil.ReadAll(asAnyBio(bio))
}

// PublicKey returns the public key embedded in the X509 certificate.
func (c *Certificate) PublicKey() (PublicKey, error) {
	pkey := C.X509_get_pubkey(c.x)
	if pkey == nil {
		return nil, errors.New("no public key found")
	}
	key := &pKey{key: pkey}
	runtime.SetFinalizer(key, func(key *pKey) {
		C.EVP_PKEY_free(key.key)
	})
	return key, nil
}

// GetSerialNumberHex returns the certificate's serial number in hex format
func (c *Certificate) GetSerialNumberHex() (serial string) {
	asn1_i := C.X509_get_serialNumber(c.x)
	bignum := C.ASN1_INTEGER_to_BN(asn1_i, nil)
	hex := C.BN_bn2hex(bignum)
	serial = C.GoString(hex)
	C.BN_free(bignum)
	C.OPENSSL_free_not_a_macro(unsafe.Pointer(hex))
	return
}
<commit_msg>Allow certificate serials to be 64 bit integers<commit_after>// Copyright (C) 2014 Space Monkey, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build cgo

package openssl

// #include <openssl/conf.h>
// #include <openssl/ssl.h>
// #include <openssl/x509v3.h>
//
// void OPENSSL_free_not_a_macro(void *ref) { OPENSSL_free(ref); }
//
import "C"

import (
	"errors"
	"io/ioutil"
	"runtime"
	"time"
	"unsafe"
)

type EVP_MD int

const (
	EVP_NULL      EVP_MD = iota
	EVP_MD5       EVP_MD = iota
	EVP_SHA       EVP_MD = iota
	EVP_SHA1      EVP_MD = iota
	EVP_DSS       EVP_MD = iota
	EVP_DSS1      EVP_MD = iota
	EVP_MDC2      EVP_MD = iota
	EVP_RIPEMD160 EVP_MD = iota
	EVP_SHA224    EVP_MD = iota
	EVP_SHA256    EVP_MD = iota
	EVP_SHA384    EVP_MD = iota
	EVP_SHA512    EVP_MD = iota
)

type Certificate struct {
	x      *C.X509
	Issuer *Certificate
	ref    interface{}
	pubKey PublicKey
}

type CertificateInfo struct {
	Serial       int64
	Issued       time.Duration
	Expires      time.Duration
	Country      string
	Organization string
	CommonName   string
}

type Name struct {
	name *C.X509_NAME
}

// Allocate and return a new Name object.
func NewName() (*Name, error) {
	n := C.X509_NAME_new()
	if n == nil {
		return nil, errors.New("could not create x509 name")
	}
	name := &Name{name: n}
	runtime.SetFinalizer(name, func(n *Name) {
		C.X509_NAME_free(n.name)
	})
	return name, nil
}

// AddTextEntry appends a text entry to an X509 NAME.
func (n *Name) AddTextEntry(field, value string) error {
	cfield := C.CString(field)
	defer C.free(unsafe.Pointer(cfield))
	cvalue := (*C.uchar)(unsafe.Pointer(C.CString(value)))
	defer C.free(unsafe.Pointer(cvalue))
	ret := C.X509_NAME_add_entry_by_txt(
		n.name, cfield, C.MBSTRING_ASC, cvalue, -1, -1, 0)
	if ret != 1 {
		return errors.New("failed to add x509 name text entry")
	}
	return nil
}

// AddTextEntries allows adding multiple entries to a name in one call.
func (n *Name) AddTextEntries(entries map[string]string) error {
	for f, v := range entries {
		if err := n.AddTextEntry(f, v); err != nil {
			return err
		}
	}
	return nil
}

// NewCertificate generates a basic certificate based
// on the provided CertificateInfo struct
func NewCertificate(info *CertificateInfo, key PublicKey) (*Certificate, error) {
	c := &Certificate{x: C.X509_new()}
	runtime.SetFinalizer(c, func(c *Certificate) {
		C.X509_free(c.x)
	})

	name, err := c.GetSubjectName()
	if err != nil {
		return nil, err
	}
	err = name.AddTextEntries(map[string]string{
		"C":  info.Country,
		"O":  info.Organization,
		"CN": info.CommonName,
	})
	if err != nil {
		return nil, err
	}
	// self-issue for now
	if err := c.SetIssuerName(name); err != nil {
		return nil, err
	}
	if err := c.SetSerial(info.Serial); err != nil {
		return nil, err
	}
	if err := c.SetIssueDate(info.Issued); err != nil {
		return nil, err
	}
	if err := c.SetExpireDate(info.Expires); err != nil {
		return nil, err
	}
	if err := c.SetPubKey(key); err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Certificate) GetSubjectName() (*Name, error) {
	n := C.X509_get_subject_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get subject name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) GetIssuerName() (*Name, error) {
	n := C.X509_get_issuer_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get issuer name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) SetSubjectName(name *Name) error {
	if C.X509_set_subject_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetIssuer updates the stored Issuer cert
// and the internal x509 Issuer Name of a certificate.
// The stored Issuer reference is used when adding extensions.
func (c *Certificate) SetIssuer(issuer *Certificate) error {
	name, err := issuer.GetSubjectName()
	if err != nil {
		return err
	}
	if err = c.SetIssuerName(name); err != nil {
		return err
	}
	c.Issuer = issuer
	return nil
}

// SetIssuerName populates the issuer name of a certificate.
// Use SetIssuer instead, if possible.
func (c *Certificate) SetIssuerName(name *Name) error {
	if C.X509_set_issuer_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetSerial sets the serial of a certificate.
func (c *Certificate) SetSerial(serial int64) error {
	if C.ASN1_INTEGER_set(C.X509_get_serialNumber(c.x), C.long(serial)) != 1 {
		return errors.New("failed to set serial")
	}
	return nil
}

// SetIssueDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetIssueDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notBefore, offset)
	if result == nil {
		return errors.New("failed to set issue date")
	}
	return nil
}

// SetExpireDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetExpireDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notAfter, offset)
	if result == nil {
		return errors.New("failed to set expire date")
	}
	return nil
}

// SetPubKey assigns a new public key to a certificate.
func (c *Certificate) SetPubKey(pubKey PublicKey) error {
	c.pubKey = pubKey
	if C.X509_set_pubkey(c.x, pubKey.evpPKey()) != 1 {
		return errors.New("failed to set public key")
	}
	return nil
}

// Sign a certificate using a private key and a digest name.
// Accepted digest names are 'sha256', 'sha384', and 'sha512'.
func (c *Certificate) Sign(privKey PrivateKey, digest EVP_MD) error {
	var md *C.EVP_MD
	switch digest {
	// please don't use these digest functions
	case EVP_NULL:
		md = C.EVP_md_null()
	case EVP_MD5:
		md = C.EVP_md5()
	case EVP_SHA:
		md = C.EVP_sha()
	case EVP_SHA1:
		md = C.EVP_sha1()
	case EVP_DSS:
		md = C.EVP_dss()
	case EVP_DSS1:
		md = C.EVP_dss1()
	case EVP_RIPEMD160:
		md = C.EVP_ripemd160()
	case EVP_SHA224:
		md = C.EVP_sha224()
	// you actually want one of these
	case EVP_SHA256:
		md = C.EVP_sha256()
	case EVP_SHA384:
		md = C.EVP_sha384()
	case EVP_SHA512:
		md = C.EVP_sha512()
	}
	if C.X509_sign(c.x, privKey.evpPKey(), md) <= 0 {
		return errors.New("failed to sign certificate")
	}
	return nil
}

// Add an extension to a certificate.
// Extension constants are NID_* as found in openssl.
func (c *Certificate) AddExtension(nid NID, value string) error {
	issuer := c
	if c.Issuer != nil {
		issuer = c.Issuer
	}
	var ctx C.X509V3_CTX
	C.X509V3_set_ctx(&ctx, c.x, issuer.x, nil, nil, 0)
	ex := C.X509V3_EXT_conf_nid(nil, &ctx, C.int(nid), C.CString(value))
	if ex == nil {
		return errors.New("failed to create x509v3 extension")
	}
	defer C.X509_EXTENSION_free(ex)
	if C.X509_add_ext(c.x, ex, -1) <= 0 {
		return errors.New("failed to add x509v3 extension")
	}
	return nil
}

// Wraps AddExtension using a map of NID to text extension.
// Will return without finishing if it encounters an error.
func (c *Certificate) AddExtensions(extensions map[NID]string) error {
	for nid, value := range extensions {
		if err := c.AddExtension(nid, value); err != nil {
			return err
		}
	}
	return nil
}

// LoadCertificateFromPEM loads an X509 certificate from a PEM-encoded block.
func LoadCertificateFromPEM(pem_block []byte) (*Certificate, error) {
	if len(pem_block) == 0 {
		return nil, errors.New("empty pem block")
	}
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bio := C.BIO_new_mem_buf(unsafe.Pointer(&pem_block[0]),
		C.int(len(pem_block)))
	cert := C.PEM_read_bio_X509(bio, nil, nil, nil)
	C.BIO_free(bio)
	if cert == nil {
		return nil, errorFromErrorQueue()
	}
	x := &Certificate{x: cert}
	runtime.SetFinalizer(x, func(x *Certificate) {
		C.X509_free(x.x)
	})
	return x, nil
}

// MarshalPEM converts the X509 certificate to PEM-encoded format
func (c *Certificate) MarshalPEM() (pem_block []byte, err error) {
	bio := C.BIO_new(C.BIO_s_mem())
	if bio == nil {
		return nil, errors.New("failed to allocate memory BIO")
	}
	defer C.BIO_free(bio)
	if int(C.PEM_write_bio_X509(bio, c.x)) != 1 {
		return nil, errors.New("failed dumping certificate")
	}
	return ioutil.ReadAll(asAnyBio(bio))
}

// PublicKey returns the public key embedded in the X509 certificate.
func (c *Certificate) PublicKey() (PublicKey, error) {
	pkey := C.X509_get_pubkey(c.x)
	if pkey == nil {
		return nil, errors.New("no public key found")
	}
	key := &pKey{key: pkey}
	runtime.SetFinalizer(key, func(key *pKey) {
		C.EVP_PKEY_free(key.key)
	})
	return key, nil
}

// GetSerialNumberHex returns the certificate's serial number in hex format
func (c *Certificate) GetSerialNumberHex() (serial string) {
	asn1_i := C.X509_get_serialNumber(c.x)
	bignum := C.ASN1_INTEGER_to_BN(asn1_i, nil)
	hex := C.BN_bn2hex(bignum)
	serial = C.GoString(hex)
	C.BN_free(bignum)
	C.OPENSSL_free_not_a_macro(unsafe.Pointer(hex))
	return
}
<|endoftext|>


<commit_before>package telebot

import "strconv"

// User object represents a Telegram user, bot
type User struct {
	ID int `json:"id"`

	FirstName    string `json:"first_name"`
	LastName     string `json:"last_name"`
	Username     string `json:"username"`
	LanguageCode string `json:"language_code"`
	IsBot        bool   `json:"is_bot"`
}

// Recipient returns user ID (see Recipient interface).
func (u *User) Recipient() string {
	return strconv.Itoa(u.ID)
}

// Chat object represents a Telegram user, bot, group or a channel.
type Chat struct {
	ID int64 `json:"id"`

	// See telebot.ChatType and consts.
	Type ChatType `json:"type"`

	// Won't be there for ChatPrivate.
	Title string `json:"title"`

	FirstName string `json:"first_name"`
	LastName  string `json:"last_name"`
	Username  string `json:"username"`
}

// Recipient returns chat ID (see Recipient interface).
func (c *Chat) Recipient() string {
	return strconv.FormatInt(c.ID, 10)
}

// ChatMember object represents information about a single chat member.
type ChatMember struct {
	Rights

	User *User        `json:"user"`
	Role MemberStatus `json:"status"`

	// Date when restrictions will be lifted for the user, unix time.
	//
	// If user is restricted for more than 366 days or less than
	// 30 seconds from the current time, they are considered to be
	// restricted forever.
	//
	// Use tb.Forever().
	//
	RestrictedUntil int64 `json:"until_date,omitempty"`
}
<commit_msg>chat: update all types<commit_after>package telebot

import "strconv"

// User object represents a Telegram user, bot.
type User struct {
	ID int `json:"id"`

	FirstName    string `json:"first_name"`
	LastName     string `json:"last_name"`
	Username     string `json:"username"`
	LanguageCode string `json:"language_code"`
	IsBot        bool   `json:"is_bot"`

	// Returns only in getMe
	CanJoinGroups   bool `json:"can_join_groups"`
	CanReadMessages bool `json:"can_read_all_group_messages"`
	SupportsInline  bool `json:"supports_inline_queries"`
}

// Recipient returns user ID (see Recipient interface).
func (u *User) Recipient() string {
	return strconv.Itoa(u.ID)
}

// Chat object represents a Telegram user, bot, group or a channel.
type Chat struct {
	ID int64 `json:"id"`

	// See telebot.ChatType and consts.
	Type ChatType `json:"type"`

	// Won't be there for ChatPrivate.
	Title string `json:"title"`

	FirstName string `json:"first_name"`
	LastName  string `json:"last_name"`
	Username  string `json:"username"`

	// Returns only in getChat
	Photo            *ChatPhoto `json:"photo,omitempty"`
	Description      string     `json:"description,omitempty"`
	InviteLink       string     `json:"invite_link,omitempty"`
	PinnedMessage    *Message   `json:"pinned_message,omitempty"`
	Permissions      *Rights    `json:"permissions,omitempty"`
	SlowMode         int        `json:"slow_mode_delay,omitempty"`
	StickerSet       string     `json:"sticker_set_name,omitempty"`
	CanSetStickerSet bool       `json:"can_set_sticker_set,omitempty"`
}

// ChatPhoto object represents a chat photo.
type ChatPhoto struct {
	// File identifiers of small (160x160) chat photo
	SmallFileID       string `json:"small_file_id"`
	SmallFileUniqueID string `json:"small_file_unique_id"`

	// File identifiers of big (640x640) chat photo
	BigFileID       string `json:"big_file_id"`
	BigFileUniqueID string `json:"big_file_unique_id"`
}

// Recipient returns chat ID (see Recipient interface).
func (c *Chat) Recipient() string {
	return strconv.FormatInt(c.ID, 10)
}

// ChatMember object represents information about a single chat member.
type ChatMember struct {
	Rights

	User  *User        `json:"user"`
	Role  MemberStatus `json:"status"`
	Title string       `json:"custom_title"`

	// Date when restrictions will be lifted for the user, unix time.
	//
	// If user is restricted for more than 366 days or less than
	// 30 seconds from the current time, they are considered to be
	// restricted forever.
	//
	// Use tb.Forever().
	//
	RestrictedUntil int64 `json:"until_date,omitempty"`
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 (
	"github.com/jacobsa/aws/exp/sdb"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"sync"
)

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type integrationTest struct {
	db sdb.SimpleDB

	mutex           sync.Mutex
	domainsToDelete []sdb.Domain  // Protected by mutex
}

func (t *integrationTest) SetUp(i *TestInfo) {
	var err error

	// Open a connection.
	t.db, err = sdb.NewSimpleDB(g_region, g_accessKey)
	AssertEq(nil, err)
}

func (t *integrationTest) ensureDeleted(d sdb.Domain) {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	t.domainsToDelete = append(t.domainsToDelete, d)
}

func (t *integrationTest) TearDown() {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	// Delete each of the domains created during the test.
	for _, d := range t.domainsToDelete {
		ExpectEq(nil, t.db.DeleteDomain(d), "Domain: %s", d.Name())
	}
}

////////////////////////////////////////////////////////////////////////
// Domains
////////////////////////////////////////////////////////////////////////

type DomainsTest struct {
	integrationTest
}

func init() { RegisterTestSuite(&DomainsTest{}) }

func (t *DomainsTest) InvalidAccessKey() {
	// Open a connection with an unknown key ID.
	wrongKey := g_accessKey
	wrongKey.Id += "taco"

	db, err := sdb.NewSimpleDB(g_region, wrongKey)
	AssertEq(nil, err)

	// Attempt to create a domain.
	_, err = db.OpenDomain("some_domain")

	ExpectThat(err, Error(HasSubstr("403")))
	ExpectThat(err, Error(HasSubstr("Key Id")))
	ExpectThat(err, Error(HasSubstr("exist")))
}

func (t *DomainsTest) DomainsHaveIndependentItems() {
	ExpectEq("TODO", "")
}

func (t *DomainsTest) Delete() {
	ExpectEq("TODO", "")
}

func (t *DomainsTest) OpeningTwiceDoesntDeleteExistingItems() {
	ExpectEq("TODO", "")
}

func (t *DomainsTest) DeleteTwice() {
	ExpectEq("TODO", "")
}

////////////////////////////////////////////////////////////////////////
// Items
////////////////////////////////////////////////////////////////////////

type ItemsTest struct {
	integrationTest
}

func init() { RegisterTestSuite(&ItemsTest{}) }

func (t *ItemsTest) WrongAccessKeySecret() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8ItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) PutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenBatchGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetForNonExistentItem() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetForNonExistentItems() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedValuePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedNonExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SuccessfulPreconditions() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteAllAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchDelete() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidSelectQuery() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectAll() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectCount() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithPredicates() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithSortOrder() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithLimit() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectEmptyResultSet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectLargeResultSet() {
	ExpectEq("TODO", "")
}
<commit_msg>DomainsTest.SeparatelyNamedDomainsHaveIndependentItems<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 (
	"github.com/jacobsa/aws/exp/sdb"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"sync"
)

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type integrationTest struct {
	db sdb.SimpleDB

	mutex           sync.Mutex
	domainsToDelete []sdb.Domain  // Protected by mutex
}

func (t *integrationTest) SetUp(i *TestInfo) {
	var err error

	// Open a connection.
	t.db, err = sdb.NewSimpleDB(g_region, g_accessKey)
	AssertEq(nil, err)
}

func (t *integrationTest) ensureDeleted(d sdb.Domain) {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	t.domainsToDelete = append(t.domainsToDelete, d)
}

func (t *integrationTest) TearDown() {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	// Delete each of the domains created during the test.
	for _, d := range t.domainsToDelete {
		ExpectEq(nil, t.db.DeleteDomain(d), "Domain: %s", d.Name())
	}
}

////////////////////////////////////////////////////////////////////////
// Domains
////////////////////////////////////////////////////////////////////////

type DomainsTest struct {
	integrationTest
}

func init() { RegisterTestSuite(&DomainsTest{}) }

func (t *DomainsTest) InvalidAccessKey() {
	// Open a connection with an unknown key ID.
	wrongKey := g_accessKey
	wrongKey.Id += "taco"

	db, err := sdb.NewSimpleDB(g_region, wrongKey)
	AssertEq(nil, err)

	// Attempt to create a domain.
	_, err = db.OpenDomain("some_domain")

	ExpectThat(err, Error(HasSubstr("403")))
	ExpectThat(err, Error(HasSubstr("Key Id")))
	ExpectThat(err, Error(HasSubstr("exist")))
}

func (t *DomainsTest) SeparatelyNamedDomainsHaveIndependentItems() {
	var err error

	// Open two domains.
	domain0, err := t.db.OpenDomain("taco")
	AssertEq(nil, err)
	t.ensureDeleted(domain0)

	domain1, err := t.db.OpenDomain("burrito")
	AssertEq(nil, err)
	t.ensureDeleted(domain1)

	// Set up an item in the first.
	itemName := sdb.ItemName("some_item")
	err = domain0.PutAttributes(
		itemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: "enchilada", Value: "Queso"},
		},
		[]sdb.Precondition{},
	)

	AssertEq(nil, err)

	// Get attributes for the same name in the other domain. There should be
	// none.
	attrs, err := domain1.GetAttributes(itemName, true, []string{})
	AssertEq(nil, err)

	ExpectThat(attrs, ElementsAre())
}

func (t *DomainsTest) IdenticallyNamedDomainsHaveIdenticalItems() {
	ExpectFalse(true, "TODO")
}

func (t *DomainsTest) Delete() {
	ExpectEq("TODO", "")
}

func (t *DomainsTest) OpeningTwiceDoesntDeleteExistingItems() {
	ExpectEq("TODO", "")
}

func (t *DomainsTest) DeleteTwice() {
	ExpectEq("TODO", "")
}

////////////////////////////////////////////////////////////////////////
// Items
////////////////////////////////////////////////////////////////////////

type ItemsTest struct {
	integrationTest
}

func init() { RegisterTestSuite(&ItemsTest{}) }

func (t *ItemsTest) WrongAccessKeySecret() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8ItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) PutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenBatchGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetForNonExistentItem() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetForNonExistentItems() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedValuePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedNonExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SuccessfulPreconditions() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteAllAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchDelete() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidSelectQuery() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectAll() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectCount() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithPredicates() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithSortOrder() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithLimit() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectEmptyResultSet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectLargeResultSet() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>package inspect

import (
	"flag"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"

	"github.com/hashicorp/consul/agent/consul/fsm"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/command/flags"
	"github.com/hashicorp/consul/snapshot"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/go-msgpack/codec"
	"github.com/hashicorp/raft"
	"github.com/mitchellh/cli"
)

func New(ui cli.Ui) *cmd {
	c := &cmd{UI: ui}
	c.init()
	return c
}

type cmd struct {
	UI     cli.Ui
	flags  *flag.FlagSet
	help   string
	format string

	// flags
	kvDetails bool
	kvDepth   int
	kvFilter  string
}

func (c *cmd) init() {
	c.flags = flag.NewFlagSet("", flag.ContinueOnError)
	c.flags.BoolVar(&c.kvDetails, "kvdetails", false,
		"Provides a detailed KV space usage breakdown for any KV data that's been stored.")
	c.flags.IntVar(&c.kvDepth, "kvdepth", 0,
		"Can only be used with -kvdetails. The key prefix depth used to breakdown KV store data. Defaults to 2.")
	c.flags.StringVar(&c.kvFilter, "kvfilter", "",
		"Can only be used with -kvdetails. Limits KV key breakdown using this prefix filter.")
	c.flags.StringVar(
		&c.format,
		"format",
		PrettyFormat,
		fmt.Sprintf("Output format {%s}", strings.Join(GetSupportedFormats(), "|")))

	c.help = flags.Usage(help, c.flags)
}

// MetadataInfo is used for passing information
// through the formatter
type MetadataInfo struct {
	ID      string
	Size    int64
	Index   uint64
	Term    uint64
	Version raft.SnapshotVersion
}

// SnapshotInfo is used for passing snapshot stat
// information between functions
type SnapshotInfo struct {
	Meta        MetadataInfo
	Stats       map[structs.MessageType]typeStats
	StatsKV     map[string]typeStats
	TotalSize   int
	TotalSizeKV int
}

// OutputFormat is used for passing information
// through the formatter
type OutputFormat struct {
	Meta        *MetadataInfo
	Stats       []typeStats
	StatsKV     []typeStats
	TotalSize   int
	TotalSizeKV int
}

func (c *cmd) Run(args []string) int {
	if err := c.flags.Parse(args); err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	var file string
	args = c.flags.Args()

	switch len(args) {
	case 0:
		c.UI.Error("Missing FILE argument")
		return 1
	case 1:
		file = args[0]
	default:
		c.UI.Error(fmt.Sprintf("Too many arguments (expected 1, got %d)", len(args)))
		return 1
	}

	// Open the file.
	f, err := os.Open(file)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error opening snapshot file: %s", err))
		return 1
	}
	defer f.Close()

	readFile, meta, err := snapshot.Read(hclog.New(nil), f)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error reading snapshot: %s", err))
	}
	defer func() {
		if err := readFile.Close(); err != nil {
			c.UI.Error(fmt.Sprintf("Failed to close temp snapshot: %v", err))
		}
		if err := os.Remove(readFile.Name()); err != nil {
			c.UI.Error(fmt.Sprintf("Failed to clean up temp snapshot: %v", err))
		}
	}()

	info, err := c.enhance(readFile)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error extracting snapshot data: %s", err))
		return 1
	}

	formatter, err := NewFormatter(c.format)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error outputting enhanced snapshot data: %s", err))
		return 1
	}
	//Generate structs for the formatter with information we read in
	metaformat := &MetadataInfo{
		ID:      meta.ID,
		Size:    meta.Size,
		Index:   meta.Index,
		Term:    meta.Term,
		Version: meta.Version,
	}

	//Restructures stats given above to be human readable
	formattedStats := generateStats(info)
	formattedStatsKV := generateKVStats(info)

	in := &OutputFormat{
		Meta:        metaformat,
		Stats:       formattedStats,
		StatsKV:     formattedStatsKV,
		TotalSize:   info.TotalSize,
		TotalSizeKV: info.TotalSizeKV,
	}

	out, err := formatter.Format(in)
	if err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	c.UI.Output(out)
	return 0
}

type typeStats struct {
	Name  string
	Sum   int
	Count int
}

func generateStats(info SnapshotInfo) []typeStats {
	ss := make([]typeStats, 0, len(info.Stats))

	for _, s := range info.Stats {
		ss = append(ss, s)
	}

	ss = sortTypeStats(ss)

	return ss
}

func generateKVStats(info SnapshotInfo) []typeStats {
	if info.StatsKV != nil {
		ks := make([]typeStats, 0, len(info.StatsKV))

		for _, s := range info.StatsKV {
			ks = append(ks, s)
		}

		ks = sortTypeStats(ks)

		return ks
	}

	return nil
}

// Sort the stat slice
func sortTypeStats(stats []typeStats) []typeStats {
	sort.Slice(stats, func(i, j int) bool {
		// sort alphabetically if size is equal
		if stats[i].Sum == stats[j].Sum {
			return stats[i].Name < stats[j].Name
		}

		return stats[i].Sum > stats[j].Sum
	})

	return stats
}

// countingReader helps keep track of the bytes we have read
// when reading snapshots
type countingReader struct {
	wrappedReader io.Reader
	read          int
}

func (r *countingReader) Read(p []byte) (n int, err error) {
	n, err = r.wrappedReader.Read(p)
	if err == nil {
		r.read += n
	}
	return n, err
}

// enhance utilizes ReadSnapshot to populate the struct with
// all of the snapshot's itemized data
func (c *cmd) enhance(file io.Reader) (SnapshotInfo, error) {
	info := SnapshotInfo{
		Stats:       make(map[structs.MessageType]typeStats),
		StatsKV:     make(map[string]typeStats),
		TotalSize:   0,
		TotalSizeKV: 0,
	}
	cr := &countingReader{wrappedReader: file}
	handler := func(header *fsm.SnapshotHeader, msg structs.MessageType, dec *codec.Decoder) error {
		name := structs.MessageType.String(msg)
		s := info.Stats[msg]
		if s.Name == "" {
			s.Name = name
		}

		var val interface{}
		err := dec.Decode(&val)
		if err != nil {
			return fmt.Errorf("failed to decode msg type %v, error %v", name, err)
		}

		size := cr.read - info.TotalSize
		s.Sum += size
		s.Count++
		info.TotalSize = cr.read
		info.Stats[msg] = s

		c.kvEnhance(s.Name, val, size, &info)

		return nil
	}
	if err := fsm.ReadSnapshot(cr, handler); err != nil {
		return info, err
	}
	return info, nil

}

func (c *cmd) kvEnhance(keyType string, val interface{}, size int, info *SnapshotInfo) {
	// automatically set kvDetails to true if a depth or filter
	// is provided. this allows the user to omit the -kvdetails
	// flag if they prefer.
	if c.kvDepth != 0 || c.kvFilter != "" {
		c.kvDetails = true
	}

	// set the default depth if one wasn't specified with -kvdepth.
	// this is used rather than the flag default to facilitate the
	// above shortcut.
	if c.kvDetails && c.kvDepth == 0 {
		c.kvDepth = 2
	}

	if c.kvDetails {
		if keyType != "KVS" {
			return
		}

		// have to coerce this into a usable type here or this won't work
		keyVal := val.(map[string]interface{})
		for k, v := range keyVal {
			// we only care about the entry on the key specifically
			// related to the key name, so skip all others
			if k != "Key" {
				continue
			}

			// check for whether a filter is specified. if it is, skip
			// any keys that don't match.
			if len(c.kvFilter) > 0 && !strings.HasPrefix(v.(string), c.kvFilter) {
				break
			}

			split := strings.Split(v.(string), "/")

			// handle the situation where the key is shorter than
			// the specified depth.
			actualDepth := c.kvDepth
			if c.kvDepth > len(split) {
				actualDepth = len(split)
			}
			prefix := strings.Join(split[0:actualDepth], "/")
			kvs := info.StatsKV[prefix]
			if kvs.Name == "" {
				kvs.Name = prefix
			}

			kvs.Sum += size
			kvs.Count++
			info.TotalSizeKV += size
			info.StatsKV[prefix] = kvs
		}
	}
}

func (c *cmd) Synopsis() string {
	return synopsis
}

func (c *cmd) Help() string {
	return c.help
}

const synopsis = "Displays information about a Consul snapshot file"
const help = `
Usage: consul snapshot inspect [options] FILE

  Displays information about a snapshot file on disk.

  To inspect the file "backup.snap":

    $ consul snapshot inspect backup.snap

  For a full list of options and examples, please see the Consul documentation.
`
<commit_msg>Add more func comments<commit_after>package inspect

import (
	"flag"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"

	"github.com/hashicorp/consul/agent/consul/fsm"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/command/flags"
	"github.com/hashicorp/consul/snapshot"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/go-msgpack/codec"
	"github.com/hashicorp/raft"
	"github.com/mitchellh/cli"
)

func New(ui cli.Ui) *cmd {
	c := &cmd{UI: ui}
	c.init()
	return c
}

type cmd struct {
	UI     cli.Ui
	flags  *flag.FlagSet
	help   string
	format string

	// flags
	kvDetails bool
	kvDepth   int
	kvFilter  string
}

func (c *cmd) init() {
	c.flags = flag.NewFlagSet("", flag.ContinueOnError)
	c.flags.BoolVar(&c.kvDetails, "kvdetails", false,
		"Provides a detailed KV space usage breakdown for any KV data that's been stored.")
	c.flags.IntVar(&c.kvDepth, "kvdepth", 0,
		"Can only be used with -kvdetails. The key prefix depth used to breakdown KV store data. Defaults to 2.")
	c.flags.StringVar(&c.kvFilter, "kvfilter", "",
		"Can only be used with -kvdetails. Limits KV key breakdown using this prefix filter.")
	c.flags.StringVar(
		&c.format,
		"format",
		PrettyFormat,
		fmt.Sprintf("Output format {%s}", strings.Join(GetSupportedFormats(), "|")))

	c.help = flags.Usage(help, c.flags)
}

// MetadataInfo is used for passing information
// through the formatter
type MetadataInfo struct {
	ID      string
	Size    int64
	Index   uint64
	Term    uint64
	Version raft.SnapshotVersion
}

// SnapshotInfo is used for passing snapshot stat
// information between functions
type SnapshotInfo struct {
	Meta        MetadataInfo
	Stats       map[structs.MessageType]typeStats
	StatsKV     map[string]typeStats
	TotalSize   int
	TotalSizeKV int
}

// OutputFormat is used for passing information
// through the formatter
type OutputFormat struct {
	Meta        *MetadataInfo
	Stats       []typeStats
	StatsKV     []typeStats
	TotalSize   int
	TotalSizeKV int
}

func (c *cmd) Run(args []string) int {
	if err := c.flags.Parse(args); err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	var file string
	args = c.flags.Args()

	switch len(args) {
	case 0:
		c.UI.Error("Missing FILE argument")
		return 1
	case 1:
		file = args[0]
	default:
		c.UI.Error(fmt.Sprintf("Too many arguments (expected 1, got %d)", len(args)))
		return 1
	}

	// Open the file.
	f, err := os.Open(file)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error opening snapshot file: %s", err))
		return 1
	}
	defer f.Close()

	readFile, meta, err := snapshot.Read(hclog.New(nil), f)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error reading snapshot: %s", err))
	}
	defer func() {
		if err := readFile.Close(); err != nil {
			c.UI.Error(fmt.Sprintf("Failed to close temp snapshot: %v", err))
		}
		if err := os.Remove(readFile.Name()); err != nil {
			c.UI.Error(fmt.Sprintf("Failed to clean up temp snapshot: %v", err))
		}
	}()

	info, err := c.enhance(readFile)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error extracting snapshot data: %s", err))
		return 1
	}

	formatter, err := NewFormatter(c.format)
	if err != nil {
		c.UI.Error(fmt.Sprintf("Error outputting enhanced snapshot data: %s", err))
		return 1
	}
	//Generate structs for the formatter with information we read in
	metaformat := &MetadataInfo{
		ID:      meta.ID,
		Size:    meta.Size,
		Index:   meta.Index,
		Term:    meta.Term,
		Version: meta.Version,
	}

	//Restructures stats given above to be human readable
	formattedStats := generateStats(info)
	formattedStatsKV := generateKVStats(info)

	in := &OutputFormat{
		Meta:        metaformat,
		Stats:       formattedStats,
		StatsKV:     formattedStatsKV,
		TotalSize:   info.TotalSize,
		TotalSizeKV: info.TotalSizeKV,
	}

	out, err := formatter.Format(in)
	if err != nil {
		c.UI.Error(err.Error())
		return 1
	}

	c.UI.Output(out)
	return 0
}

type typeStats struct {
	Name  string
	Sum   int
	Count int
}

// Generate the stats for the output struct that's
// used to produce the printed output the user sees.
func generateStats(info SnapshotInfo) []typeStats {
	ss := make([]typeStats, 0, len(info.Stats))

	for _, s := range info.Stats {
		ss = append(ss, s)
	}

	ss = sortTypeStats(ss)

	return ss
}

// Generate the KV stats for the output struct that's
// used to produce the printed output the user sees.
func generateKVStats(info SnapshotInfo) []typeStats {
	if info.StatsKV != nil {
		ks := make([]typeStats, 0, len(info.StatsKV))

		for _, s := range info.StatsKV {
			ks = append(ks, s)
		}

		ks = sortTypeStats(ks)

		return ks
	}

	return nil
}

// Sort the stat slice
func sortTypeStats(stats []typeStats) []typeStats {
	sort.Slice(stats, func(i, j int) bool {
		// sort alphabetically if size is equal
		if stats[i].Sum == stats[j].Sum {
			return stats[i].Name < stats[j].Name
		}

		return stats[i].Sum > stats[j].Sum
	})

	return stats
}

// countingReader helps keep track of the bytes we have read
// when reading snapshots
type countingReader struct {
	wrappedReader io.Reader
	read          int
}

func (r *countingReader) Read(p []byte) (n int, err error) {
	n, err = r.wrappedReader.Read(p)
	if err == nil {
		r.read += n
	}
	return n, err
}

// enhance utilizes ReadSnapshot to populate the struct with
// all of the snapshot's itemized data
func (c *cmd) enhance(file io.Reader) (SnapshotInfo, error) {
	info := SnapshotInfo{
		Stats:       make(map[structs.MessageType]typeStats),
		StatsKV:     make(map[string]typeStats),
		TotalSize:   0,
		TotalSizeKV: 0,
	}
	cr := &countingReader{wrappedReader: file}
	handler := func(header *fsm.SnapshotHeader, msg structs.MessageType, dec *codec.Decoder) error {
		name := structs.MessageType.String(msg)
		s := info.Stats[msg]
		if s.Name == "" {
			s.Name = name
		}

		var val interface{}
		err := dec.Decode(&val)
		if err != nil {
			return fmt.Errorf("failed to decode msg type %v, error %v", name, err)
		}

		size := cr.read - info.TotalSize
		s.Sum += size
		s.Count++
		info.TotalSize = cr.read
		info.Stats[msg] = s

		c.kvEnhance(s.Name, val, size, &info)

		return nil
	}
	if err := fsm.ReadSnapshot(cr, handler); err != nil {
		return info, err
	}
	return info, nil

}

// kvEnhance populates the struct with all of the snapshot's
// size information for KV data stored in it
func (c *cmd) kvEnhance(keyType string, val interface{}, size int, info *SnapshotInfo) {
	// automatically set kvDetails to true if a depth or filter
	// is provided. this allows the user to omit the -kvdetails
	// flag if they prefer.
	if c.kvDepth != 0 || c.kvFilter != "" {
		c.kvDetails = true
	}

	// set the default depth if one wasn't specified with -kvdepth.
	// this is used rather than the flag default to facilitate the
	// above shortcut.
	if c.kvDetails && c.kvDepth == 0 {
		c.kvDepth = 2
	}

	if c.kvDetails {
		if keyType != "KVS" {
			return
		}

		// have to coerce this into a usable type here or this won't work
		keyVal := val.(map[string]interface{})
		for k, v := range keyVal {
			// we only care about the entry on the key specifically
			// related to the key name, so skip all others
			if k != "Key" {
				continue
			}

			// check for whether a filter is specified. if it is, skip
			// any keys that don't match.
			if len(c.kvFilter) > 0 && !strings.HasPrefix(v.(string), c.kvFilter) {
				break
			}

			split := strings.Split(v.(string), "/")

			// handle the situation where the key is shorter than
			// the specified depth.
			actualDepth := c.kvDepth
			if c.kvDepth > len(split) {
				actualDepth = len(split)
			}
			prefix := strings.Join(split[0:actualDepth], "/")
			kvs := info.StatsKV[prefix]
			if kvs.Name == "" {
				kvs.Name = prefix
			}

			kvs.Sum += size
			kvs.Count++
			info.TotalSizeKV += size
			info.StatsKV[prefix] = kvs
		}
	}
}

func (c *cmd) Synopsis() string {
	return synopsis
}

func (c *cmd) Help() string {
	return c.help
}

const synopsis = "Displays information about a Consul snapshot file"
const help = `
Usage: consul snapshot inspect [options] FILE

  Displays information about a snapshot file on disk.

  To inspect the file "backup.snap":

    $ consul snapshot inspect backup.snap

  For a full list of options and examples, please see the Consul documentation.
`
<|endoftext|>


<commit_before>// Copyright 2015 The Mangos Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use 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 mangos

import (
	"encoding/binary"
	"io"
	"net"
)

// conn implements the Pipe interface on top of net.Conn.  The
// assumption is that transports using this have similar wire protocols,
// and conn is meant to be used as a building block.
type conn struct {
	c     net.Conn
	proto Protocol
	sock  Socket
	open  bool
	props map[string]interface{}
	maxrx int64
}

// connipc is *almost* like a regular conn, but the IPC protocol insists
// on stuffing a leading byte (valued 1) in front of messages.  This is for
// compatibility with nanomsg -- the value cannot ever be anything but 1.
type connipc struct {
	conn
}

// Recv implements the Pipe Recv method.  The message received is expected as
// a 64-bit size (network byte order) followed by the message itself.
func (p *conn) Recv() (*Message, error) {

	var sz int64
	var err error
	var msg *Message

	if err = binary.Read(p.c, binary.BigEndian, &sz); err != nil {
		return nil, err
	}

	// Limit messages to the maximum receive value, if not
	// unlimited.  This avoids a potential denaial of service.
	if sz < 0 || (p.maxrx > 0 && sz > p.maxrx) {
		return nil, ErrTooLong
	}
	msg = NewMessage(int(sz))
	msg.Body = msg.Body[0:sz]
	if _, err = io.ReadFull(p.c, msg.Body); err != nil {
		msg.Free()
		return nil, err
	}
	return msg, nil
}

// Send implements the Pipe Send method.  The message is sent as a 64-bit
// size (network byte order) followed by the message itself.
func (p *conn) Send(msg *Message) error {

	l := uint64(len(msg.Header) + len(msg.Body))

	if msg.Expired() {
		msg.Free()
		return nil
	}

	// send length header
	if err := binary.Write(p.c, binary.BigEndian, l); err != nil {
		return err
	}
	if _, err := p.c.Write(msg.Header); err != nil {
		return err
	}
	// hope this works
	if _, err := p.c.Write(msg.Body); err != nil {
		return err
	}
	msg.Free()
	return nil
}

// LocalProtocol returns our local protocol number.
func (p *conn) LocalProtocol() uint16 {
	return p.proto.Number()
}

// RemoteProtocol returns our peer's protocol number.
func (p *conn) RemoteProtocol() uint16 {
	return p.proto.PeerNumber()
}

// Close implements the Pipe Close method.
func (p *conn) Close() error {
	p.open = false
	return p.c.Close()
}

// IsOpen implements the PipeIsOpen method.
func (p *conn) IsOpen() bool {
	return p.open
}

func (p *conn) GetProp(n string) (interface{}, error) {
	if v, ok := p.props[n]; ok {
		return v, nil
	}
	return nil, ErrBadProperty
}

// NewConnPipe allocates a new Pipe using the supplied net.Conn, and
// initializes it.  It performs the handshake required at the SP layer,
// only returning the Pipe once the SP layer negotiation is complete.
//
// Stream oriented transports can utilize this to implement a Transport.
// The implementation will also need to implement PipeDialer, PipeAccepter,
// and the Transport enclosing structure.   Using this layered interface,
// the implementation needn't bother concerning itself with passing actual
// SP messages once the lower layer connection is established.
func NewConnPipe(c net.Conn, sock Socket, props ...interface{}) (Pipe, error) {
	p := &conn{c: c, proto: sock.GetProtocol(), sock: sock}

	if err := p.handshake(props); err != nil {
		return nil, err
	}

	return p, nil
}

// NewConnPipeIPC allocates a new Pipe using the IPC exchange protocol.
func NewConnPipeIPC(c net.Conn, sock Socket, props ...interface{}) (Pipe, error) {
	p := &connipc{conn: conn{c: c, proto: sock.GetProtocol(), sock: sock}}

	if err := p.handshake(props); err != nil {
		return nil, err
	}

	return p, nil
}

func (p *connipc) Send(msg *Message) error {

	l := uint64(len(msg.Header) + len(msg.Body))
	one := [1]byte{1}
	var err error

	// send length header
	if _, err = p.c.Write(one[:]); err != nil {
		return err
	}
	if err = binary.Write(p.c, binary.BigEndian, l); err != nil {
		return err
	}
	if _, err = p.c.Write(msg.Header); err != nil {
		return err
	}
	// hope this works
	if _, err = p.c.Write(msg.Body); err != nil {
		return err
	}
	msg.Free()
	return nil
}

func (p *connipc) Recv() (*Message, error) {

	var sz int64
	var err error
	var msg *Message
	var one [1]byte

	if _, err = p.c.Read(one[:]); err != nil {
		return nil, err
	}
	if err = binary.Read(p.c, binary.BigEndian, &sz); err != nil {
		return nil, err
	}

	// Limit messages to the maximum receive value, if not
	// unlimited.  This avoids a potential denaial of service.
	if sz < 0 || (p.maxrx > 0 && sz > p.maxrx) {
		return nil, ErrTooLong
	}
	msg = NewMessage(int(sz))
	msg.Body = msg.Body[0:sz]
	if _, err = io.ReadFull(p.c, msg.Body); err != nil {
		msg.Free()
		return nil, err
	}
	return msg, nil
}

// connHeader is exchanged during the initial handshake.
type connHeader struct {
	Zero    byte // must be zero
	S       byte // 'S'
	P       byte // 'P'
	Version byte // only zero at present
	Proto   uint16
	Rsvd    uint16 // always zero at present
}

// handshake establishes an SP connection between peers.  Both sides must
// send the header, then both sides must wait for the peer's header.
// As a side effect, the peer's protocol number is stored in the conn.
// Also, various properties are initialized.
func (p *conn) handshake(props []interface{}) error {
	var err error

	p.props = make(map[string]interface{})
	p.props[PropLocalAddr] = p.c.LocalAddr()
	p.props[PropRemoteAddr] = p.c.RemoteAddr()

	for len(props) >= 2 {
		switch name := props[0].(type) {
		case string:
			p.props[name] = props[1]
		default:
			return ErrBadProperty
		}
		props = props[2:]
	}

	if v, e := p.sock.GetOption(OptionMaxRecvSize); e == nil {
		// socket guarantees this is an integer
		p.maxrx = int64(v.(int))
	}

	h := connHeader{S: 'S', P: 'P', Proto: p.proto.Number()}
	if err = binary.Write(p.c, binary.BigEndian, &h); err != nil {
		return err
	}
	if err = binary.Read(p.c, binary.BigEndian, &h); err != nil {
		p.c.Close()
		return err
	}
	if h.Zero != 0 || h.S != 'S' || h.P != 'P' || h.Rsvd != 0 {
		p.c.Close()
		return ErrBadHeader
	}
	// The only version number we support at present is "0", at offset 3.
	if h.Version != 0 {
		p.c.Close()
		return ErrBadVersion
	}

	// The protocol number lives as 16-bits (big-endian) at offset 4.
	if h.Proto != p.proto.PeerNumber() {
		p.c.Close()
		return ErrBadProto
	}
	p.open = true
	return nil
}
<commit_msg>fixes #255 Make TLS ConnectionState available at RecvMsg (update copyright)<commit_after>// Copyright 2016 The Mangos Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use 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 mangos

import (
	"encoding/binary"
	"io"
	"net"
)

// conn implements the Pipe interface on top of net.Conn.  The
// assumption is that transports using this have similar wire protocols,
// and conn is meant to be used as a building block.
type conn struct {
	c     net.Conn
	proto Protocol
	sock  Socket
	open  bool
	props map[string]interface{}
	maxrx int64
}

// connipc is *almost* like a regular conn, but the IPC protocol insists
// on stuffing a leading byte (valued 1) in front of messages.  This is for
// compatibility with nanomsg -- the value cannot ever be anything but 1.
type connipc struct {
	conn
}

// Recv implements the Pipe Recv method.  The message received is expected as
// a 64-bit size (network byte order) followed by the message itself.
func (p *conn) Recv() (*Message, error) {

	var sz int64
	var err error
	var msg *Message

	if err = binary.Read(p.c, binary.BigEndian, &sz); err != nil {
		return nil, err
	}

	// Limit messages to the maximum receive value, if not
	// unlimited.  This avoids a potential denaial of service.
	if sz < 0 || (p.maxrx > 0 && sz > p.maxrx) {
		return nil, ErrTooLong
	}
	msg = NewMessage(int(sz))
	msg.Body = msg.Body[0:sz]
	if _, err = io.ReadFull(p.c, msg.Body); err != nil {
		msg.Free()
		return nil, err
	}
	return msg, nil
}

// Send implements the Pipe Send method.  The message is sent as a 64-bit
// size (network byte order) followed by the message itself.
func (p *conn) Send(msg *Message) error {

	l := uint64(len(msg.Header) + len(msg.Body))

	if msg.Expired() {
		msg.Free()
		return nil
	}

	// send length header
	if err := binary.Write(p.c, binary.BigEndian, l); err != nil {
		return err
	}
	if _, err := p.c.Write(msg.Header); err != nil {
		return err
	}
	// hope this works
	if _, err := p.c.Write(msg.Body); err != nil {
		return err
	}
	msg.Free()
	return nil
}

// LocalProtocol returns our local protocol number.
func (p *conn) LocalProtocol() uint16 {
	return p.proto.Number()
}

// RemoteProtocol returns our peer's protocol number.
func (p *conn) RemoteProtocol() uint16 {
	return p.proto.PeerNumber()
}

// Close implements the Pipe Close method.
func (p *conn) Close() error {
	p.open = false
	return p.c.Close()
}

// IsOpen implements the PipeIsOpen method.
func (p *conn) IsOpen() bool {
	return p.open
}

func (p *conn) GetProp(n string) (interface{}, error) {
	if v, ok := p.props[n]; ok {
		return v, nil
	}
	return nil, ErrBadProperty
}

// NewConnPipe allocates a new Pipe using the supplied net.Conn, and
// initializes it.  It performs the handshake required at the SP layer,
// only returning the Pipe once the SP layer negotiation is complete.
//
// Stream oriented transports can utilize this to implement a Transport.
// The implementation will also need to implement PipeDialer, PipeAccepter,
// and the Transport enclosing structure.   Using this layered interface,
// the implementation needn't bother concerning itself with passing actual
// SP messages once the lower layer connection is established.
func NewConnPipe(c net.Conn, sock Socket, props ...interface{}) (Pipe, error) {
	p := &conn{c: c, proto: sock.GetProtocol(), sock: sock}

	if err := p.handshake(props); err != nil {
		return nil, err
	}

	return p, nil
}

// NewConnPipeIPC allocates a new Pipe using the IPC exchange protocol.
func NewConnPipeIPC(c net.Conn, sock Socket, props ...interface{}) (Pipe, error) {
	p := &connipc{conn: conn{c: c, proto: sock.GetProtocol(), sock: sock}}

	if err := p.handshake(props); err != nil {
		return nil, err
	}

	return p, nil
}

func (p *connipc) Send(msg *Message) error {

	l := uint64(len(msg.Header) + len(msg.Body))
	one := [1]byte{1}
	var err error

	// send length header
	if _, err = p.c.Write(one[:]); err != nil {
		return err
	}
	if err = binary.Write(p.c, binary.BigEndian, l); err != nil {
		return err
	}
	if _, err = p.c.Write(msg.Header); err != nil {
		return err
	}
	// hope this works
	if _, err = p.c.Write(msg.Body); err != nil {
		return err
	}
	msg.Free()
	return nil
}

func (p *connipc) Recv() (*Message, error) {

	var sz int64
	var err error
	var msg *Message
	var one [1]byte

	if _, err = p.c.Read(one[:]); err != nil {
		return nil, err
	}
	if err = binary.Read(p.c, binary.BigEndian, &sz); err != nil {
		return nil, err
	}

	// Limit messages to the maximum receive value, if not
	// unlimited.  This avoids a potential denaial of service.
	if sz < 0 || (p.maxrx > 0 && sz > p.maxrx) {
		return nil, ErrTooLong
	}
	msg = NewMessage(int(sz))
	msg.Body = msg.Body[0:sz]
	if _, err = io.ReadFull(p.c, msg.Body); err != nil {
		msg.Free()
		return nil, err
	}
	return msg, nil
}

// connHeader is exchanged during the initial handshake.
type connHeader struct {
	Zero    byte // must be zero
	S       byte // 'S'
	P       byte // 'P'
	Version byte // only zero at present
	Proto   uint16
	Rsvd    uint16 // always zero at present
}

// handshake establishes an SP connection between peers.  Both sides must
// send the header, then both sides must wait for the peer's header.
// As a side effect, the peer's protocol number is stored in the conn.
// Also, various properties are initialized.
func (p *conn) handshake(props []interface{}) error {
	var err error

	p.props = make(map[string]interface{})
	p.props[PropLocalAddr] = p.c.LocalAddr()
	p.props[PropRemoteAddr] = p.c.RemoteAddr()

	for len(props) >= 2 {
		switch name := props[0].(type) {
		case string:
			p.props[name] = props[1]
		default:
			return ErrBadProperty
		}
		props = props[2:]
	}

	if v, e := p.sock.GetOption(OptionMaxRecvSize); e == nil {
		// socket guarantees this is an integer
		p.maxrx = int64(v.(int))
	}

	h := connHeader{S: 'S', P: 'P', Proto: p.proto.Number()}
	if err = binary.Write(p.c, binary.BigEndian, &h); err != nil {
		return err
	}
	if err = binary.Read(p.c, binary.BigEndian, &h); err != nil {
		p.c.Close()
		return err
	}
	if h.Zero != 0 || h.S != 'S' || h.P != 'P' || h.Rsvd != 0 {
		p.c.Close()
		return ErrBadHeader
	}
	// The only version number we support at present is "0", at offset 3.
	if h.Version != 0 {
		p.c.Close()
		return ErrBadVersion
	}

	// The protocol number lives as 16-bits (big-endian) at offset 4.
	if h.Proto != p.proto.PeerNumber() {
		p.c.Close()
		return ErrBadProto
	}
	p.open = true
	return nil
}
<|endoftext|>


<commit_before>package proxy

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
)

// Conn adds hook points into "database/sql/driver".Conn.
type Conn struct {
	Conn  driver.Conn
	Proxy *Proxy
}

// Ping verifies a connection to the database is still alive.
// It will trigger PrePing, Ping, PostPing hooks.
//
// If the original connection does not satisfy "database/sql/driver".Pinger, it does nothing.
func (conn *Conn) Ping(c context.Context) error {
	var err error
	var ctx interface{}
	hooks := conn.Proxy.getHooks(c)

	if hooks != nil {
		defer func() { hooks.postPing(c, ctx, conn, err) }()
		if ctx, err = hooks.prePing(c, conn); err != nil {
			return err
		}
	}

	if p, ok := conn.Conn.(driver.Pinger); ok {
		err = p.Ping(c)
		if err != nil {
			return err
		}
	}

	if hooks != nil {
		err = hooks.ping(c, ctx, conn)
	}
	return err
}

// Prepare returns a prepared statement which is wrapped by Stmt.
// NOT SUPPORTED: use PrepareContext instead
func (conn *Conn) Prepare(query string) (driver.Stmt, error) {
	panic("not supported")
}

// PrepareContext returns a prepared statement which is wrapped by Stmt.
func (conn *Conn) PrepareContext(c context.Context, query string) (driver.Stmt, error) {
	var ctx interface{}
	var stmt = &Stmt{
		QueryString: query,
		Proxy:       conn.Proxy,
		Conn:        conn,
	}
	var err error
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postPrepare(c, ctx, stmt, err) }()
		if ctx, err = hooks.prePrepare(c, stmt); err != nil {
			return nil, err
		}
	}

	if connCtx, ok := conn.Conn.(driver.ConnPrepareContext); ok {
		stmt.Stmt, err = connCtx.PrepareContext(c, stmt.QueryString)
	} else {
		stmt.Stmt, err = conn.Conn.Prepare(stmt.QueryString)
		if err == nil {
			select {
			default:
			case <-c.Done():
				stmt.Stmt.Close()
				return nil, c.Err()
			}
		}
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.prepare(c, ctx, stmt); err != nil {
			return nil, err
		}
	}
	return stmt, nil
}

// Close calls the original Close method.
func (conn *Conn) Close() error {
	ctx := context.Background()
	var err error
	var myctx interface{}

	if hooks := conn.Proxy.hooks; hooks != nil {
		defer func() { hooks.postClose(ctx, myctx, conn, err) }()
		if myctx, err = hooks.preClose(ctx, conn); err != nil {
			return err
		}
	}

	err = conn.Conn.Close()
	if err != nil {
		return err
	}

	if hooks := conn.Proxy.hooks; hooks != nil {
		err = hooks.close(ctx, myctx, conn)
	}
	return err
}

// Begin starts and returns a new transaction which is wrapped by Tx.
// It will trigger PreBegin, Begin, PostBegin hooks.
// NOT SUPPORTED: use BeginContext instead
func (conn *Conn) Begin() (driver.Tx, error) {
	panic("not supported")
}

// BeginTx starts and returns a new transaction which is wrapped by Tx.
// It will trigger PreBegin, Begin, PostBegin hooks.
func (conn *Conn) BeginTx(c context.Context, opts driver.TxOptions) (driver.Tx, error) {
	// set the hooks.
	var err error
	var ctx interface{}
	var tx driver.Tx
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postBegin(c, ctx, conn, err) }()
		if ctx, err = hooks.preBegin(c, conn); err != nil {
			return nil, err
		}
	}

	// call the original method.
	if connCtx, ok := conn.Conn.(driver.ConnBeginTx); ok {
		tx, err = connCtx.BeginTx(c, opts)
	} else {
		if c.Done() != context.Background().Done() {
			// the original driver does not support non-default transaction options.
			// so return error if non-default transaction is requested.
			if opts.Isolation != driver.IsolationLevel(sql.LevelDefault) {
				return nil, errors.New("proxy: driver does not support non-default isolation level")
			}
			if opts.ReadOnly {
				return nil, errors.New("proxy: driver does not support read-only transactions")
			}
		}
		tx, err = conn.Conn.Begin()
		if err == nil {
			// check the context is already done.
			select {
			default:
			case <-c.Done():
				tx.Rollback()
				err = c.Err()
			}
		}
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.begin(c, ctx, conn); err != nil {
			tx.Rollback()
			return nil, err
		}
	}

	return &Tx{
		Tx:    tx,
		Proxy: conn.Proxy,
		Conn:  conn,
		ctx:   c,
	}, nil
}

// Exec calls the original Exec method of the connection.
// It will trigger PreExec, Exec, PostExec hooks.
//
// If the original connection does not satisfy "database/sql/driver".Execer, it return ErrSkip error.
// NOT SUPPORTED: use ExecContext instead
func (conn *Conn) Exec(query string, args []driver.Value) (driver.Result, error) {
	panic("not supported")
}

// ExecContext calls the original Exec method of the connection.
// It will trigger PreExec, Exec, PostExec hooks.
//
// If the original connection does not satisfy "database/sql/driver".Execer, it return ErrSkip error.
func (conn *Conn) ExecContext(c context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	execer, ok := conn.Conn.(driver.Execer)
	if !ok {
		return nil, driver.ErrSkip
	}

	// set the hooks.
	var stmt = &Stmt{
		QueryString: query,
		Proxy:       conn.Proxy,
		Conn:        conn,
	}
	var ctx interface{}
	var err error
	var result driver.Result
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postExec(c, ctx, stmt, args, result, err) }()
		if ctx, err = hooks.preExec(c, stmt, args); err != nil {
			return nil, err
		}
	}

	// call the original method.
	if execerCtx, ok := execer.(driver.ExecerContext); ok {
		result, err = execerCtx.ExecContext(c, stmt.QueryString, args)
	} else {
		select {
		default:
		case <-c.Done():
			return nil, c.Err()
		}
		dargs, err0 := namedValuesToValues(args)
		if err0 != nil {
			return nil, err0
		}
		result, err = execer.Exec(stmt.QueryString, dargs)
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.exec(c, ctx, stmt, args, result); err != nil {
			return nil, err
		}
	}

	return result, nil
}

// Query executes a query that may return rows.
// It wil trigger PreQuery, Query, PostQuery hooks.
//
// If the original connection does not satisfy "database/sql/driver".Queryer, it return ErrSkip error.
// NOT SUPPORTED: use QueryContext instead
func (conn *Conn) Query(query string, args []driver.Value) (driver.Rows, error) {
	panic("not supported")
}

// QueryContext executes a query that may return rows.
// It wil trigger PreQuery, Query, PostQuery hooks.
//
// If the original connection does not satisfy "database/sql/driver".Queryer, it return ErrSkip error.
func (conn *Conn) QueryContext(c context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := conn.Conn.(driver.Queryer)
	if !ok {
		return nil, driver.ErrSkip
	}

	var stmt = &Stmt{
		QueryString: query,
		Proxy:       conn.Proxy,
		Conn:        conn,
	}
	var ctx interface{}
	var err error
	var rows driver.Rows
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postQuery(c, ctx, stmt, args, rows, err) }()
		if ctx, err = hooks.preQuery(c, stmt, args); err != nil {
			return nil, err
		}
	}

	// call the original method.
	if queryerCtx, ok := conn.Conn.(driver.QueryerContext); ok {
		rows, err = queryerCtx.QueryContext(c, stmt.QueryString, args)
	} else {
		select {
		default:
		case <-c.Done():
			return nil, c.Err()
		}
		dargs, err0 := namedValuesToValues(args)
		if err0 != nil {
			return nil, err0
		}
		rows, err = queryer.Query(stmt.QueryString, dargs)
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.query(c, ctx, stmt, args, rows); err != nil {
			rows.Close()
			return nil, err
		}
	}

	return rows, nil
}

// copied from sql/driver/convert.go
// defaultCheckNamedValue wraps the default ColumnConverter to have the same
// function signature as the CheckNamedValue in the driver.NamedValueChecker
// interface.
func defaultCheckNamedValue(nv *driver.NamedValue) (err error) {
	nv.Value, err = driver.DefaultParameterConverter.ConvertValue(nv.Value)
	return err
}

// CheckNamedValue for implementing driver.NamedValueChecker
// This function may be unnecessary because `proxy.Stmt` already implements `NamedValueChecker`,
// but it is implemented just in case.
func (conn *Conn) CheckNamedValue(nv *driver.NamedValue) (err error) {
	if nvc, ok := conn.Conn.(namedValueChecker); ok {
		return nvc.CheckNamedValue(nv)
	}
	// fallback to default
	return defaultCheckNamedValue(nv)
}

// sessionResetter is the same as driver.SessionResetter.
// Copied from database/sql/driver/driver.go for supporting Go 1.9
type sessionResetter interface {
	// ResetSession is called while a connection is in the connection
	// pool. No queries will run on this connection until this method returns.
	//
	// If the connection is bad this should return driver.ErrBadConn to prevent
	// the connection from being returned to the connection pool. Any other
	// error will be discarded.
	ResetSession(ctx context.Context) error
}

// ResetSession resets the state of Conn.
func (conn *Conn) ResetSession(ctx context.Context) error {
	var err error
	var myctx interface{}
	hooks := conn.Proxy.getHooks(ctx)

	if hooks != nil {
		defer func() { hooks.postResetSession(ctx, myctx, conn, err) }()
		if myctx, err = hooks.preResetSession(ctx, conn); err != nil {
			return err
		}
	}

	if sr, ok := conn.Conn.(sessionResetter); ok {
		err = sr.ResetSession(ctx)
		if err != nil {
			return err
		}
	}

	if hooks != nil {
		err = hooks.resetSession(ctx, myctx, conn)
	}
	return err
}

// the same as driver.Validator.
// Copied from database/sql/driver/driver.go for supporting Go 1.15
type validator interface {
	// IsValid is called prior to placing the connection into the
	// connection pool. The connection will be discarded if false is returned.
	IsValid() bool
}

// IsValid implements driver.Validator.
// It calls the IsValid method of the original connection.
// If the original connection does not satisfy "database/sql/driver".Validator, it always returns true.
func (conn *Conn) IsValid() bool {
	valid := true
	var myctx interface{}
	hooks := conn.Proxy.hooks
	if hooks != nil {
		// Setup PostIsValid. This needs to be a closure like this
		// or otherwise changes to the `ctx` and `conn` parameters
		// within this Open() method does not get applied at the
		// time defer is fired
		defer func() { hooks.postIsValid(myctx, conn, valid) }()
		var err error
		if myctx, err = hooks.preIsValid(conn); err != nil {
			valid = false
			return valid
		}
	}

	if v, ok := conn.Conn.(validator); ok {
		valid = v.IsValid()
	}
	if valid && hooks != nil {
		err := hooks.isValid(myctx, conn)
		valid = err == nil
	}
	return valid
}
<commit_msg>support ExecerContext interface<commit_after>package proxy

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
)

// Conn adds hook points into "database/sql/driver".Conn.
type Conn struct {
	Conn  driver.Conn
	Proxy *Proxy
}

// Ping verifies a connection to the database is still alive.
// It will trigger PrePing, Ping, PostPing hooks.
//
// If the original connection does not satisfy "database/sql/driver".Pinger, it does nothing.
func (conn *Conn) Ping(c context.Context) error {
	var err error
	var ctx interface{}
	hooks := conn.Proxy.getHooks(c)

	if hooks != nil {
		defer func() { hooks.postPing(c, ctx, conn, err) }()
		if ctx, err = hooks.prePing(c, conn); err != nil {
			return err
		}
	}

	if p, ok := conn.Conn.(driver.Pinger); ok {
		err = p.Ping(c)
		if err != nil {
			return err
		}
	}

	if hooks != nil {
		err = hooks.ping(c, ctx, conn)
	}
	return err
}

// Prepare returns a prepared statement which is wrapped by Stmt.
// NOT SUPPORTED: use PrepareContext instead
func (conn *Conn) Prepare(query string) (driver.Stmt, error) {
	panic("not supported")
}

// PrepareContext returns a prepared statement which is wrapped by Stmt.
func (conn *Conn) PrepareContext(c context.Context, query string) (driver.Stmt, error) {
	var ctx interface{}
	var stmt = &Stmt{
		QueryString: query,
		Proxy:       conn.Proxy,
		Conn:        conn,
	}
	var err error
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postPrepare(c, ctx, stmt, err) }()
		if ctx, err = hooks.prePrepare(c, stmt); err != nil {
			return nil, err
		}
	}

	if connCtx, ok := conn.Conn.(driver.ConnPrepareContext); ok {
		stmt.Stmt, err = connCtx.PrepareContext(c, stmt.QueryString)
	} else {
		stmt.Stmt, err = conn.Conn.Prepare(stmt.QueryString)
		if err == nil {
			select {
			default:
			case <-c.Done():
				stmt.Stmt.Close()
				return nil, c.Err()
			}
		}
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.prepare(c, ctx, stmt); err != nil {
			return nil, err
		}
	}
	return stmt, nil
}

// Close calls the original Close method.
func (conn *Conn) Close() error {
	ctx := context.Background()
	var err error
	var myctx interface{}

	if hooks := conn.Proxy.hooks; hooks != nil {
		defer func() { hooks.postClose(ctx, myctx, conn, err) }()
		if myctx, err = hooks.preClose(ctx, conn); err != nil {
			return err
		}
	}

	err = conn.Conn.Close()
	if err != nil {
		return err
	}

	if hooks := conn.Proxy.hooks; hooks != nil {
		err = hooks.close(ctx, myctx, conn)
	}
	return err
}

// Begin starts and returns a new transaction which is wrapped by Tx.
// It will trigger PreBegin, Begin, PostBegin hooks.
// NOT SUPPORTED: use BeginContext instead
func (conn *Conn) Begin() (driver.Tx, error) {
	panic("not supported")
}

// BeginTx starts and returns a new transaction which is wrapped by Tx.
// It will trigger PreBegin, Begin, PostBegin hooks.
func (conn *Conn) BeginTx(c context.Context, opts driver.TxOptions) (driver.Tx, error) {
	// set the hooks.
	var err error
	var ctx interface{}
	var tx driver.Tx
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postBegin(c, ctx, conn, err) }()
		if ctx, err = hooks.preBegin(c, conn); err != nil {
			return nil, err
		}
	}

	// call the original method.
	if connCtx, ok := conn.Conn.(driver.ConnBeginTx); ok {
		tx, err = connCtx.BeginTx(c, opts)
	} else {
		if c.Done() != context.Background().Done() {
			// the original driver does not support non-default transaction options.
			// so return error if non-default transaction is requested.
			if opts.Isolation != driver.IsolationLevel(sql.LevelDefault) {
				return nil, errors.New("proxy: driver does not support non-default isolation level")
			}
			if opts.ReadOnly {
				return nil, errors.New("proxy: driver does not support read-only transactions")
			}
		}
		tx, err = conn.Conn.Begin()
		if err == nil {
			// check the context is already done.
			select {
			default:
			case <-c.Done():
				tx.Rollback()
				err = c.Err()
			}
		}
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.begin(c, ctx, conn); err != nil {
			tx.Rollback()
			return nil, err
		}
	}

	return &Tx{
		Tx:    tx,
		Proxy: conn.Proxy,
		Conn:  conn,
		ctx:   c,
	}, nil
}

// Exec calls the original Exec method of the connection.
// It will trigger PreExec, Exec, PostExec hooks.
//
// If the original connection does not satisfy "database/sql/driver".Execer, it return ErrSkip error.
// NOT SUPPORTED: use ExecContext instead
func (conn *Conn) Exec(query string, args []driver.Value) (driver.Result, error) {
	panic("not supported")
}

// ExecContext calls the original ExecContext (or Exec as a fallback) method of the connection.
// It will trigger PreExec, PostExec hooks.
//
// If the original connection doesn't satisfy "database/sql/driver".ExecerContext nor "database/sql/driver".Execer, it return ErrSkip error.
func (conn *Conn) ExecContext(c context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	execer, exOk := conn.Conn.(driver.Execer)
	execerCtx, exCtxOk := conn.Conn.(driver.ExecerContext)
	if !exOk && !exCtxOk {
		return nil, driver.ErrSkip
	}

	// set the hooks.
	var stmt = &Stmt{
		QueryString: query,
		Proxy:       conn.Proxy,
		Conn:        conn,
	}
	var ctx interface{}
	var err error
	var result driver.Result
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postExec(c, ctx, stmt, args, result, err) }()
		if ctx, err = hooks.preExec(c, stmt, args); err != nil {
			return nil, err
		}
	}

	// call the original method.
	if execerCtx != nil {
		result, err = execerCtx.ExecContext(c, stmt.QueryString, args)
	} else {
		select {
		default:
		case <-c.Done():
			return nil, c.Err()
		}
		dargs, err0 := namedValuesToValues(args)
		if err0 != nil {
			return nil, err0
		}
		result, err = execer.Exec(stmt.QueryString, dargs)
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.exec(c, ctx, stmt, args, result); err != nil {
			return nil, err
		}
	}

	return result, nil
}

// Query executes a query that may return rows.
// It wil trigger PreQuery, Query, PostQuery hooks.
//
// If the original connection does not satisfy "database/sql/driver".Queryer, it return ErrSkip error.
// NOT SUPPORTED: use QueryContext instead
func (conn *Conn) Query(query string, args []driver.Value) (driver.Rows, error) {
	panic("not supported")
}

// QueryContext executes a query that may return rows.
// It wil trigger PreQuery, Query, PostQuery hooks.
//
// If the original connection does not satisfy "database/sql/driver".Queryer, it return ErrSkip error.
func (conn *Conn) QueryContext(c context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := conn.Conn.(driver.Queryer)
	if !ok {
		return nil, driver.ErrSkip
	}

	var stmt = &Stmt{
		QueryString: query,
		Proxy:       conn.Proxy,
		Conn:        conn,
	}
	var ctx interface{}
	var err error
	var rows driver.Rows
	hooks := conn.Proxy.getHooks(c)
	if hooks != nil {
		defer func() { hooks.postQuery(c, ctx, stmt, args, rows, err) }()
		if ctx, err = hooks.preQuery(c, stmt, args); err != nil {
			return nil, err
		}
	}

	// call the original method.
	if queryerCtx, ok := conn.Conn.(driver.QueryerContext); ok {
		rows, err = queryerCtx.QueryContext(c, stmt.QueryString, args)
	} else {
		select {
		default:
		case <-c.Done():
			return nil, c.Err()
		}
		dargs, err0 := namedValuesToValues(args)
		if err0 != nil {
			return nil, err0
		}
		rows, err = queryer.Query(stmt.QueryString, dargs)
	}
	if err != nil {
		return nil, err
	}

	if hooks != nil {
		if err = hooks.query(c, ctx, stmt, args, rows); err != nil {
			rows.Close()
			return nil, err
		}
	}

	return rows, nil
}

// copied from sql/driver/convert.go
// defaultCheckNamedValue wraps the default ColumnConverter to have the same
// function signature as the CheckNamedValue in the driver.NamedValueChecker
// interface.
func defaultCheckNamedValue(nv *driver.NamedValue) (err error) {
	nv.Value, err = driver.DefaultParameterConverter.ConvertValue(nv.Value)
	return err
}

// CheckNamedValue for implementing driver.NamedValueChecker
// This function may be unnecessary because `proxy.Stmt` already implements `NamedValueChecker`,
// but it is implemented just in case.
func (conn *Conn) CheckNamedValue(nv *driver.NamedValue) (err error) {
	if nvc, ok := conn.Conn.(namedValueChecker); ok {
		return nvc.CheckNamedValue(nv)
	}
	// fallback to default
	return defaultCheckNamedValue(nv)
}

// sessionResetter is the same as driver.SessionResetter.
// Copied from database/sql/driver/driver.go for supporting Go 1.9
type sessionResetter interface {
	// ResetSession is called while a connection is in the connection
	// pool. No queries will run on this connection until this method returns.
	//
	// If the connection is bad this should return driver.ErrBadConn to prevent
	// the connection from being returned to the connection pool. Any other
	// error will be discarded.
	ResetSession(ctx context.Context) error
}

// ResetSession resets the state of Conn.
func (conn *Conn) ResetSession(ctx context.Context) error {
	var err error
	var myctx interface{}
	hooks := conn.Proxy.getHooks(ctx)

	if hooks != nil {
		defer func() { hooks.postResetSession(ctx, myctx, conn, err) }()
		if myctx, err = hooks.preResetSession(ctx, conn); err != nil {
			return err
		}
	}

	if sr, ok := conn.Conn.(sessionResetter); ok {
		err = sr.ResetSession(ctx)
		if err != nil {
			return err
		}
	}

	if hooks != nil {
		err = hooks.resetSession(ctx, myctx, conn)
	}
	return err
}

// the same as driver.Validator.
// Copied from database/sql/driver/driver.go for supporting Go 1.15
type validator interface {
	// IsValid is called prior to placing the connection into the
	// connection pool. The connection will be discarded if false is returned.
	IsValid() bool
}

// IsValid implements driver.Validator.
// It calls the IsValid method of the original connection.
// If the original connection does not satisfy "database/sql/driver".Validator, it always returns true.
func (conn *Conn) IsValid() bool {
	valid := true
	var myctx interface{}
	hooks := conn.Proxy.hooks
	if hooks != nil {
		// Setup PostIsValid. This needs to be a closure like this
		// or otherwise changes to the `ctx` and `conn` parameters
		// within this Open() method does not get applied at the
		// time defer is fired
		defer func() { hooks.postIsValid(myctx, conn, valid) }()
		var err error
		if myctx, err = hooks.preIsValid(conn); err != nil {
			valid = false
			return valid
		}
	}

	if v, ok := conn.Conn.(validator); ok {
		valid = v.IsValid()
	}
	if valid && hooks != nil {
		err := hooks.isValid(myctx, conn)
		valid = err == nil
	}
	return valid
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"

	"k8s.io/klog"
)

var (
	// The directories to load profiles from.
	dirs []string
	poll = flag.Duration("poll", -1, "Poll the directories for new profiles with this interval. Values < 0 disable polling, and exit after loading the profiles.")
)

const (
	parser     = "apparmor_parser"
	apparmorfs = "/sys/kernel/security/apparmor"
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [FLAG]... [PROFILE_DIR]...\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Load the AppArmor profiles specified in the PROFILE_DIR directories.\n")
		flag.PrintDefaults()
	}
	flag.Parse()

	dirs = flag.Args()
	if len(dirs) == 0 {
		klog.Errorf("Must specify at least one directory.")
		flag.Usage()
		os.Exit(1)
	}

	// Check that the required parser binary is found.
	if _, err := exec.LookPath(parser); err != nil {
		klog.Exitf("Required binary %s not found in PATH", parser)
	}

	// Check that loaded profiles can be read.
	if _, err := getLoadedProfiles(); err != nil {
		klog.Exitf("Unable to access apparmor profiles: %v", err)
	}

	if *poll < 0 {
		runOnce()
	} else {
		pollForever()
	}
}

// No polling: run once and exit.
func runOnce() {
	if success, newProfiles := loadNewProfiles(); !success {
		if len(newProfiles) > 0 {
			klog.Exitf("Not all profiles were successfully loaded. Loaded: %v", newProfiles)
		} else {
			klog.Exit("Error loading profiles.")
		}
	} else {
		if len(newProfiles) > 0 {
			klog.Infof("Successfully loaded profiles: %v", newProfiles)
		} else {
			klog.Warning("No new profiles found.")
		}
	}
}

// Poll the directories indefinitely.
func pollForever() {
	klog.V(2).Infof("Polling %s every %s", strings.Join(dirs, ", "), poll.String())
	pollFn := func() {
		_, newProfiles := loadNewProfiles()
		if len(newProfiles) > 0 {
			klog.V(2).Infof("Successfully loaded profiles: %v", newProfiles)
		}
	}
	pollFn() // Run immediately.
	ticker := time.NewTicker(*poll)
	for range ticker.C {
		pollFn()
	}
}

func loadNewProfiles() (success bool, newProfiles []string) {
	loadedProfiles, err := getLoadedProfiles()
	if err != nil {
		klog.Errorf("Error reading loaded profiles: %v", err)
		return false, nil
	}

	success = true
	for _, dir := range dirs {
		infos, err := ioutil.ReadDir(dir)
		if err != nil {
			klog.Warningf("Error reading %s: %v", dir, err)
			success = false
			continue
		}

		for _, info := range infos {
			path := filepath.Join(dir, info.Name())
			// If directory, or symlink to a directory, skip it.
			resolvedInfo, err := resolveSymlink(dir, info)
			if err != nil {
				klog.Warningf("Error resolving symlink: %v", err)
				continue
			}
			if resolvedInfo.IsDir() {
				// Directory listing is shallow.
				klog.V(4).Infof("Skipping directory %s", path)
				continue
			}

			klog.V(4).Infof("Scanning %s for new profiles", path)
			profiles, err := getProfileNames(path)
			if err != nil {
				klog.Warningf("Error reading %s: %v", path, err)
				success = false
				continue
			}

			if unloadedProfiles(loadedProfiles, profiles) {
				if err := loadProfiles(path); err != nil {
					klog.Errorf("Could not load profiles: %v", err)
					success = false
					continue
				}
				// Add new profiles to list of loaded profiles.
				newProfiles = append(newProfiles, profiles...)
				for _, profile := range profiles {
					loadedProfiles[profile] = true
				}
			}
		}
	}

	return success, newProfiles
}

func getProfileNames(path string) ([]string, error) {
	cmd := exec.Command(parser, "--names", path)
	stderr := &bytes.Buffer{}
	cmd.Stderr = stderr
	out, err := cmd.Output()
	if err != nil {
		if stderr.Len() > 0 {
			klog.Warning(stderr.String())
		}
		return nil, fmt.Errorf("error reading profiles from %s: %v", path, err)
	}

	trimmed := strings.TrimSpace(string(out)) // Remove trailing \n
	return strings.Split(trimmed, "\n"), nil
}

func unloadedProfiles(loadedProfiles map[string]bool, profiles []string) bool {
	for _, profile := range profiles {
		if !loadedProfiles[profile] {
			return true
		}
	}
	return false
}

func loadProfiles(path string) error {
	cmd := exec.Command(parser, "--verbose", path)
	stderr := &bytes.Buffer{}
	cmd.Stderr = stderr
	out, err := cmd.Output()
	klog.V(2).Infof("Loading profiles from %s:\n%s", path, out)
	if err != nil {
		if stderr.Len() > 0 {
			klog.Warning(stderr.String())
		}
		return fmt.Errorf("error loading profiles from %s: %v", path, err)
	}
	return nil
}

// If the given fileinfo is a symlink, return the FileInfo of the target. Otherwise, return the
// given fileinfo.
func resolveSymlink(basePath string, info os.FileInfo) (os.FileInfo, error) {
	if info.Mode()&os.ModeSymlink == 0 {
		// Not a symlink.
		return info, nil
	}
	fpath := filepath.Join(basePath, info.Name())
	resolvedName, err := filepath.EvalSymlinks(fpath)
	if err != nil {
		return nil, fmt.Errorf("error resolving symlink %s: %v", fpath, err)
	}
	resolvedInfo, err := os.Stat(resolvedName)
	if err != nil {
		return nil, fmt.Errorf("error calling stat on %s: %v", resolvedName, err)
	}
	return resolvedInfo, nil
}

// TODO: This is copied from k8s.io/kubernetes/pkg/security/apparmor.getLoadedProfiles.
//       Refactor that method to expose it in a reusable way, and delete this version.
func getLoadedProfiles() (map[string]bool, error) {
	profilesPath := path.Join(apparmorfs, "profiles")
	profilesFile, err := os.Open(profilesPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open %s: %v", profilesPath, err)
	}
	defer profilesFile.Close()

	profiles := map[string]bool{}
	scanner := bufio.NewScanner(profilesFile)
	for scanner.Scan() {
		profileName := parseProfileName(scanner.Text())
		if profileName == "" {
			// Unknown line format; skip it.
			continue
		}
		profiles[profileName] = true
	}
	return profiles, nil
}

// The profiles file is formatted with one profile per line, matching a form:
//   namespace://profile-name (mode)
//   profile-name (mode)
// Where mode is {enforce, complain, kill}. The "namespace://" is only included for namespaced
// profiles. For the purposes of Kubernetes, we consider the namespace part of the profile name.
func parseProfileName(profileLine string) string {
	modeIndex := strings.IndexRune(profileLine, '(')
	if modeIndex < 0 {
		return ""
	}
	return strings.TrimSpace(profileLine[:modeIndex])
}
<commit_msg>Initialize flags to support extra flags<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"

	"k8s.io/klog"
)

var (
	// The directories to load profiles from.
	dirs []string
	poll = flag.Duration("poll", -1, "Poll the directories for new profiles with this interval. Values < 0 disable polling, and exit after loading the profiles.")
)

const (
	parser     = "apparmor_parser"
	apparmorfs = "/sys/kernel/security/apparmor"
)

func main() {
	klog.InitFlags(nil)
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [FLAG]... [PROFILE_DIR]...\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Load the AppArmor profiles specified in the PROFILE_DIR directories.\n")
		flag.PrintDefaults()
	}
	flag.Parse()

	dirs = flag.Args()
	if len(dirs) == 0 {
		klog.Errorf("Must specify at least one directory.")
		flag.Usage()
		os.Exit(1)
	}

	// Check that the required parser binary is found.
	if _, err := exec.LookPath(parser); err != nil {
		klog.Exitf("Required binary %s not found in PATH", parser)
	}

	// Check that loaded profiles can be read.
	if _, err := getLoadedProfiles(); err != nil {
		klog.Exitf("Unable to access apparmor profiles: %v", err)
	}

	if *poll < 0 {
		runOnce()
	} else {
		pollForever()
	}
}

// No polling: run once and exit.
func runOnce() {
	if success, newProfiles := loadNewProfiles(); !success {
		if len(newProfiles) > 0 {
			klog.Exitf("Not all profiles were successfully loaded. Loaded: %v", newProfiles)
		} else {
			klog.Exit("Error loading profiles.")
		}
	} else {
		if len(newProfiles) > 0 {
			klog.Infof("Successfully loaded profiles: %v", newProfiles)
		} else {
			klog.Warning("No new profiles found.")
		}
	}
}

// Poll the directories indefinitely.
func pollForever() {
	klog.V(2).Infof("Polling %s every %s", strings.Join(dirs, ", "), poll.String())
	pollFn := func() {
		_, newProfiles := loadNewProfiles()
		if len(newProfiles) > 0 {
			klog.V(2).Infof("Successfully loaded profiles: %v", newProfiles)
		}
	}
	pollFn() // Run immediately.
	ticker := time.NewTicker(*poll)
	for range ticker.C {
		pollFn()
	}
}

func loadNewProfiles() (success bool, newProfiles []string) {
	loadedProfiles, err := getLoadedProfiles()
	if err != nil {
		klog.Errorf("Error reading loaded profiles: %v", err)
		return false, nil
	}

	success = true
	for _, dir := range dirs {
		infos, err := ioutil.ReadDir(dir)
		if err != nil {
			klog.Warningf("Error reading %s: %v", dir, err)
			success = false
			continue
		}

		for _, info := range infos {
			path := filepath.Join(dir, info.Name())
			// If directory, or symlink to a directory, skip it.
			resolvedInfo, err := resolveSymlink(dir, info)
			if err != nil {
				klog.Warningf("Error resolving symlink: %v", err)
				continue
			}
			if resolvedInfo.IsDir() {
				// Directory listing is shallow.
				klog.V(4).Infof("Skipping directory %s", path)
				continue
			}

			klog.V(4).Infof("Scanning %s for new profiles", path)
			profiles, err := getProfileNames(path)
			if err != nil {
				klog.Warningf("Error reading %s: %v", path, err)
				success = false
				continue
			}

			if unloadedProfiles(loadedProfiles, profiles) {
				if err := loadProfiles(path); err != nil {
					klog.Errorf("Could not load profiles: %v", err)
					success = false
					continue
				}
				// Add new profiles to list of loaded profiles.
				newProfiles = append(newProfiles, profiles...)
				for _, profile := range profiles {
					loadedProfiles[profile] = true
				}
			}
		}
	}

	return success, newProfiles
}

func getProfileNames(path string) ([]string, error) {
	cmd := exec.Command(parser, "--names", path)
	stderr := &bytes.Buffer{}
	cmd.Stderr = stderr
	out, err := cmd.Output()
	if err != nil {
		if stderr.Len() > 0 {
			klog.Warning(stderr.String())
		}
		return nil, fmt.Errorf("error reading profiles from %s: %v", path, err)
	}

	trimmed := strings.TrimSpace(string(out)) // Remove trailing \n
	return strings.Split(trimmed, "\n"), nil
}

func unloadedProfiles(loadedProfiles map[string]bool, profiles []string) bool {
	for _, profile := range profiles {
		if !loadedProfiles[profile] {
			return true
		}
	}
	return false
}

func loadProfiles(path string) error {
	cmd := exec.Command(parser, "--verbose", path)
	stderr := &bytes.Buffer{}
	cmd.Stderr = stderr
	out, err := cmd.Output()
	klog.V(2).Infof("Loading profiles from %s:\n%s", path, out)
	if err != nil {
		if stderr.Len() > 0 {
			klog.Warning(stderr.String())
		}
		return fmt.Errorf("error loading profiles from %s: %v", path, err)
	}
	return nil
}

// If the given fileinfo is a symlink, return the FileInfo of the target. Otherwise, return the
// given fileinfo.
func resolveSymlink(basePath string, info os.FileInfo) (os.FileInfo, error) {
	if info.Mode()&os.ModeSymlink == 0 {
		// Not a symlink.
		return info, nil
	}
	fpath := filepath.Join(basePath, info.Name())
	resolvedName, err := filepath.EvalSymlinks(fpath)
	if err != nil {
		return nil, fmt.Errorf("error resolving symlink %s: %v", fpath, err)
	}
	resolvedInfo, err := os.Stat(resolvedName)
	if err != nil {
		return nil, fmt.Errorf("error calling stat on %s: %v", resolvedName, err)
	}
	return resolvedInfo, nil
}

// TODO: This is copied from k8s.io/kubernetes/pkg/security/apparmor.getLoadedProfiles.
//       Refactor that method to expose it in a reusable way, and delete this version.
func getLoadedProfiles() (map[string]bool, error) {
	profilesPath := path.Join(apparmorfs, "profiles")
	profilesFile, err := os.Open(profilesPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open %s: %v", profilesPath, err)
	}
	defer profilesFile.Close()

	profiles := map[string]bool{}
	scanner := bufio.NewScanner(profilesFile)
	for scanner.Scan() {
		profileName := parseProfileName(scanner.Text())
		if profileName == "" {
			// Unknown line format; skip it.
			continue
		}
		profiles[profileName] = true
	}
	return profiles, nil
}

// The profiles file is formatted with one profile per line, matching a form:
//   namespace://profile-name (mode)
//   profile-name (mode)
// Where mode is {enforce, complain, kill}. The "namespace://" is only included for namespaced
// profiles. For the purposes of Kubernetes, we consider the namespace part of the profile name.
func parseProfileName(profileLine string) string {
	modeIndex := strings.IndexRune(profileLine, '(')
	if modeIndex < 0 {
		return ""
	}
	return strings.TrimSpace(profileLine[:modeIndex])
}
<|endoftext|>