<commit_before>package app

import (
	"github.com/Sirupsen/logrus"
	db "github.com/theduke/go-dukedb"

	kit "github.com/theduke/go-appkit"
)

type Registry struct {
	logger          *logrus.Logger
	config          kit.Config
	defaultCache    kit.Cache
	caches          map[string]kit.Cache
	defaultBackend  db.Backend
	backends        map[string]db.Backend
	resources       map[string]kit.Resource
	emailService    kit.EmailService
	fileService     kit.FileService
	resourceService kit.ResourceService
	userService     kit.UserService
	templateEngine  kit.TemplateEngine
	taskService     kit.TaskService
}

// Ensure Registry implements kit.Registry.
var _ kit.Registry = (*Registry)(nil)

func NewRegistry() kit.Registry {
	return &Registry{
		caches:    make(map[string]kit.Cache),
		backends:  make(map[string]db.Backend),
		resources: make(map[string]kit.Resource),
	}
}

func (d *Registry) Logger() *logrus.Logger {
	return d.logger
}

func (d *Registry) SetLogger(l *logrus.Logger) {
	d.logger = l
}

func (d *Registry) Config() kit.Config {
	return d.config
}

func (d *Registry) SetConfig(c kit.Config) {
	d.config = c
}

func (d *Registry) Cache(name string) kit.Cache {
	return d.caches[name]
}

func (d *Registry) DefaultCache() kit.Cache {
	return d.defaultCache
}

func (d *Registry) SetDefaultCache(c kit.Cache) {
	d.defaultCache = c
}

func (d *Registry) Caches() map[string]kit.Cache {
	return d.caches
}

func (d *Registry) AddCache(cache kit.Cache) {
	d.caches[cache.Name()] = cache
	if d.defaultCache == nil {
		d.defaultCache = cache
	}
}

func (d *Registry) SetCaches(caches map[string]kit.Cache) {
	d.caches = caches
}

func (d *Registry) DefaultBackend() db.Backend {
	return d.defaultBackend
}

func (d *Registry) SetDefaultBackend(b db.Backend) {
	d.defaultBackend = b
}

func (d *Registry) Backend(name string) db.Backend {
	return d.backends[name]
}

func (d *Registry) Backends() map[string]db.Backend {
	return d.backends
}

func (d *Registry) AddBackend(b db.Backend) {
	d.backends[b.Name()] = b
	if d.defaultBackend == nil {
		d.defaultBackend = b
	}
}

func (d *Registry) SetBackends(backends map[string]db.Backend) {
	d.backends = backends
}

func (d *Registry) Resource(name string) kit.Resource {
	return d.resources[name]
}

func (d *Registry) Resources() map[string]kit.Resource {
	return d.resources
}

func (d *Registry) AddResource(res kit.Resource) {
	d.resources[res.Collection()] = res
}

func (d *Registry) SetResources(resources map[string]kit.Resource) {
	d.resources = resources
}

func (d *Registry) EmailService() kit.EmailService {
	return d.emailService
}

func (d *Registry) SetEmailService(s kit.EmailService) {
	d.emailService = s
}

func (d *Registry) FileService() kit.FileService {
	return d.fileService
}

func (d *Registry) SetFileService(s kit.FileService) {
	d.fileService = s
}

func (d *Registry) ResourceService() kit.ResourceService {
	return d.resourceService
}

func (d *Registry) SetResourceService(s kit.ResourceService) {
	d.resourceService = s
}

func (d *Registry) UserService() kit.UserService {
	return d.userService
}

func (d *Registry) SetUserService(s kit.UserService) {
	d.userService = s
}

func (d *Registry) TemplateEngine() kit.TemplateEngine {
	return d.templateEngine
}

func (d *Registry) SetTemplateEngine(e kit.TemplateEngine) {
	d.templateEngine = e
}

func (d *Registry) TaskService() kit.TaskService {
	return d.taskService
}

func (d *Registry) SetTaskService(s kit.TaskService) {
	d.taskService = s
}
<commit_msg>updated registry to support app and custom services<commit_after>package app

import (
	"github.com/Sirupsen/logrus"
	db "github.com/theduke/go-dukedb"

	kit "github.com/theduke/go-appkit"
)

type Registry struct {
	app             kit.App
	logger          *logrus.Logger
	config          kit.Config
	defaultCache    kit.Cache
	caches          map[string]kit.Cache
	defaultBackend  db.Backend
	backends        map[string]db.Backend
	resources       map[string]kit.Resource
	emailService    kit.EmailService
	fileService     kit.FileService
	resourceService kit.ResourceService
	userService     kit.UserService
	templateEngine  kit.TemplateEngine
	taskService     kit.TaskService

	values map[string]interface{}
}

// Ensure Registry implements kit.Registry.
var _ kit.Registry = (*Registry)(nil)

func NewRegistry() kit.Registry {
	return &Registry{
		caches:    make(map[string]kit.Cache),
		backends:  make(map[string]db.Backend),
		resources: make(map[string]kit.Resource),
		values:    make(map[string]interface{}),
	}
}

func (d *Registry) App() kit.App {
	return d.app
}

func (d *Registry) SetApp(x kit.App) {
	d.app = x
}

func (d *Registry) Logger() *logrus.Logger {
	return d.logger
}

func (d *Registry) SetLogger(l *logrus.Logger) {
	d.logger = l
}

func (d *Registry) Config() kit.Config {
	return d.config
}

func (d *Registry) SetConfig(c kit.Config) {
	d.config = c
}

func (d *Registry) Cache(name string) kit.Cache {
	return d.caches[name]
}

func (d *Registry) DefaultCache() kit.Cache {
	return d.defaultCache
}

func (d *Registry) SetDefaultCache(c kit.Cache) {
	d.defaultCache = c
}

func (d *Registry) Caches() map[string]kit.Cache {
	return d.caches
}

func (d *Registry) AddCache(cache kit.Cache) {
	d.caches[cache.Name()] = cache
	if d.defaultCache == nil {
		d.defaultCache = cache
	}
}

func (d *Registry) SetCaches(caches map[string]kit.Cache) {
	d.caches = caches
}

func (d *Registry) DefaultBackend() db.Backend {
	return d.defaultBackend
}

func (d *Registry) SetDefaultBackend(b db.Backend) {
	d.defaultBackend = b
}

func (d *Registry) Backend(name string) db.Backend {
	return d.backends[name]
}

func (d *Registry) Backends() map[string]db.Backend {
	return d.backends
}

func (d *Registry) AddBackend(b db.Backend) {
	d.backends[b.Name()] = b
	if d.defaultBackend == nil {
		d.defaultBackend = b
	}
}

func (d *Registry) SetBackends(backends map[string]db.Backend) {
	d.backends = backends
}

func (d *Registry) Resource(name string) kit.Resource {
	return d.resources[name]
}

func (d *Registry) Resources() map[string]kit.Resource {
	return d.resources
}

func (d *Registry) AddResource(res kit.Resource) {
	d.resources[res.Collection()] = res
}

func (d *Registry) SetResources(resources map[string]kit.Resource) {
	d.resources = resources
}

func (d *Registry) EmailService() kit.EmailService {
	return d.emailService
}

func (d *Registry) SetEmailService(s kit.EmailService) {
	d.emailService = s
}

func (d *Registry) FileService() kit.FileService {
	return d.fileService
}

func (d *Registry) SetFileService(s kit.FileService) {
	d.fileService = s
}

func (d *Registry) ResourceService() kit.ResourceService {
	return d.resourceService
}

func (d *Registry) SetResourceService(s kit.ResourceService) {
	d.resourceService = s
}

func (d *Registry) UserService() kit.UserService {
	return d.userService
}

func (d *Registry) SetUserService(s kit.UserService) {
	d.userService = s
}

func (d *Registry) TemplateEngine() kit.TemplateEngine {
	return d.templateEngine
}

func (d *Registry) SetTemplateEngine(e kit.TemplateEngine) {
	d.templateEngine = e
}

func (d *Registry) TaskService() kit.TaskService {
	return d.taskService
}

func (d *Registry) SetTaskService(s kit.TaskService) {
	d.taskService = s
}

func (d *Registry) Get(name string) interface{} {
	return d.values[name]
}

func (d *Registry) Set(name string, val interface{}) {
	d.values[name] = val
}
<|endoftext|>


<commit_before>package transit

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/hashicorp/vault/logical"
)

const (
	shared    = false
	exclusive = true
)

type lockManager struct {
	// A lock for each named key
	locks map[string]*sync.RWMutex

	// A mutex for the map itself
	locksMutex sync.RWMutex

	// If caching is enabled, the map of name to in-memory policy cache
	cache map[string]*Policy

	// Used for global locking, and as the cache map mutex
	globalMutex sync.RWMutex
}

func newLockManager(cacheDisabled bool) *lockManager {
	lm := &lockManager{
		locks: map[string]*sync.RWMutex{},
	}
	if !cacheDisabled {
		lm.cache = map[string]*Policy{}
	}
	return lm
}

func (lm *lockManager) CacheActive() bool {
	return lm.cache != nil
}

func (lm *lockManager) lockAll(name string) {
	lm.globalMutex.Lock()
	lm.lockPolicy(name, exclusive)
}

func (lm *lockManager) UnlockAll(name string) {
	lm.UnlockPolicy(name, exclusive)
	lm.globalMutex.Unlock()
}

func (lm *lockManager) lockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	if lock != nil {
		// We want to give this up before locking the lock, but it's safe --
		// the only time we ever write to a value in this map is the first time
		// we access the value, so it won't be changing out from under us
		lm.locksMutex.RUnlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lm.locksMutex.RUnlock()
	lm.locksMutex.Lock()

	// Don't defer the unlock call because if we get a valid lock below we want
	// to release the lock mutex right away to avoid the possibility of
	// deadlock by trying to grab the second lock

	// Check to make sure it hasn't been created since
	lock = lm.locks[name]
	if lock != nil {
		lm.locksMutex.Unlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lock = &sync.RWMutex{}
	lm.locks[name] = lock
	lm.locksMutex.Unlock()
	if writeLock {
		lock.Lock()
	} else {
		lock.RLock()
	}
}

func (lm *lockManager) UnlockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	lm.locksMutex.RUnlock()

	if writeLock {
		lock.Unlock()
	} else {
		lock.RUnlock()
	}
}

func (lm *lockManager) GetPolicy(storage logical.Storage, name string) (*Policy, bool, error) {
	p, lt, _, err := lm.getPolicyCommon(storage, name, false, false)
	return p, lt, err
}

func (lm *lockManager) GetPolicyUpsert(storage logical.Storage, name string, derived bool) (*Policy, bool, bool, error) {
	return lm.getPolicyCommon(storage, name, true, derived)
}

// When the function returns, a lock will be held on the policy if err == nil.
// The type of lock will be indicated by the return value. It is the caller's
// responsibility to unlock.
func (lm *lockManager) getPolicyCommon(storage logical.Storage, name string, upsert, derived bool) (p *Policy, lockType bool, upserted bool, err error) {
	// If we are using a cache, lock it now to avoid having to do really
	// complicated lock juggling as we call various functions. We'll also defer
	// the store into the cache.
	lockType = shared
	lm.lockPolicy(name, shared)

	if lm.CacheActive() {
		lm.globalMutex.RLock()
		p = lm.cache[name]
		if p != nil {
			lm.globalMutex.RUnlock()
			return
		}
		lm.globalMutex.RUnlock()

		// When we return, since we didn't have the policy in the cache, if
		// there was no error, write the value in.
		defer func() {
			lm.globalMutex.Lock()
			defer lm.globalMutex.Unlock()
			// Make sure a policy didn't appear. If so, it will only be set if
			// there was no error, so now just clear the error and return that
			// policy.
			exp := lm.cache[name]
			if exp != nil {
				upserted = false
				err = nil
				p = exp
				return
			}

			if err == nil {
				lm.cache[name] = p
			}
		}()
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		lm.UnlockPolicy(name, shared)
		return
	}

	if p == nil {
		if !upsert {
			lm.UnlockPolicy(name, shared)
			return
		}

		// Get an exlusive lock; on success, check again to ensure that no
		// policy exists. Note that if we are using a cache we will already be
		// serializing this entire code path and it's currently the only one
		// that generates policies, so we don't need to check the cache here;
		// simply checking the disk again is sufficient.
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.lockPolicy(name, exclusive)

		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			return
		}
		if p != nil {
			return
		}

		upserted = true

		p = &Policy{
			Name:       name,
			CipherMode: "aes-gcm",
			Derived:    derived,
		}
		if derived {
			p.KDFMode = kdfMode
		}

		err = p.rotate(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			p = nil
		}

		// We don't need to worry about upgrading since it will be a new policy
		return
	}

	if p.needsUpgrade() {
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.lockPolicy(name, exclusive)

		// Reload the policy with the write lock to ensure we still need the upgrade
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			return
		}
		if p == nil {
			defer lm.UnlockPolicy(name, exclusive)
			err = fmt.Errorf("error reloading policy for upgrade")
			return
		}

		if !p.needsUpgrade() {
			// Already happened, return the newly loaded policy
			return
		}

		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) DeletePolicy(storage logical.Storage, name string) error {
	lm.lockAll(name)
	defer lm.UnlockAll(name)

	var p *Policy
	var err error

	if lm.CacheActive() {
		p = lm.cache[name]
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	} else {
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			return err
		}
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	}

	if !p.DeletionAllowed {
		return fmt.Errorf("deletion is not allowed for this policy")
	}

	err = storage.Delete("policy/" + name)
	if err != nil {
		return fmt.Errorf("error deleting policy %s: %s", name, err)
	}

	err = storage.Delete("archive/" + name)
	if err != nil {
		return fmt.Errorf("error deleting archive %s: %s", name, err)
	}

	if lm.CacheActive() {
		delete(lm.cache, name)
	}

	return nil
}

// When this function returns it's the responsibility of the caller to call
// UnlockPolicy if err is nil and policy is not nil
func (lm *lockManager) RefreshPolicy(storage logical.Storage, name string) (p *Policy, err error) {
	lm.lockPolicy(name, exclusive)

	if lm.CacheActive() {
		p = lm.cache[name]
		if p != nil {
			return
		}
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
		return
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		defer lm.UnlockPolicy(name, exclusive)
		return
	}

	if p == nil {
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
	}

	if p.needsUpgrade() {
		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) getStoredPolicy(storage logical.Storage, name string) (*Policy, error) {
	// Check if the policy already exists
	raw, err := storage.Get("policy/" + name)
	if err != nil {
		return nil, err
	}
	if raw == nil {
		return nil, nil
	}

	// Decode the policy
	policy := &Policy{
		Keys: KeyEntryMap{},
	}
	err = json.Unmarshal(raw.Value, policy)
	if err != nil {
		return nil, err
	}

	return policy, nil
}
<commit_msg>Remove some deferring<commit_after>package transit

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/hashicorp/vault/logical"
)

const (
	shared    = false
	exclusive = true
)

type lockManager struct {
	// A lock for each named key
	locks map[string]*sync.RWMutex

	// A mutex for the map itself
	locksMutex sync.RWMutex

	// If caching is enabled, the map of name to in-memory policy cache
	cache map[string]*Policy

	// Used for global locking, and as the cache map mutex
	globalMutex sync.RWMutex
}

func newLockManager(cacheDisabled bool) *lockManager {
	lm := &lockManager{
		locks: map[string]*sync.RWMutex{},
	}
	if !cacheDisabled {
		lm.cache = map[string]*Policy{}
	}
	return lm
}

func (lm *lockManager) CacheActive() bool {
	return lm.cache != nil
}

func (lm *lockManager) lockAll(name string) {
	lm.globalMutex.Lock()
	lm.lockPolicy(name, exclusive)
}

func (lm *lockManager) UnlockAll(name string) {
	lm.UnlockPolicy(name, exclusive)
	lm.globalMutex.Unlock()
}

func (lm *lockManager) lockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	if lock != nil {
		// We want to give this up before locking the lock, but it's safe --
		// the only time we ever write to a value in this map is the first time
		// we access the value, so it won't be changing out from under us
		lm.locksMutex.RUnlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lm.locksMutex.RUnlock()
	lm.locksMutex.Lock()

	// Don't defer the unlock call because if we get a valid lock below we want
	// to release the lock mutex right away to avoid the possibility of
	// deadlock by trying to grab the second lock

	// Check to make sure it hasn't been created since
	lock = lm.locks[name]
	if lock != nil {
		lm.locksMutex.Unlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lock = &sync.RWMutex{}
	lm.locks[name] = lock
	lm.locksMutex.Unlock()
	if writeLock {
		lock.Lock()
	} else {
		lock.RLock()
	}
}

func (lm *lockManager) UnlockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	lm.locksMutex.RUnlock()

	if writeLock {
		lock.Unlock()
	} else {
		lock.RUnlock()
	}
}

func (lm *lockManager) GetPolicy(storage logical.Storage, name string) (*Policy, bool, error) {
	p, lt, _, err := lm.getPolicyCommon(storage, name, false, false)
	return p, lt, err
}

func (lm *lockManager) GetPolicyUpsert(storage logical.Storage, name string, derived bool) (*Policy, bool, bool, error) {
	return lm.getPolicyCommon(storage, name, true, derived)
}

// When the function returns, a lock will be held on the policy if err == nil.
// The type of lock will be indicated by the return value. It is the caller's
// responsibility to unlock.
func (lm *lockManager) getPolicyCommon(storage logical.Storage, name string, upsert, derived bool) (p *Policy, lockType bool, upserted bool, err error) {
	// If we are using a cache, lock it now to avoid having to do really
	// complicated lock juggling as we call various functions. We'll also defer
	// the store into the cache.
	lockType = shared
	lm.lockPolicy(name, shared)

	if lm.CacheActive() {
		lm.globalMutex.RLock()
		p = lm.cache[name]
		if p != nil {
			lm.globalMutex.RUnlock()
			return
		}
		lm.globalMutex.RUnlock()

		// When we return, since we didn't have the policy in the cache, if
		// there was no error, write the value in.
		defer func() {
			lm.globalMutex.Lock()
			defer lm.globalMutex.Unlock()
			// Make sure a policy didn't appear. If so, it will only be set if
			// there was no error, so now just clear the error and return that
			// policy.
			exp := lm.cache[name]
			if exp != nil {
				upserted = false
				err = nil
				p = exp
				return
			}

			if err == nil {
				lm.cache[name] = p
			}
		}()
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		lm.UnlockPolicy(name, shared)
		return
	}

	if p == nil {
		if !upsert {
			lm.UnlockPolicy(name, shared)
			return
		}

		// Get an exlusive lock; on success, check again to ensure that no
		// policy exists. Note that if we are using a cache we will already be
		// serializing this entire code path and it's currently the only one
		// that generates policies, so we don't need to check the cache here;
		// simply checking the disk again is sufficient.
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.lockPolicy(name, exclusive)

		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			return
		}
		if p != nil {
			return
		}

		upserted = true

		p = &Policy{
			Name:       name,
			CipherMode: "aes-gcm",
			Derived:    derived,
		}
		if derived {
			p.KDFMode = kdfMode
		}

		err = p.rotate(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			p = nil
		}

		// We don't need to worry about upgrading since it will be a new policy
		return
	}

	if p.needsUpgrade() {
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.lockPolicy(name, exclusive)

		// Reload the policy with the write lock to ensure we still need the upgrade
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			lm.UnlockPolicy(name, exclusive)
			return
		}
		if p == nil {
			err = fmt.Errorf("error reloading policy for upgrade")
			lm.UnlockPolicy(name, exclusive)
			return
		}

		if !p.needsUpgrade() {
			// Already happened, return the newly loaded policy
			return
		}

		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) DeletePolicy(storage logical.Storage, name string) error {
	lm.lockAll(name)
	defer lm.UnlockAll(name)

	var p *Policy
	var err error

	if lm.CacheActive() {
		p = lm.cache[name]
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	} else {
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			return err
		}
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	}

	if !p.DeletionAllowed {
		return fmt.Errorf("deletion is not allowed for this policy")
	}

	err = storage.Delete("policy/" + name)
	if err != nil {
		return fmt.Errorf("error deleting policy %s: %s", name, err)
	}

	err = storage.Delete("archive/" + name)
	if err != nil {
		return fmt.Errorf("error deleting archive %s: %s", name, err)
	}

	if lm.CacheActive() {
		delete(lm.cache, name)
	}

	return nil
}

// When this function returns it's the responsibility of the caller to call
// UnlockPolicy if err is nil and policy is not nil
func (lm *lockManager) RefreshPolicy(storage logical.Storage, name string) (p *Policy, err error) {
	lm.lockPolicy(name, exclusive)

	if lm.CacheActive() {
		p = lm.cache[name]
		if p != nil {
			return
		}
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
		return
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		lm.UnlockPolicy(name, exclusive)
		return
	}

	if p == nil {
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
	}

	if p.needsUpgrade() {
		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) getStoredPolicy(storage logical.Storage, name string) (*Policy, error) {
	// Check if the policy already exists
	raw, err := storage.Get("policy/" + name)
	if err != nil {
		return nil, err
	}
	if raw == nil {
		return nil, nil
	}

	// Decode the policy
	policy := &Policy{
		Keys: KeyEntryMap{},
	}
	err = json.Unmarshal(raw.Value, policy)
	if err != nil {
		return nil, err
	}

	return policy, nil
}
<|endoftext|>


<commit_before>/*
Package github_flavored_markdown provides a GitHub Flavored Markdown renderer
with fenced code block highlighting, clickable header anchor links.

The functionality should be equivalent to the GitHub Markdown API endpoint specified at
https://developer.github.com/v3/markdown/#render-a-markdown-document-in-raw-mode, except
the rendering is performed locally.

See example below for how to generate a complete HTML page, including CSS styles.
*/
package github_flavored_markdown

import (
	"bytes"
	"fmt"
	"html"
	"io"
	"regexp"
	"sort"
	"strings"
	"text/template"

	"github.com/microcosm-cc/bluemonday"
	"github.com/russross/blackfriday"
	"github.com/shurcooL/sanitized_anchor_name"
	"github.com/shurcooL/go/highlight_go"
	"github.com/shurcooL/go/u/u7"
	"github.com/sourcegraph/annotate"
	"github.com/sourcegraph/syntaxhighlight"
)

// Markdown renders GitHub Flavored Markdown text.
func Markdown(text []byte) []byte {
	htmlFlags := 0
	renderer := &renderer{Html: blackfriday.HtmlRenderer(htmlFlags, "", "").(*blackfriday.Html)}

	// Parser extensions for GitHub Flavored Markdown.
	extensions := 0
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS
	//extensions |= blackfriday.EXTENSION_HARD_LINE_BREAK

	unsanitized := blackfriday.Markdown(text, renderer, extensions)

	// GitHub Flavored Markdown-like sanitization policy.
	p := bluemonday.UGCPolicy()
	p.AllowAttrs("class").Matching(bluemonday.SpaceSeparatedTokens).OnElements("div", "span")
	p.AllowAttrs("class", "name").Matching(bluemonday.SpaceSeparatedTokens).OnElements("a")
	p.AllowAttrs("rel").Matching(regexp.MustCompile(`^nofollow$`)).OnElements("a")
	p.AllowAttrs("aria-hidden").Matching(regexp.MustCompile(`^true$`)).OnElements("a")
	p.AllowDataURIImages()

	return p.SanitizeBytes(unsanitized)
}

type renderer struct {
	*blackfriday.Html
}

// GitHub Flavored Markdown header with clickable and hidden anchor.
func (_ *renderer) Header(out *bytes.Buffer, text func() bool, level int, _ string) {
	marker := out.Len()
	doubleSpace(out)

	if !text() {
		out.Truncate(marker)
		return
	}

	textString := out.String()[marker:]
	out.Truncate(marker)

	anchorName := sanitized_anchor_name.Create(html.UnescapeString(textString))

	out.WriteString(fmt.Sprintf(`<h%d><a name="%s" class="anchor" href="#%s" rel="nofollow" aria-hidden="true"><span class="octicon octicon-link"></span></a>`, level, anchorName, anchorName))
	out.WriteString(textString)
	out.WriteString(fmt.Sprintf("</h%d>\n", level))
}

// TODO: Clean up and improve this code.
// GitHub Flavored Markdown fenced code block with highlighting.
func (_ *renderer) BlockCode(out *bytes.Buffer, text []byte, lang string) {
	doubleSpace(out)

	// parse out the language name
	count := 0
	for _, elt := range strings.Fields(lang) {
		if elt[0] == '.' {
			elt = elt[1:]
		}
		if len(elt) == 0 {
			continue
		}
		out.WriteString(`<div class="highlight highlight-`)
		attrEscape(out, []byte(elt))
		lang = elt
		out.WriteString(`"><pre>`)
		count++
		break
	}

	if count == 0 {
		out.WriteString("
")
	}

	if highlightedCode, ok := highlightCode(text, lang); ok {
		out.Write(highlightedCode)
	} else {
		attrEscape(out, text)
	}

	if count == 0 {
		out.WriteString("
\n")
	} else {
		out.WriteString("

\n")
	}
}

var gfmHtmlConfig = syntaxhighlight.HTMLConfig{
	String:        "s",
	Keyword:       "k",
	Comment:       "c",
	Type:          "n",
	Literal:       "o",
	Punctuation:   "p",
	Plaintext:     "n",
	Tag:           "tag",
	HTMLTag:       "htm",
	HTMLAttrName:  "atn",
	HTMLAttrValue: "atv",
	Decimal:       "m",
}

// TODO: Support highlighting for more languages.
func highlightCode(src []byte, lang string) (highlightedCode []byte, ok bool) {
	switch lang {
	case "Go":
		var buf bytes.Buffer
		err := highlight_go.Print(src, &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "Go-old":
		var buf bytes.Buffer
		err := syntaxhighlight.Print(syntaxhighlight.NewScanner(src), &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "diff":
		switch 2 {
		default:
			var buf bytes.Buffer
			err := u7.Print(u7.NewScanner(src), &buf)
			if err != nil {
				return nil, false
			}
			return buf.Bytes(), true
		case 1:
			lines := bytes.Split(src, []byte("\n"))
			return bytes.Join(lines, []byte("\n")), true
		case 2:
			anns, err := u7.Annotate(src)
			if err != nil {
				return nil, false
			}

			lines := bytes.Split(src, []byte("\n"))
			lineStarts := make([]int, len(lines))
			var offset int
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				lineStarts[lineIndex] = offset
				offset += len(lines[lineIndex]) + 1
			}

			lastDel, lastIns := -1, -1
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				var lineFirstChar byte
				if len(lines[lineIndex]) > 0 {
					lineFirstChar = lines[lineIndex][0]
				}
				switch lineFirstChar {
				case '+':
					if lastIns == -1 {
						lastIns = lineIndex
					}
				case '-':
					if lastDel == -1 {
						lastDel = lineIndex
					}
				default:
					if lastDel != -1 || lastIns != -1 {
						if lastDel == -1 {
							lastDel = lastIns
						} else if lastIns == -1 {
							lastIns = lineIndex
						}

						beginOffsetLeft := lineStarts[lastDel]
						endOffsetLeft := lineStarts[lastIns]
						beginOffsetRight := lineStarts[lastIns]
						endOffsetRight := lineStarts[lineIndex]

						anns = append(anns, &annotate.Annotation{Start: beginOffsetLeft, End: endOffsetLeft, Left: []byte(`<span class="gd input-block">`), Right: []byte(`</span>`), WantInner: 0})
						anns = append(anns, &annotate.Annotation{Start: beginOffsetRight, End: endOffsetRight, Left: []byte(`<span class="gi input-block">`), Right: []byte(`</span>`), WantInner: 0})

						if '@' != lineFirstChar {
							//leftContent := string(src[beginOffsetLeft:endOffsetLeft])
							//rightContent := string(src[beginOffsetRight:endOffsetRight])
							// This is needed to filter out the "-" and "+" at the beginning of each line from being highlighted.
							// TODO: Still not completely filtered out.
							leftContent := ""
							for line := lastDel; line < lastIns; line++ {
								leftContent += "\x00" + string(lines[line][1:]) + "\n"
							}
							rightContent := ""
							for line := lastIns; line < lineIndex; line++ {
								rightContent += "\x00" + string(lines[line][1:]) + "\n"
							}

							var sectionSegments [2][]*annotate.Annotation
							u7.HighlightedDiffFunc(leftContent, rightContent, &sectionSegments, [2]int{beginOffsetLeft, beginOffsetRight})

							anns = append(anns, sectionSegments[0]...)
							anns = append(anns, sectionSegments[1]...)
						}
					}
					lastDel, lastIns = -1, -1
				}
			}

			sort.Sort(anns)

			out, err := annotate.Annotate(src, anns, func(w io.Writer, b []byte) { template.HTMLEscape(w, b) })
			if err != nil {
				return nil, false
			}

			return out, true
		}
	default:
		return nil, false
	}
}

// Unexported blackfriday helpers.

func doubleSpace(out *bytes.Buffer) {
	if out.Len() > 0 {
		out.WriteByte('\n')
	}
}

func escapeSingleChar(char byte) (string, bool) {
	if char == '"' {
		return """, true
	}
	if char == '&' {
		return "&", true
	}
	if char == '<' {
		return "<", true
	}
	if char == '>' {
		return ">", true
	}
	return "", false
}

func attrEscape(out *bytes.Buffer, src []byte) {
	org := 0
	for i, ch := range src {
		if entity, ok := escapeSingleChar(ch); ok {
			if i > org {
				// copy all the normal characters since the last escape
				out.Write(src[org:i])
			}
			org = i + 1
			out.WriteString(entity)
		}
	}
	if org < len(src) {
		out.Write(src[org:])
	}
}
<commit_msg>gofmt.<commit_after>/*
Package github_flavored_markdown provides a GitHub Flavored Markdown renderer
with fenced code block highlighting, clickable header anchor links.

The functionality should be equivalent to the GitHub Markdown API endpoint specified at
https://developer.github.com/v3/markdown/#render-a-markdown-document-in-raw-mode, except
the rendering is performed locally.

See example below for how to generate a complete HTML page, including CSS styles.
*/
package github_flavored_markdown

import (
	"bytes"
	"fmt"
	"html"
	"io"
	"regexp"
	"sort"
	"strings"
	"text/template"

	"github.com/microcosm-cc/bluemonday"
	"github.com/russross/blackfriday"
	"github.com/shurcooL/go/highlight_go"
	"github.com/shurcooL/go/u/u7"
	"github.com/shurcooL/sanitized_anchor_name"
	"github.com/sourcegraph/annotate"
	"github.com/sourcegraph/syntaxhighlight"
)

// Markdown renders GitHub Flavored Markdown text.
func Markdown(text []byte) []byte {
	htmlFlags := 0
	renderer := &renderer{Html: blackfriday.HtmlRenderer(htmlFlags, "", "").(*blackfriday.Html)}

	// Parser extensions for GitHub Flavored Markdown.
	extensions := 0
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS
	//extensions |= blackfriday.EXTENSION_HARD_LINE_BREAK

	unsanitized := blackfriday.Markdown(text, renderer, extensions)

	// GitHub Flavored Markdown-like sanitization policy.
	p := bluemonday.UGCPolicy()
	p.AllowAttrs("class").Matching(bluemonday.SpaceSeparatedTokens).OnElements("div", "span")
	p.AllowAttrs("class", "name").Matching(bluemonday.SpaceSeparatedTokens).OnElements("a")
	p.AllowAttrs("rel").Matching(regexp.MustCompile(`^nofollow$`)).OnElements("a")
	p.AllowAttrs("aria-hidden").Matching(regexp.MustCompile(`^true$`)).OnElements("a")
	p.AllowDataURIImages()

	return p.SanitizeBytes(unsanitized)
}

type renderer struct {
	*blackfriday.Html
}

// GitHub Flavored Markdown header with clickable and hidden anchor.
func (_ *renderer) Header(out *bytes.Buffer, text func() bool, level int, _ string) {
	marker := out.Len()
	doubleSpace(out)

	if !text() {
		out.Truncate(marker)
		return
	}

	textString := out.String()[marker:]
	out.Truncate(marker)

	anchorName := sanitized_anchor_name.Create(html.UnescapeString(textString))

	out.WriteString(fmt.Sprintf(`<h%d><a name="%s" class="anchor" href="#%s" rel="nofollow" aria-hidden="true"><span class="octicon octicon-link"></span></a>`, level, anchorName, anchorName))
	out.WriteString(textString)
	out.WriteString(fmt.Sprintf("</h%d>\n", level))
}

// TODO: Clean up and improve this code.
// GitHub Flavored Markdown fenced code block with highlighting.
func (_ *renderer) BlockCode(out *bytes.Buffer, text []byte, lang string) {
	doubleSpace(out)

	// parse out the language name
	count := 0
	for _, elt := range strings.Fields(lang) {
		if elt[0] == '.' {
			elt = elt[1:]
		}
		if len(elt) == 0 {
			continue
		}
		out.WriteString(`<div class="highlight highlight-`)
		attrEscape(out, []byte(elt))
		lang = elt
		out.WriteString(`"><pre>`)
		count++
		break
	}

	if count == 0 {
		out.WriteString("
")
	}

	if highlightedCode, ok := highlightCode(text, lang); ok {
		out.Write(highlightedCode)
	} else {
		attrEscape(out, text)
	}

	if count == 0 {
		out.WriteString("
\n")
	} else {
		out.WriteString("

\n")
	}
}

var gfmHtmlConfig = syntaxhighlight.HTMLConfig{
	String:        "s",
	Keyword:       "k",
	Comment:       "c",
	Type:          "n",
	Literal:       "o",
	Punctuation:   "p",
	Plaintext:     "n",
	Tag:           "tag",
	HTMLTag:       "htm",
	HTMLAttrName:  "atn",
	HTMLAttrValue: "atv",
	Decimal:       "m",
}

// TODO: Support highlighting for more languages.
func highlightCode(src []byte, lang string) (highlightedCode []byte, ok bool) {
	switch lang {
	case "Go":
		var buf bytes.Buffer
		err := highlight_go.Print(src, &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "Go-old":
		var buf bytes.Buffer
		err := syntaxhighlight.Print(syntaxhighlight.NewScanner(src), &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "diff":
		switch 2 {
		default:
			var buf bytes.Buffer
			err := u7.Print(u7.NewScanner(src), &buf)
			if err != nil {
				return nil, false
			}
			return buf.Bytes(), true
		case 1:
			lines := bytes.Split(src, []byte("\n"))
			return bytes.Join(lines, []byte("\n")), true
		case 2:
			anns, err := u7.Annotate(src)
			if err != nil {
				return nil, false
			}

			lines := bytes.Split(src, []byte("\n"))
			lineStarts := make([]int, len(lines))
			var offset int
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				lineStarts[lineIndex] = offset
				offset += len(lines[lineIndex]) + 1
			}

			lastDel, lastIns := -1, -1
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				var lineFirstChar byte
				if len(lines[lineIndex]) > 0 {
					lineFirstChar = lines[lineIndex][0]
				}
				switch lineFirstChar {
				case '+':
					if lastIns == -1 {
						lastIns = lineIndex
					}
				case '-':
					if lastDel == -1 {
						lastDel = lineIndex
					}
				default:
					if lastDel != -1 || lastIns != -1 {
						if lastDel == -1 {
							lastDel = lastIns
						} else if lastIns == -1 {
							lastIns = lineIndex
						}

						beginOffsetLeft := lineStarts[lastDel]
						endOffsetLeft := lineStarts[lastIns]
						beginOffsetRight := lineStarts[lastIns]
						endOffsetRight := lineStarts[lineIndex]

						anns = append(anns, &annotate.Annotation{Start: beginOffsetLeft, End: endOffsetLeft, Left: []byte(`<span class="gd input-block">`), Right: []byte(`</span>`), WantInner: 0})
						anns = append(anns, &annotate.Annotation{Start: beginOffsetRight, End: endOffsetRight, Left: []byte(`<span class="gi input-block">`), Right: []byte(`</span>`), WantInner: 0})

						if '@' != lineFirstChar {
							//leftContent := string(src[beginOffsetLeft:endOffsetLeft])
							//rightContent := string(src[beginOffsetRight:endOffsetRight])
							// This is needed to filter out the "-" and "+" at the beginning of each line from being highlighted.
							// TODO: Still not completely filtered out.
							leftContent := ""
							for line := lastDel; line < lastIns; line++ {
								leftContent += "\x00" + string(lines[line][1:]) + "\n"
							}
							rightContent := ""
							for line := lastIns; line < lineIndex; line++ {
								rightContent += "\x00" + string(lines[line][1:]) + "\n"
							}

							var sectionSegments [2][]*annotate.Annotation
							u7.HighlightedDiffFunc(leftContent, rightContent, &sectionSegments, [2]int{beginOffsetLeft, beginOffsetRight})

							anns = append(anns, sectionSegments[0]...)
							anns = append(anns, sectionSegments[1]...)
						}
					}
					lastDel, lastIns = -1, -1
				}
			}

			sort.Sort(anns)

			out, err := annotate.Annotate(src, anns, func(w io.Writer, b []byte) { template.HTMLEscape(w, b) })
			if err != nil {
				return nil, false
			}

			return out, true
		}
	default:
		return nil, false
	}
}

// Unexported blackfriday helpers.

func doubleSpace(out *bytes.Buffer) {
	if out.Len() > 0 {
		out.WriteByte('\n')
	}
}

func escapeSingleChar(char byte) (string, bool) {
	if char == '"' {
		return """, true
	}
	if char == '&' {
		return "&", true
	}
	if char == '<' {
		return "<", true
	}
	if char == '>' {
		return ">", true
	}
	return "", false
}

func attrEscape(out *bytes.Buffer, src []byte) {
	org := 0
	for i, ch := range src {
		if entity, ok := escapeSingleChar(ch); ok {
			if i > org {
				// copy all the normal characters since the last escape
				out.Write(src[org:i])
			}
			org = i + 1
			out.WriteString(entity)
		}
	}
	if org < len(src) {
		out.Write(src[org:])
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/jurgen-kluft/go-home/config"
	"github.com/jurgen-kluft/go-home/micro-service"
	"github.com/jurgen-kluft/go-icloud-calendar"
)

// Calendar ...
type Calendar struct {
	name         string
	config       *config.CalendarConfig
	sensors      map[string]config.Csensor
	sensorStates map[string]*config.SensorState
	cals         []*icalendar.Calendar
	update       time.Time
	service      *microservice.Service
}

func new() *Calendar {
	c := &Calendar{}
	c.name = "calendar"
	c.sensors = map[string]config.Csensor{}
	c.sensorStates = map[string]*config.SensorState{}
	return c
}

// initialize  ... create a new Calendar from the given JSON configuration
func (c *Calendar) initialize(jsondata []byte) (err error) {
	c.config, err = config.CalendarConfigFromJSON(jsondata)
	if err != nil {
		c.service.Logger.LogError(c.name, err.Error())
	}
	//c.ccal.print()
	for _, sn := range c.config.Sensors {
		ekey := strings.ToLower(sn.Name)
		c.sensors[ekey] = sn
		sensor := &config.SensorState{Name: ekey, Time: time.Now()}
		if sn.Type == "string" {
			sensor.AddStringAttr(ekey, sn.State)
		} else if sn.Type == "float" {
			value, _ := strconv.ParseFloat(sn.State, 64)
			sensor.AddFloatAttr(ekey, value)
		} else if sn.Type == "bool" {
			value, _ := strconv.ParseBool(sn.State)
			sensor.AddBoolAttr(ekey, value)
		}
		c.sensorStates[ekey] = sensor
	}

	for _, cal := range c.config.Calendars {
		if strings.HasPrefix(cal.URL.String, "http") {
			c.cals = append(c.cals, icalendar.NewURLCalendar(cal.URL.String))
		} else if strings.HasPrefix(cal.URL.String, "file") {
			c.cals = append(c.cals, icalendar.NewFileCalendar(cal.URL.String))
		} else {
			c.service.Logger.LogError(c.name, fmt.Sprintf("Unknown calendar source '%s'", cal.URL))
		}
		c.service.Logger.LogInfo(c.name, cal.URL.String)
	}

	c.update = time.Now()

	return err
}

func (c *Calendar) updateSensorStates(when time.Time) error {
	//fmt.Printf("Update calendar events: '%d'\n", len(c.cals))

	for _, cal := range c.cals {
		//fmt.Printf("Update calendar events: '%s'\n", cal.Name)

		eventsForDay := cal.GetEventsByDate(when)
		for _, e := range eventsForDay {
			var dname string
			var dstate string
			title := strings.Replace(e.Summary, ":", ".", 3)
			title = strings.Replace(title, "=", " = ", 1)
			n, err := fmt.Sscanf(title, "%s = %s", &dname, &dstate)
			if n == 2 && err == nil {
				//fmt.Printf("Parsed: '%s' - '%s' - '%s' - '%s'\n", domain, dproduct, dname, dstate)
				dname = strings.ToLower(strings.Trim(dname, " "))
				dstate = strings.ToLower(strings.Trim(dstate, " "))
				ekey := dname

				sensor, exists := c.sensorStates[ekey]
				if exists {
					sensor.Time = time.Now()
					if sensor.StringAttrs != nil {
						sensor.StringAttrs[0].Value = dstate
					} else if sensor.IntAttrs != nil {
						value, _ := strconv.ParseInt(dstate, 10, 64)
						sensor.IntAttrs[0].Value = value
					} else if sensor.FloatAttrs != nil {
						value, _ := strconv.ParseFloat(dstate, 64)
						sensor.FloatAttrs[0].Value = value
					} else if sensor.BoolAttrs != nil {
						value, _ := strconv.ParseBool(dstate)
						sensor.BoolAttrs[0].Value = value
					}
				}
			}
		}
	}
	return nil
}

func (c *Calendar) load() (err error) {
	for _, cal := range c.cals {
		err = cal.Load()
	}
	return err
}

func timeInRange(when time.Time, rangeBegin time.Time, rangeEnd time.Time) bool {
	t := when.Unix()
	rb := rangeBegin.Unix()
	re := rangeEnd.Unix()
	return t >= rb && t < re
}

func weekOrWeekEndStartEnd(now time.Time) (weekend bool, westart, weend, wdstart, wdend time.Time) {
	day := now.Day()

	westart = now
	weend = now

	wdstart = now
	wdend = now

	if now.Weekday() == time.Friday || now.Weekday() == time.Saturday || now.Weekday() == time.Sunday {
		if now.Weekday() == time.Saturday {
			day--
		} else if now.Weekday() == time.Sunday {
			day -= 2
		}
		westart = time.Date(now.Year(), now.Month(), day, 18, 0, 0, 0, now.Location())
		weend = time.Date(now.Year(), now.Month(), day+2, 18, 0, 0, 0, now.Location())
		if timeInRange(now, westart, weend) {
			weekend = true
			wdstart = weend
			wdend = time.Date(wdstart.Year(), wdstart.Month(), wdstart.Day()+5, 18, 0, 0, 0, now.Location())
		} else {
			weekend = false
			if now.Weekday() == time.Friday {
				wdend = westart
				wdstart = time.Date(wdend.Year(), wdend.Month(), wdend.Day()-4, 18, 0, 0, 0, now.Location())
			} else {
				wdstart = weend
				wdend = time.Date(wdstart.Year(), wdstart.Month(), wdstart.Day()+5, 18, 0, 0, 0, now.Location())
			}
		}
	} else {
		if now.Weekday() == time.Monday {
			day--
		} else if now.Weekday() == time.Tuesday {
			day -= 2
		} else if now.Weekday() == time.Wednesday {
			day -= 3
		} else if now.Weekday() == time.Thursday {
			day -= 4
		} else if now.Weekday() == time.Friday {
			day -= 5
		}

		wdstart = time.Date(now.Year(), now.Month(), day, 18, 0, 0, 0, now.Location())
		wdend = time.Date(now.Year(), now.Month(), day+4, 18, 0, 0, 0, now.Location())

		weekend = false
		westart = wdend
		weend = time.Date(westart.Year(), westart.Month(), westart.Day()+2, 18, 0, 0, 0, now.Location())
	}

	return weekend, westart, weend, wdstart, wdend
}

func (c *Calendar) applyRulesToSensorStates() {
	for _, p := range c.config.Rules {
		var sensor *config.SensorState
		var ifsensor *config.SensorState
		var exists bool
		sensor, exists = c.sensorStates[p.Key]
		if exists && sensor.StringAttrs != nil {
			for _, ifthen := range p.IfThen {
				ifsensor, exists = c.sensorStates[ifthen.Key]
				if exists && ifsensor.StringAttrs != nil {
					ifthenValue := ifsensor.StringAttrs[0].Value
					if ifthenValue == ifthen.State {
						sensor.StringAttrs[0].Value = p.State
					}
				} else if exists && ifsensor.BoolAttrs != nil {
					ifthenValue := ifsensor.BoolAttrs[0].Value
					if ifthenValue && ("true" == ifthen.State) {
						sensor.StringAttrs[0].Value = p.State
					} else if !ifthenValue && ("false" == ifthen.State) {
						sensor.StringAttrs[0].Value = p.State
					}
				} else {
					c.service.Logger.LogError(c.name, fmt.Sprintf("Logical error when applying rules to sensor states (%s)", p.Key+", "+p.State))
				}
			}
		}
	}
}

func (c *Calendar) publishSensorState(name string, sensorjsonbytes []byte) {
	c.service.Pubsub.Publish(fmt.Sprintf("state/sensor/%s/", name), sensorjsonbytes)
}

// Process will update 'events' from the calendar
func (c *Calendar) Process() {
	var err error
	now := time.Now()

	if now.Unix() >= c.update.Unix() {
		// Download again after 15 minutes
		c.update = time.Unix(now.Unix()+1*60, 0)

		// Download calendars
		// fmt.Println("CALENDAR: LOAD")
		err := c.load()
		if err != nil {
			c.service.Logger.LogError(c.name, err.Error())
			return
		}
	}

	// Other general states
	weekendsensor, exists := c.sensorStates["weekend"]
	if exists {
		if weekendsensor.BoolAttrs != nil {
			weekend, _, _, _, _ := weekOrWeekEndStartEnd(now)
			weekendsensor.BoolAttrs[0].Value = weekend
			sensorjsonbytes, err := weekendsensor.ToJSON()
			if err == nil {
				c.publishSensorState("weekend", sensorjsonbytes)
			} else {
				c.service.Logger.LogError(c.name, err.Error())
			}
		}
	}

	// Update sensors and apply configured rules to sensors
	err = c.updateSensorStates(now)
	if err == nil {
		c.applyRulesToSensorStates()

		// Publish sensor states
		for _, ss := range c.sensorStates {
			sensorjsonbytes, err := ss.ToJSON()
			if err == nil {
				c.publishSensorState(ss.Name, sensorjsonbytes)
			} else {
				c.service.Logger.LogError(c.name, err.Error())
			}
		}
	} else {
		c.service.Logger.LogError(c.name, err.Error())
	}
}

func main() {
	register := []string{"config/calendar/"}
	subscribe := []string{"config/calendar/", "config/request/"}

	m := microservice.New("calendar")
	m.RegisterAndSubscribe(register, subscribe)

	c := new()
	c.service = m

	m.RegisterHandler("config/calendar/", func(m *microservice.Service, topic string, msg []byte) bool {
		m.Logger.LogInfo(m.Name, "received configuration")
		err := c.initialize(msg)
		if err != nil {
			c = nil
			m.Logger.LogError(c.name, err.Error())
		} else {
			// Register emitter channel for every sensor
			for _, ss := range c.sensorStates {
				if err = m.Register(fmt.Sprintf("state/sensor/%s/", ss.Name)); err != nil {
					m.Logger.LogError("pubsub", err.Error())
				}
			}
		}
		return true
	})

	tickCount := 0
	m.RegisterHandler("tick/", func(m *microservice.Service, topic string, msg []byte) bool {
		if tickCount%5 == 0 {
			if c != nil && c.config == nil {
				m.Pubsub.PublishStr("config/request/", m.Name)
			}
		}
		if tickCount%30 == 0 {
			if c != nil && c.config != nil {
				c.Process()
			}
		}
		tickCount++
		return true
	})

	m.Loop()
}
<commit_msg>* enable some printf<commit_after>package main

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/jurgen-kluft/go-home/config"
	"github.com/jurgen-kluft/go-home/micro-service"
	"github.com/jurgen-kluft/go-icloud-calendar"
)

// Calendar ...
type Calendar struct {
	name         string
	config       *config.CalendarConfig
	sensors      map[string]config.Csensor
	sensorStates map[string]*config.SensorState
	cals         []*icalendar.Calendar
	update       time.Time
	service      *microservice.Service
}

func new() *Calendar {
	c := &Calendar{}
	c.name = "calendar"
	c.sensors = map[string]config.Csensor{}
	c.sensorStates = map[string]*config.SensorState{}
	return c
}

// initialize  ... create a new Calendar from the given JSON configuration
func (c *Calendar) initialize(jsondata []byte) (err error) {
	c.config, err = config.CalendarConfigFromJSON(jsondata)
	if err != nil {
		c.service.Logger.LogError(c.name, err.Error())
	}
	//c.ccal.print()
	for _, sn := range c.config.Sensors {
		ekey := strings.ToLower(sn.Name)
		c.sensors[ekey] = sn
		sensor := &config.SensorState{Name: ekey, Time: time.Now()}
		if sn.Type == "string" {
			sensor.AddStringAttr(ekey, sn.State)
		} else if sn.Type == "float" {
			value, _ := strconv.ParseFloat(sn.State, 64)
			sensor.AddFloatAttr(ekey, value)
		} else if sn.Type == "bool" {
			value, _ := strconv.ParseBool(sn.State)
			sensor.AddBoolAttr(ekey, value)
		}
		c.sensorStates[ekey] = sensor
	}

	for _, cal := range c.config.Calendars {
		if strings.HasPrefix(cal.URL.String, "http") {
			c.cals = append(c.cals, icalendar.NewURLCalendar(cal.URL.String))
		} else if strings.HasPrefix(cal.URL.String, "file") {
			c.cals = append(c.cals, icalendar.NewFileCalendar(cal.URL.String))
		} else {
			c.service.Logger.LogError(c.name, fmt.Sprintf("Unknown calendar source '%s'", cal.URL))
		}
		c.service.Logger.LogInfo(c.name, cal.URL.String)
	}

	c.update = time.Now()

	return err
}

func (c *Calendar) updateSensorStates(when time.Time) error {
	fmt.Printf("Update calendar events: '%d'\n", len(c.cals))

	for _, cal := range c.cals {
		fmt.Printf("Update calendar events: '%s'\n", cal.Name)

		eventsForDay := cal.GetEventsByDate(when)
		for _, e := range eventsForDay {
			var dname string
			var dstate string
			title := strings.Replace(e.Summary, ":", ".", 3)
			title = strings.Replace(title, "=", " = ", 1)
			n, err := fmt.Sscanf(title, "%s = %s", &dname, &dstate)
			if n == 2 && err == nil {
				fmt.Printf("Parsed: '%s' - '%s' - '%s' - '%s'\n", domain, dproduct, dname, dstate)
				dname = strings.ToLower(strings.Trim(dname, " "))
				dstate = strings.ToLower(strings.Trim(dstate, " "))
				ekey := dname

				sensor, exists := c.sensorStates[ekey]
				if exists {
					sensor.Time = time.Now()
					if sensor.StringAttrs != nil {
						sensor.StringAttrs[0].Value = dstate
					} else if sensor.IntAttrs != nil {
						value, _ := strconv.ParseInt(dstate, 10, 64)
						sensor.IntAttrs[0].Value = value
					} else if sensor.FloatAttrs != nil {
						value, _ := strconv.ParseFloat(dstate, 64)
						sensor.FloatAttrs[0].Value = value
					} else if sensor.BoolAttrs != nil {
						value, _ := strconv.ParseBool(dstate)
						sensor.BoolAttrs[0].Value = value
					}
				}
			}
		}
	}
	return nil
}

func (c *Calendar) load() (err error) {
	for _, cal := range c.cals {
		err = cal.Load()
	}
	return err
}

func timeInRange(when time.Time, rangeBegin time.Time, rangeEnd time.Time) bool {
	t := when.Unix()
	rb := rangeBegin.Unix()
	re := rangeEnd.Unix()
	return t >= rb && t < re
}

func weekOrWeekEndStartEnd(now time.Time) (weekend bool, westart, weend, wdstart, wdend time.Time) {
	day := now.Day()

	westart = now
	weend = now

	wdstart = now
	wdend = now

	if now.Weekday() == time.Friday || now.Weekday() == time.Saturday || now.Weekday() == time.Sunday {
		if now.Weekday() == time.Saturday {
			day--
		} else if now.Weekday() == time.Sunday {
			day -= 2
		}
		westart = time.Date(now.Year(), now.Month(), day, 18, 0, 0, 0, now.Location())
		weend = time.Date(now.Year(), now.Month(), day+2, 18, 0, 0, 0, now.Location())
		if timeInRange(now, westart, weend) {
			weekend = true
			wdstart = weend
			wdend = time.Date(wdstart.Year(), wdstart.Month(), wdstart.Day()+5, 18, 0, 0, 0, now.Location())
		} else {
			weekend = false
			if now.Weekday() == time.Friday {
				wdend = westart
				wdstart = time.Date(wdend.Year(), wdend.Month(), wdend.Day()-4, 18, 0, 0, 0, now.Location())
			} else {
				wdstart = weend
				wdend = time.Date(wdstart.Year(), wdstart.Month(), wdstart.Day()+5, 18, 0, 0, 0, now.Location())
			}
		}
	} else {
		if now.Weekday() == time.Monday {
			day--
		} else if now.Weekday() == time.Tuesday {
			day -= 2
		} else if now.Weekday() == time.Wednesday {
			day -= 3
		} else if now.Weekday() == time.Thursday {
			day -= 4
		} else if now.Weekday() == time.Friday {
			day -= 5
		}

		wdstart = time.Date(now.Year(), now.Month(), day, 18, 0, 0, 0, now.Location())
		wdend = time.Date(now.Year(), now.Month(), day+4, 18, 0, 0, 0, now.Location())

		weekend = false
		westart = wdend
		weend = time.Date(westart.Year(), westart.Month(), westart.Day()+2, 18, 0, 0, 0, now.Location())
	}

	return weekend, westart, weend, wdstart, wdend
}

func (c *Calendar) applyRulesToSensorStates() {
	for _, p := range c.config.Rules {
		var sensor *config.SensorState
		var ifsensor *config.SensorState
		var exists bool
		sensor, exists = c.sensorStates[p.Key]
		if exists && sensor.StringAttrs != nil {
			for _, ifthen := range p.IfThen {
				ifsensor, exists = c.sensorStates[ifthen.Key]
				if exists && ifsensor.StringAttrs != nil {
					ifthenValue := ifsensor.StringAttrs[0].Value
					if ifthenValue == ifthen.State {
						sensor.StringAttrs[0].Value = p.State
					}
				} else if exists && ifsensor.BoolAttrs != nil {
					ifthenValue := ifsensor.BoolAttrs[0].Value
					if ifthenValue && ("true" == ifthen.State) {
						sensor.StringAttrs[0].Value = p.State
					} else if !ifthenValue && ("false" == ifthen.State) {
						sensor.StringAttrs[0].Value = p.State
					}
				} else {
					c.service.Logger.LogError(c.name, fmt.Sprintf("Logical error when applying rules to sensor states (%s)", p.Key+", "+p.State))
				}
			}
		}
	}
}

func (c *Calendar) publishSensorState(name string, sensorjsonbytes []byte) {
	c.service.Pubsub.Publish(fmt.Sprintf("state/sensor/%s/", name), sensorjsonbytes)
}

// Process will update 'events' from the calendar
func (c *Calendar) Process() {
	var err error
	now := time.Now()

	if now.Unix() >= c.update.Unix() {
		// Download again after 15 minutes
		c.update = time.Unix(now.Unix()+1*60, 0)

		// Download calendars
		fmt.Println("CALENDAR: LOAD")
		err := c.load()
		if err != nil {
			c.service.Logger.LogError(c.name, err.Error())
			return
		}
	}

	// Other general states
	weekendsensor, exists := c.sensorStates["weekend"]
	if exists {
		if weekendsensor.BoolAttrs != nil {
			weekend, _, _, _, _ := weekOrWeekEndStartEnd(now)
			weekendsensor.BoolAttrs[0].Value = weekend
			sensorjsonbytes, err := weekendsensor.ToJSON()
			if err == nil {
				c.publishSensorState("weekend", sensorjsonbytes)
			} else {
				c.service.Logger.LogError(c.name, err.Error())
			}
		}
	}

	// Update sensors and apply configured rules to sensors
	err = c.updateSensorStates(now)
	if err == nil {
		c.applyRulesToSensorStates()

		// Publish sensor states
		for _, ss := range c.sensorStates {
			sensorjsonbytes, err := ss.ToJSON()
			if err == nil {
				c.publishSensorState(ss.Name, sensorjsonbytes)
			} else {
				c.service.Logger.LogError(c.name, err.Error())
			}
		}
	} else {
		c.service.Logger.LogError(c.name, err.Error())
	}
}

func main() {
	register := []string{"config/calendar/"}
	subscribe := []string{"config/calendar/", "config/request/"}

	m := microservice.New("calendar")
	m.RegisterAndSubscribe(register, subscribe)

	c := new()
	c.service = m

	m.RegisterHandler("config/calendar/", func(m *microservice.Service, topic string, msg []byte) bool {
		m.Logger.LogInfo(m.Name, "received configuration")
		err := c.initialize(msg)
		if err != nil {
			c = nil
			m.Logger.LogError(c.name, err.Error())
		} else {
			// Register emitter channel for every sensor
			for _, ss := range c.sensorStates {
				if err = m.Register(fmt.Sprintf("state/sensor/%s/", ss.Name)); err != nil {
					m.Logger.LogError("pubsub", err.Error())
				}
			}
		}
		return true
	})

	tickCount := 0
	m.RegisterHandler("tick/", func(m *microservice.Service, topic string, msg []byte) bool {
		if tickCount%5 == 0 {
			if c != nil && c.config == nil {
				m.Pubsub.PublishStr("config/request/", m.Name)
			}
		}
		if tickCount%30 == 0 {
			if c != nil && c.config != nil {
				c.Process()
			}
		}
		tickCount++
		return true
	})

	m.Loop()
}
<|endoftext|>


<commit_before>// Package calendar generates and has utlilty functions to generate an HTML calendar for use in nlgids.
package calendar

import (
	"fmt"
	"sort"
	"time"
)

// TODO: next button should have a ref to point to something.

var avail = [...]string{"past", "busy", "free"}

type Available int

const (
	past Available = iota
	busy
	free
)

func (a Available) String() string { return avail[a] }

// Calendar holds the HTML that makes up the calendar. Each
// day is indexed by the 12 o' clock night time time.Time.
// All date are in the UTC timezone.
type Calendar struct {
	days map[time.Time]Available
	// TODO: add date we are going see, for the month heading at the top of the calendar.
}

type times []time.Time

func (t times) Len() int           { return len(t) }
func (t times) Less(i, j int) bool { return t[i].Before(t[j]) }
func (t times) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }

func (c *Calendar) heading() string {
	// lang!
	return "
Sun	Mon	Tue	Wed	Thu	Fri	Sat
\n"
}

func (c *Calendar) Header() string {
	// Todo insert links 'n such.
	return `

bla

`
}

func (c *Calendar) Footer() string {
	return `

`
}

func (c *Calendar) openTR() string  { return "
\n" }
func (c *Calendar) closeTR() string { return "
\n" }

func (c *Calendar) entry(t time.Time) string {
	// Make something like:
	// 3
	d := c.days[t]
	day := fmt.Sprintf("%02d", t.Day())
	class := fmt.Sprintf("\t", d)
	close := "	\n"
	href := ""
	switch d {
	case free:
		date := fmt.Sprintf("%4d-%02d-%02d", t.Year(), t.Month(), t.Day())
		href = fmt.Sprintf("%s", date, day)
	case busy:
		href = day
	case past:
		href = day
	}
	s := class + href + close
	return s
}

func (c *Calendar) HTML() string {
	s := c.Header()
	s += "
\n"
	s += c.html()
	s += "
"
	s += c.Footer()
	return s
}

func (c *Calendar) html() string {
	keys := times{}
	for k := range c.days {
		keys = append(keys, k)
	}

	s := c.heading()
	sort.Sort(keys)
	i := 0
	for _, k := range keys {
		if i%7 == 0 {
			if i > 0 {
				s += c.closeTR()
			}
			s += c.closeTR()
		}
		s += c.entry(k)
		i++
	}
	s += c.closeTR()
	return s
}

// New creates a new month calendar based on d. D must
// be in the form: YYYY-MM-DD.
func New(d string) (*Calendar, error) {
	date, err := time.Parse("2006-01-02", d)
	if err != nil {
		return nil, err
	}
	cal := &Calendar{days: make(map[time.Time]Available)}

	now := time.Now()
	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)
	first := time.Date(date.Year(), date.Month(), 1, 0, 0, 0, 0, time.UTC)
	last := time.Date(date.Year(), date.Month()+1, 1, 0, 0, 0, 0, time.UTC)
	last = last.Add(-24 * time.Hour)

	// Add the remaining days of the previous month.
	for i := 0; i < int(first.Weekday()); i++ {
		lastMonthDay := first.AddDate(0, 0, -1*(i+1))
		cal.days[lastMonthDay] = free

		if lastMonthDay.Before(today) {
			cal.days[lastMonthDay] = past
		}
	}

	// Loop from i to lastDay and add the entire month.
	for i := 1; i <= last.Day(); i++ {
		day := time.Date(date.Year(), date.Month(), i, 0, 0, 0, 0, time.UTC)

		cal.days[day] = free

		if day.Before(today) {
			cal.days[day] = past
		}
	}

	// These are dates in the new month.
	j := 1
	for i := int(last.Weekday()) + 1; i < 7; i++ {
		nextMonthDay := last.AddDate(0, 0, j)
		cal.days[nextMonthDay] = free

		if nextMonthDay.Before(today) {
			cal.days[nextMonthDay] = past
		}

		j++
	}

	return cal, nil
}
<commit_msg>make it a btn-block<commit_after>// Package calendar generates and has utlilty functions to generate an HTML calendar for use in nlgids.
package calendar

import (
	"fmt"
	"sort"
	"time"
)

// TODO: next button should have a ref to point to something.

var avail = [...]string{"past", "busy", "free"}

type Available int

const (
	past Available = iota
	busy
	free
)

func (a Available) String() string { return avail[a] }

// Calendar holds the HTML that makes up the calendar. Each
// day is indexed by the 12 o' clock night time time.Time.
// All date are in the UTC timezone.
type Calendar struct {
	days map[time.Time]Available
	// TODO: add date we are going see, for the month heading at the top of the calendar.
}

type times []time.Time

func (t times) Len() int           { return len(t) }
func (t times) Less(i, j int) bool { return t[i].Before(t[j]) }
func (t times) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }

func (c *Calendar) heading() string {
	// lang!
	return "
Sun	Mon	Tue	Wed	Thu	Fri	Sat
\n"
}

func (c *Calendar) Header() string {
	// Todo insert links 'n such.
	return `

bla

`
}

func (c *Calendar) Footer() string {
	return `

`
}

func (c *Calendar) openTR() string  { return "
\n" }
func (c *Calendar) closeTR() string { return "
\n" }

func (c *Calendar) entry(t time.Time) string {
	// Make something like:
	// 3
	d := c.days[t]
	day := fmt.Sprintf("%02d", t.Day())
	class := fmt.Sprintf("\t", d)
	close := "	\n"
	href := ""
	switch d {
	case free:
		date := fmt.Sprintf("%4d-%02d-%02d", t.Year(), t.Month(), t.Day())
		href = fmt.Sprintf("%s", date, day)
		class= fmt.Sprintf("\t", d)
	case busy:
		href = day
	case past:
		href = day
	}
	s := class + href + close
	return s
}

func (c *Calendar) HTML() string {
	s := c.Header()
	s += "
\n"
	s += c.html()
	s += "
"
	s += c.Footer()
	return s
}

func (c *Calendar) html() string {
	keys := times{}
	for k := range c.days {
		keys = append(keys, k)
	}

	s := c.heading()
	sort.Sort(keys)
	i := 0
	for _, k := range keys {
		if i%7 == 0 {
			if i > 0 {
				s += c.closeTR()
			}
			s += c.closeTR()
		}
		s += c.entry(k)
		i++
	}
	s += c.closeTR()
	return s
}

// New creates a new month calendar based on d. D must
// be in the form: YYYY-MM-DD.
func New(d string) (*Calendar, error) {
	date, err := time.Parse("2006-01-02", d)
	if err != nil {
		return nil, err
	}
	cal := &Calendar{days: make(map[time.Time]Available)}

	now := time.Now()
	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)
	first := time.Date(date.Year(), date.Month(), 1, 0, 0, 0, 0, time.UTC)
	last := time.Date(date.Year(), date.Month()+1, 1, 0, 0, 0, 0, time.UTC)
	last = last.Add(-24 * time.Hour)

	// Add the remaining days of the previous month.
	for i := 0; i < int(first.Weekday()); i++ {
		lastMonthDay := first.AddDate(0, 0, -1*(i+1))
		cal.days[lastMonthDay] = free

		if lastMonthDay.Before(today) {
			cal.days[lastMonthDay] = past
		}
	}

	// Loop from i to lastDay and add the entire month.
	for i := 1; i <= last.Day(); i++ {
		day := time.Date(date.Year(), date.Month(), i, 0, 0, 0, 0, time.UTC)

		cal.days[day] = free

		if day.Before(today) {
			cal.days[day] = past
		}
	}

	// These are dates in the new month.
	j := 1
	for i := int(last.Weekday()) + 1; i < 7; i++ {
		nextMonthDay := last.AddDate(0, 0, j)
		cal.days[nextMonthDay] = free

		if nextMonthDay.Before(today) {
			cal.days[nextMonthDay] = past
		}

		j++
	}

	return cal, nil
}
<|endoftext|>


<commit_before>package assert

func ShouldPanic(t Exam, fn func()) {
	t.Log("Expecting panic error.")
	defer func() {
		t.Log("Expecting panic error.")
		err := recover()
		t.Log(err)
		if err == nil {
			t.Log("Expected panic occured.")
			t.Fail()
		}
	}()
	//	fn()
	t.Log("Expected panic.")
	//	t.Fail()
}

func ShouldNotPanic(t Exam, fn func()) {
	defer func() {
		err := recover()
		if err != nil {
			t.Log("Expect func() not to panic.", err)
			t.Fail()
		}
	}()
	fn()
}
<commit_msg>Fixing panic.<commit_after>package assert

func ShouldPanic(t Exam, fn func()) {
	defer func() {
		err := recover()
		t.Log(err)
		if err == nil {
			t.Log("Expected panic occured.")
			t.Fail()
		}
	}()
	fn()
	t.Log("Expecting panic.")
	t.Fail()
}

func ShouldNotPanic(t Exam, fn func()) {
	defer func() {
		err := recover()
		if err != nil {
			t.Log("Expect func() not to panic.", err)
			t.Fail()
		}
	}()
	fn()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/rainforestapp/rainforest-cli/rainforest"
	"github.com/urfave/cli"
)

const (
	// Version of the app in SemVer
	version = "2.0.0"
)

var (
	// Build info to be set while building using:
	// go build -ldflags "-X main.build 'build details'"
	build string

	// Rainforest API client
	api *rainforest.Client

	// default output for printing resource tables
	tablesOut io.Writer = os.Stdout

	// Run status polling interval
	runStatusPollInterval = time.Second * 5

	// Batch size (number of rows) for tabular var upload
	tabularBatchSize = 50
	// Concurrent connections when uploading CSV rows
	tabularConcurrency = 1
)

// notImplemented is a placholder function for actions that are not yet implemented.
// In the current iteration of cli.
func notImplemented(c *cli.Context) error {
	fmt.Println("Not yet implemented. Check it out soon!")
	return nil
}

// cliContext is an interface providing context of running application
// i.e. command line options and flags. One of the types that provides the interface is
// cli.Context, the other is fakeCLIContext which is used for testing.
type cliContext interface {
	String(flag string) (val string)
	StringSlice(flag string) (vals []string)
	Bool(flag string) (val bool)
	Int(flag string) (val int)

	Args() (args cli.Args)
}

// main is an entry point of the app. It sets up the new cli app, and defines the API.
func main() {
	app := cli.NewApp()
	app.Usage = "Rainforest QA CLI - https://www.rainforestqa.com/"
	app.Version = version

	// Before running any of the commands we init the API Client
	app.Before = func(c *cli.Context) error {
		api = rainforest.NewClient(c.String("token"))

		// Set the User-Agent that will be used for api calls
		if build != "" {
			api.UserAgent = "rainforest-cli/" + version + " build: " + build
		} else {
			api.UserAgent = "rainforest-cli/" + version
		}

		return nil
	}

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "token, t",
			Usage:  "API token. You can find it at https://app.rainforestqa.com/settings/integrations",
			EnvVar: "RAINFOREST_API_TOKEN",
		},
		cli.IntFlag{
			Name:   "threads",
			Usage:  "Used to customize the amount of concurrent HTTP connections to use.",
			Value:  16,
			EnvVar: "RAINFOREST_THREADS",
		},
	}

	app.Commands = []cli.Command{
		{
			Name:    "run",
			Aliases: []string{"r"},
			Usage:   "Run your tests on Rainforest",
			Action:  startRun,
			Description: "Runs your tests on Rainforest platform. " +
				"You need to specify list of test IDs to run or use keyword 'all'. " +
				"Alternatively you can use one of the filtering options.",
			ArgsUsage: "[test IDs]",
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "tag",
					Usage: "filter tests by `TAG`. Can be used multiple times for filtering by multiple tags.",
				},
				cli.StringFlag{
					Name:  "site, site-id",
					Usage: "filter tests by a specific site. You can see a list of your `SITE-ID`s with sites command.",
				},
				cli.StringFlag{
					Name:  "folder",
					Usage: "filter tests by a specific folder. You can see a list of your `FOLDER-ID`s with folders command.",
				},
				cli.StringSliceFlag{
					Name: "browser, browsers",
					Usage: "specify the `BROWSER` you wish to run against. This overrides test level settings." +
						"Can be used multiple times to run against multiple browsers.",
				},
				cli.StringFlag{
					Name:  "environment-id",
					Usage: "run your tests using specified `ENVIRONMENT`. Otherwise it will use your default one.",
				},
				cli.StringFlag{
					Name:  "crowd",
					Value: "default",
					Usage: "run your tests using specified `CROWD`. Available choices are: default or on_premise_crowd. " +
						"Contact your CSM for more details.",
				},
				cli.StringFlag{
					Name: "conflict",
					Usage: "use the abort option to abort any runs in the same environment or " +
						"use the abort-all option to abort all runs in progress.",
				},
				cli.BoolFlag{
					Name: "bg, background",
					Usage: "run in the background. This option makes cli return after succesfully starting a run, " +
						"without waiting for the run results.",
				},
				cli.BoolFlag{
					Name: "fail-fast, ff",
					Usage: "fail the build as soon as the first failed result comes in. " +
						"If you don't pass this it will wait until 100% of the run is done. Use with --fg.",
				},
				cli.StringFlag{
					Name: "custom-url",
					Usage: "use a custom `URL` for this run. Example use case: an ad-hoc QA environment with Fourchette. " +
						"You will need to specify a site_id too for this to work.",
				},
				cli.BoolFlag{
					Name: "git-trigger",
					Usage: "only trigger a run when the last commit (for a git repo in the current working directory) " +
						"contains @rainforest and a list of one or more tags. rainforest-cli exits with 0 otherwise.",
				},
				cli.StringFlag{
					Name:  "description",
					Usage: "add arbitrary `DESCRIPTION` to the run.",
				},
				cli.StringFlag{
					Name:  "junit-file",
					Usage: "Create a JUnit XML report `FILE` with the specified name. Must be run in foreground mode.",
				},
				cli.StringFlag{
					Name:  "import-variable-name",
					Usage: "`NAME` of the tabular variable to be created or updated.",
				},
				cli.StringFlag{
					Name:  "import-variable-csv-file",
					Usage: "`PATH` to the CSV file to be uploaded.",
				},
				cli.BoolFlag{
					Name:  "overwrite-variable",
					Usage: "If the flag is set, named variable will be updated.",
				},
				cli.BoolFlag{
					Name:  "single-use",
					Usage: "This option marks uploaded variable as single-use",
				},
				cli.StringFlag{
					Name:  "wait, reattach",
					Usage: "monitor existing run with `RUN_ID` instead of starting a new one.",
				},
			},
		},
		{
			Name:      "new",
			Usage:     "Create a new RFML test",
			ArgsUsage: "[name]",
			Description: "Create new Rainforest test in RFML format (Rainforest Markup Language). " +
				"You may also specify a custom test title or file name.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` at which to create new test.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
			},
			Action: notImplemented,
		},
		{
			Name:      "validate",
			Usage:     "Validate your RFML tests",
			ArgsUsage: " ",
			Description: "Validate your tests for syntax and correct RFML ids for embedded tests. " +
				"If API token is set it'll validate your tests against server data as well.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` where to look for a tests to validate.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
			},
			Action: notImplemented,
		},
		{
			Name:      "upload",
			Usage:     "Upload your RFML tests",
			ArgsUsage: "[path to RFML file]",
			Description: "Uploads specified test to Rainforest. " +
				"If no filepath is given it uploads all RFML tests.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` where to look for a tests to upload.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
			},
			Action: notImplemented,
		},
		{
			Name:        "rm",
			Usage:       "Remove an RFML test locally and remotely",
			ArgsUsage:   "[path to RFML file]",
			Description: "Remove RFML file and remove test from Rainforest test suite.",
			Action:      notImplemented,
		},
		{
			Name: "download",
			// Left for legacy reason, should nuke?
			Aliases:   []string{"export"},
			Usage:     "Download your remote Rainforest tests to RFML",
			ArgsUsage: "[test IDs]",
			Description: "Download your remote tests from Rainforest to RFML. " +
				"You need to specify list of test IDs to download or use keyword 'all'. " +
				"Alternatively you can use one of the filtering options.",
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "tag",
					Usage: "filter tests by `TAG`. Can be used multiple times for filtering by multiple tags.",
				},
				cli.StringFlag{
					Name:  "site, site-id",
					Usage: "filter tests by a specific site. You can see a list of your `SITE-ID`s with sites command.",
				},
				cli.StringFlag{
					Name:  "folder",
					Usage: "filter tests by a specific folder. You can see a list of your `FOLDER-ID`s with folders command.",
				},
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` at which to save all the downloaded tests.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
				cli.BoolFlag{
					Name:  "embed-tests",
					Usage: "download your tests without extracting the steps of an embedded test.",
				},
			},
			Action: notImplemented,
		},
		{
			Name:        "csv-upload",
			Usage:       "Create or update tabular var from CSV.",
			Description: "Upload a CSV file to create or update tabular variables.",
			ArgsUsage:   "[path to CSV file]",
			Flags: []cli.Flag{
				cli.StringFlag{
					// Alternative name left for legacy reason.
					Name:  "name, import-variable-name",
					Usage: "`NAME` of the tabular variable to be created or updated.",
				},
				cli.BoolFlag{
					Name:  "overwrite-variable, overwrite",
					Usage: "If the flag is set, named variable will be updated.",
				},
				cli.BoolFlag{
					Name:  "single-use",
					Usage: "This option marks uploaded variable as single-use",
				},
				// Left here for legacy reason, but imho we should move that to args
				cli.StringFlag{
					Name:  "csv-file, import-variable-csv-file",
					Usage: "DEPRECATED: `PATH` to the CSV file to be uploaded. Since v2 please provide the path as an argument.",
				},
			},
			Action: func(c *cli.Context) error {
				return csvUpload(c, api)
			},
		},
		{
			Name:  "report",
			Usage: "Create a JUnit report from your run results",
			Description: "Creates a JUnit report from your specified run." +
				"You can specify output file using options, otherwise report will be generated to STDOUT",
			ArgsUsage: "[run ID]",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "junit-file, out, o",
					Usage: "`PATH` of file to which save the JUnit report.",
				},
				// Left here for legacy reason, but imho we should move that to args
				cli.StringFlag{
					Name:  "run-id",
					Usage: "DEPRECATED: ID of a run for which to generate results. Since v2 please provide the run ID as an argument.",
				},
			},
			Action: notImplemented,
		},
		{
			Name:  "sites",
			Usage: "Lists available sites",
			Action: func(c *cli.Context) error {
				return printSites(c, api)
			},
		},
		{
			Name:  "folders",
			Usage: "Lists available folders",
			Action: func(c *cli.Context) error {
				return printFolders(c, api)
			},
		},
		{
			Name:  "browsers",
			Usage: "Lists available browsers",
			Action: func(c *cli.Context) error {
				return printBrowsers(c, api)
			},
		},
	}

	app.Run(os.Args)
}
<commit_msg>add timestamps to errors<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"time"

	"github.com/rainforestapp/rainforest-cli/rainforest"
	"github.com/urfave/cli"
)

const (
	// Version of the app in SemVer
	version = "2.0.0"
)

var (
	// Build info to be set while building using:
	// go build -ldflags "-X main.build 'build details'"
	build string

	// Rainforest API client
	api *rainforest.Client

	// default output for printing resource tables
	tablesOut io.Writer = os.Stdout

	// Run status polling interval
	runStatusPollInterval = time.Second * 5

	// Batch size (number of rows) for tabular var upload
	tabularBatchSize = 50
	// Concurrent connections when uploading CSV rows
	tabularConcurrency = 1
)

// notImplemented is a placholder function for actions that are not yet implemented.
// In the current iteration of cli.
func notImplemented(c *cli.Context) error {
	fmt.Println("Not yet implemented. Check it out soon!")
	return nil
}

// cliContext is an interface providing context of running application
// i.e. command line options and flags. One of the types that provides the interface is
// cli.Context, the other is fakeCLIContext which is used for testing.
type cliContext interface {
	String(flag string) (val string)
	StringSlice(flag string) (vals []string)
	Bool(flag string) (val bool)
	Int(flag string) (val int)

	Args() (args cli.Args)
}

// Create custom writer which will use timestamps
type logWriter struct{}

func (l *logWriter) Write(p []byte) (int, error) {
	log.Printf("%s", p)
	return len(p), nil
}

// main is an entry point of the app. It sets up the new cli app, and defines the API.
func main() {
	app := cli.NewApp()
	app.Usage = "Rainforest QA CLI - https://www.rainforestqa.com/"
	app.Version = version
	// Use our custom writer to print our errors with timestamps
	cli.ErrWriter = &logWriter{}

	// Before running any of the commands we init the API Client
	app.Before = func(c *cli.Context) error {
		api = rainforest.NewClient(c.String("token"))

		// Set the User-Agent that will be used for api calls
		if build != "" {
			api.UserAgent = "rainforest-cli/" + version + " build: " + build
		} else {
			api.UserAgent = "rainforest-cli/" + version
		}

		return nil
	}

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "token, t",
			Usage:  "API token. You can find it at https://app.rainforestqa.com/settings/integrations",
			EnvVar: "RAINFOREST_API_TOKEN",
		},
		cli.IntFlag{
			Name:   "threads",
			Usage:  "Used to customize the amount of concurrent HTTP connections to use.",
			Value:  16,
			EnvVar: "RAINFOREST_THREADS",
		},
	}

	app.Commands = []cli.Command{
		{
			Name:    "run",
			Aliases: []string{"r"},
			Usage:   "Run your tests on Rainforest",
			Action:  startRun,
			Description: "Runs your tests on Rainforest platform. " +
				"You need to specify list of test IDs to run or use keyword 'all'. " +
				"Alternatively you can use one of the filtering options.",
			ArgsUsage: "[test IDs]",
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "tag",
					Usage: "filter tests by `TAG`. Can be used multiple times for filtering by multiple tags.",
				},
				cli.StringFlag{
					Name:  "site, site-id",
					Usage: "filter tests by a specific site. You can see a list of your `SITE-ID`s with sites command.",
				},
				cli.StringFlag{
					Name:  "folder",
					Usage: "filter tests by a specific folder. You can see a list of your `FOLDER-ID`s with folders command.",
				},
				cli.StringSliceFlag{
					Name: "browser, browsers",
					Usage: "specify the `BROWSER` you wish to run against. This overrides test level settings." +
						"Can be used multiple times to run against multiple browsers.",
				},
				cli.StringFlag{
					Name:  "environment-id",
					Usage: "run your tests using specified `ENVIRONMENT`. Otherwise it will use your default one.",
				},
				cli.StringFlag{
					Name:  "crowd",
					Value: "default",
					Usage: "run your tests using specified `CROWD`. Available choices are: default or on_premise_crowd. " +
						"Contact your CSM for more details.",
				},
				cli.StringFlag{
					Name: "conflict",
					Usage: "use the abort option to abort any runs in the same environment or " +
						"use the abort-all option to abort all runs in progress.",
				},
				cli.BoolFlag{
					Name: "bg, background",
					Usage: "run in the background. This option makes cli return after succesfully starting a run, " +
						"without waiting for the run results.",
				},
				cli.BoolFlag{
					Name: "fail-fast, ff",
					Usage: "fail the build as soon as the first failed result comes in. " +
						"If you don't pass this it will wait until 100% of the run is done. Use with --fg.",
				},
				cli.StringFlag{
					Name: "custom-url",
					Usage: "use a custom `URL` for this run. Example use case: an ad-hoc QA environment with Fourchette. " +
						"You will need to specify a site_id too for this to work.",
				},
				cli.BoolFlag{
					Name: "git-trigger",
					Usage: "only trigger a run when the last commit (for a git repo in the current working directory) " +
						"contains @rainforest and a list of one or more tags. rainforest-cli exits with 0 otherwise.",
				},
				cli.StringFlag{
					Name:  "description",
					Usage: "add arbitrary `DESCRIPTION` to the run.",
				},
				cli.StringFlag{
					Name:  "junit-file",
					Usage: "Create a JUnit XML report `FILE` with the specified name. Must be run in foreground mode.",
				},
				cli.StringFlag{
					Name:  "import-variable-name",
					Usage: "`NAME` of the tabular variable to be created or updated.",
				},
				cli.StringFlag{
					Name:  "import-variable-csv-file",
					Usage: "`PATH` to the CSV file to be uploaded.",
				},
				cli.BoolFlag{
					Name:  "overwrite-variable",
					Usage: "If the flag is set, named variable will be updated.",
				},
				cli.BoolFlag{
					Name:  "single-use",
					Usage: "This option marks uploaded variable as single-use",
				},
				cli.StringFlag{
					Name:  "wait, reattach",
					Usage: "monitor existing run with `RUN_ID` instead of starting a new one.",
				},
			},
		},
		{
			Name:      "new",
			Usage:     "Create a new RFML test",
			ArgsUsage: "[name]",
			Description: "Create new Rainforest test in RFML format (Rainforest Markup Language). " +
				"You may also specify a custom test title or file name.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` at which to create new test.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
			},
			Action: notImplemented,
		},
		{
			Name:      "validate",
			Usage:     "Validate your RFML tests",
			ArgsUsage: " ",
			Description: "Validate your tests for syntax and correct RFML ids for embedded tests. " +
				"If API token is set it'll validate your tests against server data as well.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` where to look for a tests to validate.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
			},
			Action: notImplemented,
		},
		{
			Name:      "upload",
			Usage:     "Upload your RFML tests",
			ArgsUsage: "[path to RFML file]",
			Description: "Uploads specified test to Rainforest. " +
				"If no filepath is given it uploads all RFML tests.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` where to look for a tests to upload.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
			},
			Action: notImplemented,
		},
		{
			Name:        "rm",
			Usage:       "Remove an RFML test locally and remotely",
			ArgsUsage:   "[path to RFML file]",
			Description: "Remove RFML file and remove test from Rainforest test suite.",
			Action:      notImplemented,
		},
		{
			Name: "download",
			// Left for legacy reason, should nuke?
			Aliases:   []string{"export"},
			Usage:     "Download your remote Rainforest tests to RFML",
			ArgsUsage: "[test IDs]",
			Description: "Download your remote tests from Rainforest to RFML. " +
				"You need to specify list of test IDs to download or use keyword 'all'. " +
				"Alternatively you can use one of the filtering options.",
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "tag",
					Usage: "filter tests by `TAG`. Can be used multiple times for filtering by multiple tags.",
				},
				cli.StringFlag{
					Name:  "site, site-id",
					Usage: "filter tests by a specific site. You can see a list of your `SITE-ID`s with sites command.",
				},
				cli.StringFlag{
					Name:  "folder",
					Usage: "filter tests by a specific folder. You can see a list of your `FOLDER-ID`s with folders command.",
				},
				cli.StringFlag{
					Name:   "test-folder",
					Value:  "./spec/rainforest/",
					Usage:  "`PATH` at which to save all the downloaded tests.",
					EnvVar: "RAINFOREST_TEST_FOLDER",
				},
				cli.BoolFlag{
					Name:  "embed-tests",
					Usage: "download your tests without extracting the steps of an embedded test.",
				},
			},
			Action: notImplemented,
		},
		{
			Name:        "csv-upload",
			Usage:       "Create or update tabular var from CSV.",
			Description: "Upload a CSV file to create or update tabular variables.",
			ArgsUsage:   "[path to CSV file]",
			Flags: []cli.Flag{
				cli.StringFlag{
					// Alternative name left for legacy reason.
					Name:  "name, import-variable-name",
					Usage: "`NAME` of the tabular variable to be created or updated.",
				},
				cli.BoolFlag{
					Name:  "overwrite-variable, overwrite",
					Usage: "If the flag is set, named variable will be updated.",
				},
				cli.BoolFlag{
					Name:  "single-use",
					Usage: "This option marks uploaded variable as single-use",
				},
				// Left here for legacy reason, but imho we should move that to args
				cli.StringFlag{
					Name:  "csv-file, import-variable-csv-file",
					Usage: "DEPRECATED: `PATH` to the CSV file to be uploaded. Since v2 please provide the path as an argument.",
				},
			},
			Action: func(c *cli.Context) error {
				return csvUpload(c, api)
			},
		},
		{
			Name:  "report",
			Usage: "Create a JUnit report from your run results",
			Description: "Creates a JUnit report from your specified run." +
				"You can specify output file using options, otherwise report will be generated to STDOUT",
			ArgsUsage: "[run ID]",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "junit-file, out, o",
					Usage: "`PATH` of file to which save the JUnit report.",
				},
				// Left here for legacy reason, but imho we should move that to args
				cli.StringFlag{
					Name:  "run-id",
					Usage: "DEPRECATED: ID of a run for which to generate results. Since v2 please provide the run ID as an argument.",
				},
			},
			Action: notImplemented,
		},
		{
			Name:  "sites",
			Usage: "Lists available sites",
			Action: func(c *cli.Context) error {
				return printSites(c, api)
			},
		},
		{
			Name:  "folders",
			Usage: "Lists available folders",
			Action: func(c *cli.Context) error {
				return printFolders(c, api)
			},
		},
		{
			Name:  "browsers",
			Usage: "Lists available browsers",
			Action: func(c *cli.Context) error {
				return printBrowsers(c, api)
			},
		},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>// Package tunnel implements an io.ReadWriter that encrypts it's data.
// Technically it performs a Elliptic Curve Diffie Hellman key exchange
// before the first read or write (or triggered manually using Exchange())
//
// All communication over the tunnel is encrypted with AES using CFB mode.
package security

import (
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"
	"io"

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

////////////////////////////////////////////////////////////////////
// The following code has been taken from go-ecdh:              //
// https://github.com/tang0th/go-ecdh/blob/master/curve25519.go //
// It's here to prevent another external dependency.            //
////////////////////////////////////////////////////////////////////

func generateKey(rand io.Reader) (crypto.PrivateKey, crypto.PublicKey, error) {
	var pub, priv [32]byte
	var err error

	_, err = io.ReadFull(rand, priv[:])
	if err != nil {
		return nil, nil, err
	}

	curve25519.ScalarBaseMult(&pub, &priv)
	return &priv, &pub, nil
}

func marshal(p crypto.PublicKey) []byte {
	pub := p.(*[32]byte)
	return pub[:]
}

func unmarshal(data []byte) (crypto.PublicKey, bool) {
	var pub [32]byte
	if len(data) != 32 {
		return nil, false
	}

	copy(pub[:], data)
	return &pub, true
}

func generateSharedSecret(privKey crypto.PrivateKey, pubKey crypto.PublicKey) ([]byte, error) {
	var priv, pub, secret *[32]byte

	priv = privKey.(*[32]byte)
	pub = pubKey.(*[32]byte)
	secret = new([32]byte)

	curve25519.ScalarMult(secret, priv, pub)
	return secret[:], nil
}

type ecdhTunnel struct {
	// Underlying ReadWriter
	io.ReadWriter

	// Elliptic Curve Diffie Hellman keys:
	privKey crypto.PrivateKey
	pubKey  crypto.PublicKey

	crypted io.ReadWriter
}

// NewEllipticTunnel creates an io.ReadWriter that transparently encrypts all data.
func NewEllipticTunnel(rw io.ReadWriter) (*ecdhTunnel, error) {
	tnl := &ecdhTunnel{
		ReadWriter: rw,
	}

	var err error
	tnl.privKey, tnl.pubKey, err = generateKey(rand.Reader)
	if err != nil {
		return nil, err
	}

	return tnl, nil
}

// Exchange triggers the Diffie Hellman key exchange manually.
func (tnl *ecdhTunnel) Exchange() error {
	if tnl.crypted != nil {
		return nil
	}

	pubKeyBuf := marshal(tnl.pubKey)
	if _, err := tnl.ReadWriter.Write(pubKeyBuf); err != nil {
		return err
	}

	partnerBuf := make([]byte, len(pubKeyBuf))
	if _, err := tnl.ReadWriter.Read(partnerBuf); err != nil {
		return err
	}

	partnerKey, ok := unmarshal(partnerBuf)
	if !ok {
		return fmt.Errorf("Partner key unmarshal failed")
	}

	secret, err := generateSharedSecret(tnl.privKey, partnerKey)
	if err != nil {
		return err
	}

	// Transform the secret to a usable 32 byte key:
	key := Scrypt(secret, secret[:16], 32)
	inv := Scrypt(secret, secret[16:], aes.BlockSize)

	rw, err := WrapReadWriter(inv, key, tnl.ReadWriter)
	if err != nil {
		return err
	}

	tnl.crypted = rw
	return nil
}

type rwCapsule struct {
	io.Reader
	io.Writer
}

func WrapReadWriter(iv, key []byte, rw io.ReadWriter) (io.ReadWriter, error) {
	blockCipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	streamW := &cipher.StreamWriter{
		S: cipher.NewCFBEncrypter(blockCipher, iv),
		W: rw,
	}

	streamR := &cipher.StreamReader{
		S: cipher.NewCFBDecrypter(blockCipher, iv),
		R: rw,
	}

	return rwCapsule{streamR, streamW}, nil
}

// Read decrypts underlying data using CFB and will trigger a key exchange
// if this was not done yet for this session.
func (tnl *ecdhTunnel) Read(buf []byte) (int, error) {
	if err := tnl.Exchange(); err != nil {
		return 0, err
	}

	return tnl.crypted.Read(buf)
}

// Write encrypts incoming data using CFB and will trigger a key exchange
// if this was not done yet for this session.
func (tnl *ecdhTunnel) Write(buf []byte) (int, error) {
	if err := tnl.Exchange(); err != nil {
		return 0, err
	}

	n, e := tnl.crypted.Write(buf)
	return n, e
}
<commit_msg>ecdh: Fix key generation from upstream (see https://github.com/wsddn/go-ecdh/issues/3)<commit_after>// Package tunnel implements an io.ReadWriter that encrypts it's data.
// Technically it performs a Elliptic Curve Diffie Hellman key exchange
// before the first read or write (or triggered manually using Exchange())
//
// All communication over the tunnel is encrypted with AES using CFB mode.
package security

import (
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"
	"io"

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

////////////////////////////////////////////////////////////////////
// The following code has been taken from go-ecdh:              //
// https://github.com/tang0th/go-ecdh/blob/master/curve25519.go //
// It's here to prevent another external dependency.            //
////////////////////////////////////////////////////////////////////

func generateKey(rand io.Reader) (crypto.PrivateKey, crypto.PublicKey, error) {
	var pub, priv [32]byte
	var err error

	_, err = io.ReadFull(rand, priv[:])
	if err != nil {
		return nil, nil, err
	}

	priv[0] &= 248
	priv[31] &= 127
	priv[31] |= 64

	curve25519.ScalarBaseMult(&pub, &priv)
	return &priv, &pub, nil
}

func marshal(p crypto.PublicKey) []byte {
	pub := p.(*[32]byte)
	return pub[:]
}

func unmarshal(data []byte) (crypto.PublicKey, bool) {
	var pub [32]byte
	if len(data) != 32 {
		return nil, false
	}

	copy(pub[:], data)
	return &pub, true
}

func generateSharedSecret(privKey crypto.PrivateKey, pubKey crypto.PublicKey) ([]byte, error) {
	var priv, pub, secret *[32]byte

	priv = privKey.(*[32]byte)
	pub = pubKey.(*[32]byte)
	secret = new([32]byte)

	curve25519.ScalarMult(secret, priv, pub)
	return secret[:], nil
}

type ecdhTunnel struct {
	// Underlying ReadWriter
	io.ReadWriter

	// Elliptic Curve Diffie Hellman keys:
	privKey crypto.PrivateKey
	pubKey  crypto.PublicKey

	crypted io.ReadWriter
}

// NewEllipticTunnel creates an io.ReadWriter that transparently encrypts all data.
func NewEllipticTunnel(rw io.ReadWriter) (*ecdhTunnel, error) {
	tnl := &ecdhTunnel{
		ReadWriter: rw,
	}

	var err error
	tnl.privKey, tnl.pubKey, err = generateKey(rand.Reader)
	if err != nil {
		return nil, err
	}

	return tnl, nil
}

// Exchange triggers the Diffie Hellman key exchange manually.
func (tnl *ecdhTunnel) Exchange() error {
	if tnl.crypted != nil {
		return nil
	}

	pubKeyBuf := marshal(tnl.pubKey)
	if _, err := tnl.ReadWriter.Write(pubKeyBuf); err != nil {
		return err
	}

	partnerBuf := make([]byte, len(pubKeyBuf))
	if _, err := tnl.ReadWriter.Read(partnerBuf); err != nil {
		return err
	}

	partnerKey, ok := unmarshal(partnerBuf)
	if !ok {
		return fmt.Errorf("Partner key unmarshal failed")
	}

	secret, err := generateSharedSecret(tnl.privKey, partnerKey)
	if err != nil {
		return err
	}

	// Transform the secret to a usable 32 byte key:
	key := Scrypt(secret, secret[:16], 32)
	inv := Scrypt(secret, secret[16:], aes.BlockSize)

	rw, err := WrapReadWriter(inv, key, tnl.ReadWriter)
	if err != nil {
		return err
	}

	tnl.crypted = rw
	return nil
}

type rwCapsule struct {
	io.Reader
	io.Writer
}

func WrapReadWriter(iv, key []byte, rw io.ReadWriter) (io.ReadWriter, error) {
	blockCipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	streamW := &cipher.StreamWriter{
		S: cipher.NewCFBEncrypter(blockCipher, iv),
		W: rw,
	}

	streamR := &cipher.StreamReader{
		S: cipher.NewCFBDecrypter(blockCipher, iv),
		R: rw,
	}

	return rwCapsule{streamR, streamW}, nil
}

// Read decrypts underlying data using CFB and will trigger a key exchange
// if this was not done yet for this session.
func (tnl *ecdhTunnel) Read(buf []byte) (int, error) {
	if err := tnl.Exchange(); err != nil {
		return 0, err
	}

	return tnl.crypted.Read(buf)
}

// Write encrypts incoming data using CFB and will trigger a key exchange
// if this was not done yet for this session.
func (tnl *ecdhTunnel) Write(buf []byte) (int, error) {
	if err := tnl.Exchange(); err != nil {
		return 0, err
	}

	n, e := tnl.crypted.Write(buf)
	return n, e
}
<|endoftext|>