<commit_before>// Copyright 2016 Liam Stanley <me@liamstanley.io>. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

package girc

import (
	"net"
	"strings"
	"sync"
	"time"
)

// state represents the actively-changing variables within the client
// runtime.
type state struct {
	// m is a RW mutex lock, used to guard the state from goroutines causing
	// corruption.
	mu sync.RWMutex

	// reader is the socket buffer reader from the IRC server.
	reader *ircDecoder
	// reader is the socket buffer write to the IRC server.
	writer *ircEncoder
	// conn is a net.Conn reference to the IRC server.
	conn net.Conn

	// connected is true if we're actively connected to a server.
	connected bool
	// connTime is the time at which the client has connected to a server.
	connTime *time.Time
	// quitting is used to determine if we've finished quitting/cleaning up.
	quitting bool
	// reconnecting lets the internal state know a reconnect is occurring.
	reconnecting bool
	// nick is the tracker for our nickname on the server.
	nick string
	// channels represents all channels we're active in.
	channels map[string]*Channel
}

// User represents an IRC user and the state attached to them.
type User struct {
	// Nick is the users current nickname.
	Nick string
	// Ident is the users username/ident. Ident is commonly prefixed with a
	// "~", which indicates that they do not have a identd server setup for
	// authentication.
	Ident string
	// Host is the visible host of the users connection that the server has
	// provided to us for their connection. May not always be accurate due to
	// many networks spoofing/hiding parts of the hostname for privacy
	// reasons.
	Host string
	// Name is the users "realname" or full name. Commonly contains links
	// to the IRC client being used, or something of non-importance. May also
	// be empty.
	Name string
	// FirstSeen represents the first time that the user was seen by the
	// client for the given channel.
	FirstSeen time.Time
}

// Channel represents an IRC channel and the state attached to it.
type Channel struct {
	// Name of the channel. Must be rfc compliant. Always represented as
	// lower-case, to ensure that the channel is only being tracked once.
	Name string
	// Topic of the channel.
	Topic string
	// users represents the users that we can currently see within the
	// channel.
	users map[string]*User
	// Joined represents the first time that the client joined the channel.
	Joined time.Time
}

// newState returns a clean client state.
func newState() *state {
	s := &state{}

	s.channels = make(map[string]*Channel)
	s.connected = false

	return s
}

// createChanIfNotExists creates the channel in state, if not already done.
// Always use state.mu for transaction.
func (s *state) createChanIfNotExists(name string) (channel *Channel) {
	// Not a valid channel.
	if !IsValidChannel(name) {
		return nil
	}

	name = strings.ToLower(name)
	if _, ok := s.channels[name]; !ok {
		channel = &Channel{
			Name:   name,
			users:  make(map[string]*User),
			Joined: time.Now(),
		}
		s.channels[name] = channel
	} else {
		channel = s.channels[name]
	}

	return channel
}

// deleteChannel removes the channel from state, if not already done. Always
// use state.mu for transaction.
func (s *state) deleteChannel(name string) {
	channel := s.createChanIfNotExists(name)
	if channel == nil {
		return
	}

	if _, ok := s.channels[channel.Name]; ok {
		delete(s.channels, channel.Name)
	}
}

// createUserIfNotExists creates the channel and user in state, if not already
// done. Always use state.mu for transaction.
func (s *state) createUserIfNotExists(channelName, nick string) (user *User) {
	if !IsValidNick(nick) {
		return nil
	}

	channel := s.createChanIfNotExists(channelName)
	if channel == nil {
		return nil
	}

	if _, ok := channel.users[nick]; ok {
		return channel.users[nick]
	}

	user = &User{Nick: nick, FirstSeen: time.Now()}
	channel.users[nick] = user

	return user
}

// deleteUser removes the user from channel state. Always use state.mu for
// transaction.
func (s *state) deleteUser(nick string) {
	if !IsValidNick(nick) {
		return
	}

	for k := range s.channels {
		// Check to see if they're in this channel.
		if _, ok := s.channels[k].users[nick]; !ok {
			continue
		}

		delete(s.channels[k].users, nick)
	}
}

// renameUser renames the user in state, in all locations where relevant.
// Always use state.mu for transaction.
func (s *state) renameUser(from, to string) {
	if !IsValidNick(from) || !IsValidNick(to) {
		return
	}

	for k := range s.channels {
		// Check to see if they're in this channel.
		if _, ok := s.channels[k].users[from]; !ok {
			continue
		}

		// Take the actual reference to the pointer.
		source := *s.channels[k].users[from]

		// Update the nick field (as we not only have a key, but a matching
		// struct field).
		source.Nick = to

		// Delete the old reference.
		delete(s.channels[k].users, from)

		// In with the new.
		s.channels[k].users[to] = &source
	}
}
<commit_msg>implement User LastActive, Lifetime(), Active() and IsActive()<commit_after>// Copyright 2016 Liam Stanley <me@liamstanley.io>. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

package girc

import (
	"net"
	"strings"
	"sync"
	"time"
)

// state represents the actively-changing variables within the client
// runtime.
type state struct {
	// m is a RW mutex lock, used to guard the state from goroutines causing
	// corruption.
	mu sync.RWMutex

	// reader is the socket buffer reader from the IRC server.
	reader *ircDecoder
	// reader is the socket buffer write to the IRC server.
	writer *ircEncoder
	// conn is a net.Conn reference to the IRC server.
	conn net.Conn

	// connected is true if we're actively connected to a server.
	connected bool
	// connTime is the time at which the client has connected to a server.
	connTime *time.Time
	// quitting is used to determine if we've finished quitting/cleaning up.
	quitting bool
	// reconnecting lets the internal state know a reconnect is occurring.
	reconnecting bool
	// nick is the tracker for our nickname on the server.
	nick string
	// channels represents all channels we're active in.
	channels map[string]*Channel
}

// User represents an IRC user and the state attached to them.
type User struct {
	// Nick is the users current nickname.
	Nick string
	// Ident is the users username/ident. Ident is commonly prefixed with a
	// "~", which indicates that they do not have a identd server setup for
	// authentication.
	Ident string
	// Host is the visible host of the users connection that the server has
	// provided to us for their connection. May not always be accurate due to
	// many networks spoofing/hiding parts of the hostname for privacy
	// reasons.
	Host string
	// Name is the users "realname" or full name. Commonly contains links
	// to the IRC client being used, or something of non-importance. May also
	// be empty.
	Name string
	// FirstSeen represents the first time that the user was seen by the
	// client for the given channel.
	FirstSeen time.Time

	// LastActive represents the last time that we saw the user active,
	// which could be during nickname change, message, channel join, etc.
	LastActive time.Time
}

// Lifetime represents the amount of time that has passed since we have first
// seen the user.
func (u *User) Lifetime() time.Duration {
	return time.Since(u.FirstSeen)
}

// Active represents the the amount of time that has passed since we have
// last seen the user.
func (u *User) Active() time.Duration {
	return time.Since(u.LastActive)
}

// IsActive returns true if they were active within the last 30 minutes.
func (u *User) IsActive() bool {
	return u.Active() < (time.Minute * 30)
}

// Channel represents an IRC channel and the state attached to it.
type Channel struct {
	// Name of the channel. Must be rfc compliant. Always represented as
	// lower-case, to ensure that the channel is only being tracked once.
	Name string
	// Topic of the channel.
	Topic string
	// users represents the users that we can currently see within the
	// channel.
	users map[string]*User
	// Joined represents the first time that the client joined the channel.
	Joined time.Time
}

// newState returns a clean client state.
func newState() *state {
	s := &state{}

	s.channels = make(map[string]*Channel)
	s.connected = false

	return s
}

// createChanIfNotExists creates the channel in state, if not already done.
// Always use state.mu for transaction.
func (s *state) createChanIfNotExists(name string) (channel *Channel) {
	// Not a valid channel.
	if !IsValidChannel(name) {
		return nil
	}

	name = strings.ToLower(name)
	if _, ok := s.channels[name]; !ok {
		channel = &Channel{
			Name:   name,
			users:  make(map[string]*User),
			Joined: time.Now(),
		}
		s.channels[name] = channel
	} else {
		channel = s.channels[name]
	}

	return channel
}

// deleteChannel removes the channel from state, if not already done. Always
// use state.mu for transaction.
func (s *state) deleteChannel(name string) {
	channel := s.createChanIfNotExists(name)
	if channel == nil {
		return
	}

	if _, ok := s.channels[channel.Name]; ok {
		delete(s.channels, channel.Name)
	}
}

// createUserIfNotExists creates the channel and user in state, if not already
// done. Always use state.mu for transaction.
func (s *state) createUserIfNotExists(channelName, nick string) (user *User) {
	if !IsValidNick(nick) {
		return nil
	}

	channel := s.createChanIfNotExists(channelName)
	if channel == nil {
		return nil
	}

	if _, ok := channel.users[nick]; ok {
		channel.users[nick].LastActive = time.Now()
		return channel.users[nick]
	}

	user = &User{Nick: nick, FirstSeen: time.Now(), LastActive: time.Now()}
	channel.users[nick] = user

	return user
}

// deleteUser removes the user from channel state. Always use state.mu for
// transaction.
func (s *state) deleteUser(nick string) {
	if !IsValidNick(nick) {
		return
	}

	for k := range s.channels {
		// Check to see if they're in this channel.
		if _, ok := s.channels[k].users[nick]; !ok {
			continue
		}

		delete(s.channels[k].users, nick)
	}
}

// renameUser renames the user in state, in all locations where relevant.
// Always use state.mu for transaction.
func (s *state) renameUser(from, to string) {
	if !IsValidNick(from) || !IsValidNick(to) {
		return
	}

	for k := range s.channels {
		// Check to see if they're in this channel.
		if _, ok := s.channels[k].users[from]; !ok {
			continue
		}

		// Take the actual reference to the pointer.
		source := *s.channels[k].users[from]

		// Update the nick field (as we not only have a key, but a matching
		// struct field).
		source.Nick = to
		source.LastActive = time.Now()

		// Delete the old reference.
		delete(s.channels[k].users, from)

		// In with the new.
		s.channels[k].users[to] = &source
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"time"
)

// In-memory representation of a secret.
type StoreEntry struct {
	Secret    string    `json:"secret"`
	MaxClicks int       `json:"max_clicks"`
	Clicks    int       `json:"clicks"`
	DateAdded time.Time `json:"date_added"`
}

// Secret augmented with computed fields.
type StoreEntryInfo struct {
	StoreEntry
	Id        string `json:"id"`
	PathQuery string `json:"path_query"`
	Url       string `json:"url"`
	ApiUrl    string `json:"api_url"`
}

type secretStore map[string]StoreEntry

// hashStruct returns a hash from an arbitrary structure, usable in a URL.
func hashStruct(data interface{}) (hash string) {
	hashBytes := sha256.Sum256([]byte(fmt.Sprintf("%#v", data)))
	hash = base64.URLEncoding.EncodeToString(hashBytes[:])
	return
}

// AddEntry adds a secret to the store.
func (st secretStore) AddEntry(e StoreEntry, id string) string {
	e.DateAdded = time.Now()
	if id == "" {
		id = hashStruct(e)
	}
	st[id] = e
	return id
}

// NewEntry adds a new secret to the store. Set id to ""
// to have it auto-generated by hashing the entry.
func (st secretStore) NewEntry(secret string, maxclicks int, id string) string {
	return st.AddEntry(StoreEntry{secret, maxclicks, 0, time.Time{}}, id)
}

// GetEntry retrieves a secret from the store.
func (st secretStore) GetEntry(id string) (se StoreEntry, ok bool) {
	se, ok = st[id]
	return
}

// GetEntryInfo wraps GetEntry and adds some computed fields.
func (st secretStore) GetEntryInfo(id string) (si StoreEntryInfo, ok bool) {
	entry, ok := st.GetEntry(id)
	pathQuery := uGet + "?id=" + id
	url := schemeHost + listen + pathQuery
	apiurl := schemeHost + listen + uApiGet + id
	return StoreEntryInfo{entry, id, pathQuery, url, apiurl}, ok
}

// GetEntryInfo wraps GetEntry and adds some computed fields. In addition it
// hides the "secret" value.
func (st secretStore) GetEntryInfoHidden(id string) (si StoreEntryInfo, ok bool) {
	si, ok = st.GetEntryInfo(id)
	si.Secret = "#HIDDEN#"
	return
}

// Click increases the click counter for an entry.
func (st secretStore) Click(id string) {
	entry, ok := st.GetEntry(id)
	if ok {
		if entry.Clicks < entry.MaxClicks-1 {
			entry.Clicks += 1
			st[id] = entry
		} else {
			delete(st, id)
		}
	}
	return
}
<commit_msg>avoid padding characters in ids<commit_after>package main

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"time"
)

// In-memory representation of a secret.
type StoreEntry struct {
	Secret    string    `json:"secret"`
	MaxClicks int       `json:"max_clicks"`
	Clicks    int       `json:"clicks"`
	DateAdded time.Time `json:"date_added"`
}

// Secret augmented with computed fields.
type StoreEntryInfo struct {
	StoreEntry
	Id        string `json:"id"`
	PathQuery string `json:"path_query"`
	Url       string `json:"url"`
	ApiUrl    string `json:"api_url"`
}

type secretStore map[string]StoreEntry

// hashStruct returns a hash from an arbitrary structure, usable in a URL.
func hashStruct(data interface{}) (hash string) {
	hashBytes := sha256.Sum256([]byte(fmt.Sprintf("%#v", data)))
	hash = base64.RawURLEncoding.EncodeToString(hashBytes[:])
	return
}

// AddEntry adds a secret to the store.
func (st secretStore) AddEntry(e StoreEntry, id string) string {
	e.DateAdded = time.Now()
	if id == "" {
		id = hashStruct(e)
	}
	st[id] = e
	return id
}

// NewEntry adds a new secret to the store. Set id to ""
// to have it auto-generated by hashing the entry.
func (st secretStore) NewEntry(secret string, maxclicks int, id string) string {
	return st.AddEntry(StoreEntry{secret, maxclicks, 0, time.Time{}}, id)
}

// GetEntry retrieves a secret from the store.
func (st secretStore) GetEntry(id string) (se StoreEntry, ok bool) {
	se, ok = st[id]
	return
}

// GetEntryInfo wraps GetEntry and adds some computed fields.
func (st secretStore) GetEntryInfo(id string) (si StoreEntryInfo, ok bool) {
	entry, ok := st.GetEntry(id)
	pathQuery := uGet + "?id=" + id
	url := schemeHost + listen + pathQuery
	apiurl := schemeHost + listen + uApiGet + id
	return StoreEntryInfo{entry, id, pathQuery, url, apiurl}, ok
}

// GetEntryInfo wraps GetEntry and adds some computed fields. In addition it
// hides the "secret" value.
func (st secretStore) GetEntryInfoHidden(id string) (si StoreEntryInfo, ok bool) {
	si, ok = st.GetEntryInfo(id)
	si.Secret = "#HIDDEN#"
	return
}

// Click increases the click counter for an entry.
func (st secretStore) Click(id string) {
	entry, ok := st.GetEntry(id)
	if ok {
		if entry.Clicks < entry.MaxClicks-1 {
			entry.Clicks += 1
			st[id] = entry
		} else {
			delete(st, id)
		}
	}
	return
}
<|endoftext|>


<commit_before>// Package store automatically configures a database to store structured information in an sql database
package store

import (
	"database/sql"
	"errors"
	"reflect"
	"strings"
	"sync"

	_ "github.com/mxk/go-sqlite/sqlite3"
)

const (
	add = iota
	get
	update
	remove
	getPage
	count
)

type field struct {
	isStruct bool
	pos      int
	name     string
}

type typeInfo struct {
	primary    int
	fields     []field
	statements []*sql.Stmt
}

type Store struct {
	db    *sql.DB
	types map[string]typeInfo
	mutex sync.Mutex
}

func New(dataSourceName string) (*Store, error) {
	db, err := sql.Open("sqlite3", dataSourceName)
	if err != nil {
		return nil, err
	}
	return &Store{
		db:    db,
		types: make(map[string]typeInfo),
	}, nil
}

func (s *Store) Close() error {
	err := s.db.Close()
	s.db = nil
	return err
}

func (s *Store) Register(i interface{}) error {
	if s.db == nil {
		return ErrDBClosed
	} else if !isPointerStruct(i) {
		return ErrNoPointerStruct
	}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.defineType(i)
}

func (s *Store) defineType(i interface{}) error {
	name := typeName(i)
	if _, ok := s.types[name]; ok {
		return nil
	}

	s.types[name] = typeInfo{}

	v := reflect.ValueOf(i).Elem()
	numFields := v.Type().NumField()
	fields := make([]field, 0, numFields)
	id := 0
	idType := 0

	for n := 0; n < numFields; n++ {
		f := v.Type().Field(n)
		if f.PkgPath != "" { // not exported
			continue
		}
		fieldName := f.Name
		if fn := f.Tag.Get("store"); fn != "" {
			fieldName = fn
		}
		if fieldName == "-" { // Skip field
			continue
		}
		tmp := strings.ToLower(fieldName)
		for _, tf := range fields {
			if strings.ToLower(tf.name) == tmp {
				return ErrDuplicateColumn
			}
		}
		isPointer := f.Type.Kind() == reflect.Ptr
		var iface interface{}
		if isPointer {
			iface = v.Field(n).Interface()
		} else {
			iface = v.Field(n).Addr().Interface()
		}
		isStruct := false
		if isPointerStruct(iface) {
			s.defineType(iface)
			isStruct = true
		} else if !isValidType(iface) {
			continue
		}
		if isValidKeyType(iface) {
			if idType < 3 && f.Tag.Get("key") == "1" {
				idType = 3
				id = len(fields)
			} else if idType < 2 && strings.ToLower(fieldName) == "id" {
				idType = 2
				id = len(fields)
			} else if idType < 1 {
				idType = 1
				id = len(fields)
			}
		}
		fields = append(fields, field{
			isStruct,
			n,
			fieldName,
		})
	}
	if idType == 0 {
		return ErrNoKey
	}
	s.types[name] = typeInfo{
		primary: id,
	}

	// create statements
	var (
		sqlVars, sqlParams, setSQLParams, tableVars string
		doneFirst, doneFirstNonKey                  bool
	)

	for pos, f := range fields {
		if doneFirst {
			tableVars += ", "
		} else {
			doneFirst = true
		}
		if pos != id {
			if doneFirstNonKey {
				sqlVars += ", "
				setSQLParams += ", "
				sqlParams += ", "
			} else {
				doneFirstNonKey = true
			}
		}
		var varType string
		if f.isStruct {
			varType = "INTEGER"
		} else {
			varType = getType(i, f.pos)
		}
		tableVars += "[" + f.name + "] " + varType
		if pos == id {
			tableVars += " PRIMARY KEY AUTOINCREMENT"
		} else {
			sqlVars += "[" + f.name + "]"
			setSQLParams += "[" + f.name + "] = ?"
			sqlParams += "?"
		}
	}

	statements := make([]*sql.Stmt, 6)

	sql := "CREATE TABLE IF NOT EXISTS [" + name + "](" + tableVars + ");"
	_, err := s.db.Exec(sql)
	if err != nil {
		return err
	}

	sql = "INSERT INTO [" + name + "] (" + sqlVars + ") VALUES (" + sqlParams + ");"
	stmt, err := s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[add] = stmt

	sql = "SELECT " + sqlVars + " FROM [" + name + "] WHERE [" + fields[id].name + "] = ? LIMIT 1;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[get] = stmt

	sql = "UPDATE [" + name + "] SET " + setSQLParams + " WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[update] = stmt

	sql = "DELETE FROM [" + name + "] WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[remove] = stmt

	sql = "SELECT [" + fields[id].name + "] FROM [" + name + "] ORDER BY [" + fields[id].name + "] LIMIT ? OFFSET ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[getPage] = stmt

	sql = "SELECT COUNT(1) FROM [" + name + "];"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[count] = stmt

	s.types[name] = typeInfo{
		primary:    id,
		fields:     fields,
		statements: statements,
	}
	return nil
}

func (s *Store) Set(is ...interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	var toSet []interface{}
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return ErrUnregisteredType
		}
		toSet = toSet[:0]
		err := s.set(i, &t, &toSet)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) set(i interface{}, t *typeInfo, toSet *[]interface{}) error {
	for _, oi := range *toSet {
		if oi == i {
			return nil
		}
	}
	(*toSet) = append(*toSet, i)
	id := t.GetID(i)
	isUpdate := id != 0
	vars := make([]interface{}, 0, len(t.fields))
	for pos, f := range t.fields {
		if pos == t.primary {
			continue
		}
		if f.isStruct {
			ni := getFieldPointer(i, f.pos)
			nt := s.types[typeName(ni)]
			err := s.set(ni, &nt, toSet)
			if err != nil {
				return err
			}
			vars = append(vars, getField(ni, nt.fields[nt.primary].pos))
		} else {
			vars = append(vars, getField(i, f.pos))
		}
	}
	if isUpdate {
		vars = append(vars, id)
		_, err := t.statements[update].Exec(vars...)
		return err
	}
	r, err := t.statements[add].Exec(vars...)
	if err != nil {
		return err
	}
	lid, err := r.LastInsertId()
	if err != nil {
		return err
	}
	t.SetID(i, lid)
	return nil
}

func (s *Store) Get(is ...interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.get(is...)
}
func (s *Store) get(is ...interface{}) error {
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return ErrUnregisteredType
		}
		id := t.GetID(i)
		if id == 0 {
			continue
		}
		vars := make([]interface{}, 0, len(t.fields)-1)
		var toGet []interface{}
		for pos, f := range t.fields {
			if pos == t.primary {
				continue
			}
			if f.isStruct {
				ni := getFieldPointer(i, f.pos)
				nt := s.types[typeName(ni)]
				toGet = append(toGet, ni)
				vars = append(vars, getFieldPointer(ni, nt.fields[nt.primary].pos))
			} else {
				vars = append(vars, getFieldPointer(i, f.pos))
			}
		}
		row := t.statements[get].QueryRow(id)
		err := row.Scan(vars...)
		if err == sql.ErrNoRows {
			t.SetID(i, 0)
		} else if err != nil {
			return err
		} else if len(toGet) > 0 {
			if err = s.get(toGet...); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *Store) GetPage(is []interface{}, offset int) (int, error) {
	if len(is) == 0 {
		return 0, nil
	}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	t := s.types[typeName(is[0])]
	rows, err := t.statements[getPage].Query(len(is), offset)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	return s.getPage(is, rows)
}

func (s *Store) getPage(is []interface{}, rows *sql.Rows) (int, error) {
	t := s.types[typeName(is[0])]
	n := 0
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err != nil {
			return 0, err
		}
		t.SetID(is[n], id)
		n++
	}
	is = is[:n]
	if err := rows.Err(); err == sql.ErrNoRows {
		return 0, nil
	} else if err != nil {
		return 0, err
	} else if err = s.get(is...); err != nil {
		return 0, err
	}
	return n, nil
}

func (s *Store) Remove(is ...interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return ErrUnregisteredType
		}
		_, err := t.statements[remove].Exec(t.GetID(i))
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) Count(i interface{}) (int, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if !isPointerStruct(i) {
		return 0, ErrNoPointerStruct
	}
	name := typeName(i)
	stmt := s.types[name].statements[count]
	res, err := stmt.Query()
	if err != nil {
		return 0, err
	}
	num := 0
	err = res.Scan(&num)
	return num, err
}

// Errors

var (
	ErrDBClosed         = errors.New("database already closed")
	ErrNoPointerStruct  = errors.New("given variable is not a pointer to a struct")
	ErrNoKey            = errors.New("could not determine key")
	ErrDuplicateColumn  = errors.New("duplicate column name found")
	ErrUnregisteredType = errors.New("type not registered")
	ErrInvalidType      = errors.New("invalid type")
)
<commit_msg>Added fallthrough for set when 'updating' an non-existant id<commit_after>// Package store automatically configures a database to store structured information in an sql database
package store

import (
	"database/sql"
	"errors"
	"reflect"
	"strings"
	"sync"

	_ "github.com/mxk/go-sqlite/sqlite3"
)

const (
	add = iota
	get
	update
	remove
	getPage
	count
)

type field struct {
	isStruct bool
	pos      int
	name     string
}

type typeInfo struct {
	primary    int
	fields     []field
	statements []*sql.Stmt
}

type Store struct {
	db    *sql.DB
	types map[string]typeInfo
	mutex sync.Mutex
}

func New(dataSourceName string) (*Store, error) {
	db, err := sql.Open("sqlite3", dataSourceName)
	if err != nil {
		return nil, err
	}
	return &Store{
		db:    db,
		types: make(map[string]typeInfo),
	}, nil
}

func (s *Store) Close() error {
	err := s.db.Close()
	s.db = nil
	return err
}

func (s *Store) Register(i interface{}) error {
	if s.db == nil {
		return ErrDBClosed
	} else if !isPointerStruct(i) {
		return ErrNoPointerStruct
	}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.defineType(i)
}

func (s *Store) defineType(i interface{}) error {
	name := typeName(i)
	if _, ok := s.types[name]; ok {
		return nil
	}

	s.types[name] = typeInfo{}

	v := reflect.ValueOf(i).Elem()
	numFields := v.Type().NumField()
	fields := make([]field, 0, numFields)
	id := 0
	idType := 0

	for n := 0; n < numFields; n++ {
		f := v.Type().Field(n)
		if f.PkgPath != "" { // not exported
			continue
		}
		fieldName := f.Name
		if fn := f.Tag.Get("store"); fn != "" {
			fieldName = fn
		}
		if fieldName == "-" { // Skip field
			continue
		}
		tmp := strings.ToLower(fieldName)
		for _, tf := range fields {
			if strings.ToLower(tf.name) == tmp {
				return ErrDuplicateColumn
			}
		}
		isPointer := f.Type.Kind() == reflect.Ptr
		var iface interface{}
		if isPointer {
			iface = v.Field(n).Interface()
		} else {
			iface = v.Field(n).Addr().Interface()
		}
		isStruct := false
		if isPointerStruct(iface) {
			s.defineType(iface)
			isStruct = true
		} else if !isValidType(iface) {
			continue
		}
		if isValidKeyType(iface) {
			if idType < 3 && f.Tag.Get("key") == "1" {
				idType = 3
				id = len(fields)
			} else if idType < 2 && strings.ToLower(fieldName) == "id" {
				idType = 2
				id = len(fields)
			} else if idType < 1 {
				idType = 1
				id = len(fields)
			}
		}
		fields = append(fields, field{
			isStruct,
			n,
			fieldName,
		})
	}
	if idType == 0 {
		return ErrNoKey
	}
	s.types[name] = typeInfo{
		primary: id,
	}

	// create statements
	var (
		sqlVars, sqlParams, setSQLParams, tableVars string
		doneFirst, doneFirstNonKey                  bool
	)

	for pos, f := range fields {
		if doneFirst {
			tableVars += ", "
		} else {
			doneFirst = true
		}
		if pos != id {
			if doneFirstNonKey {
				sqlVars += ", "
				setSQLParams += ", "
				sqlParams += ", "
			} else {
				doneFirstNonKey = true
			}
		}
		var varType string
		if f.isStruct {
			varType = "INTEGER"
		} else {
			varType = getType(i, f.pos)
		}
		tableVars += "[" + f.name + "] " + varType
		if pos == id {
			tableVars += " PRIMARY KEY AUTOINCREMENT"
		} else {
			sqlVars += "[" + f.name + "]"
			setSQLParams += "[" + f.name + "] = ?"
			sqlParams += "?"
		}
	}

	statements := make([]*sql.Stmt, 6)

	sql := "CREATE TABLE IF NOT EXISTS [" + name + "](" + tableVars + ");"
	_, err := s.db.Exec(sql)
	if err != nil {
		return err
	}

	sql = "INSERT INTO [" + name + "] (" + sqlVars + ") VALUES (" + sqlParams + ");"
	stmt, err := s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[add] = stmt

	sql = "SELECT " + sqlVars + " FROM [" + name + "] WHERE [" + fields[id].name + "] = ? LIMIT 1;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[get] = stmt

	sql = "UPDATE [" + name + "] SET " + setSQLParams + " WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[update] = stmt

	sql = "DELETE FROM [" + name + "] WHERE [" + fields[id].name + "] = ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[remove] = stmt

	sql = "SELECT [" + fields[id].name + "] FROM [" + name + "] ORDER BY [" + fields[id].name + "] LIMIT ? OFFSET ?;"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[getPage] = stmt

	sql = "SELECT COUNT(1) FROM [" + name + "];"
	stmt, err = s.db.Prepare(sql)
	if err != nil {
		return err
	}
	statements[count] = stmt

	s.types[name] = typeInfo{
		primary:    id,
		fields:     fields,
		statements: statements,
	}
	return nil
}

func (s *Store) Set(is ...interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	var toSet []interface{}
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return ErrUnregisteredType
		}
		toSet = toSet[:0]
		err := s.set(i, &t, &toSet)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) set(i interface{}, t *typeInfo, toSet *[]interface{}) error {
	for _, oi := range *toSet {
		if oi == i {
			return nil
		}
	}
	(*toSet) = append(*toSet, i)
	id := t.GetID(i)
	isUpdate := id != 0
	vars := make([]interface{}, 0, len(t.fields))
	for pos, f := range t.fields {
		if pos == t.primary {
			continue
		}
		if f.isStruct {
			ni := getFieldPointer(i, f.pos)
			nt := s.types[typeName(ni)]
			err := s.set(ni, &nt, toSet)
			if err != nil {
				return err
			}
			vars = append(vars, getField(ni, nt.fields[nt.primary].pos))
		} else {
			vars = append(vars, getField(i, f.pos))
		}
	}
	if isUpdate {
		r, err := t.statements[update].Exec(append(vars, id)...)
		if r.RowsAffected() > 0 || err != nil {
			return err
		}
		// id wasn't found, so insert...
	}
	r, err := t.statements[add].Exec(vars...)
	if err != nil {
		return err
	}
	lid, err := r.LastInsertId()
	if err != nil {
		return err
	}
	t.SetID(i, lid)
	return nil
}

func (s *Store) Get(is ...interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.get(is...)
}
func (s *Store) get(is ...interface{}) error {
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return ErrUnregisteredType
		}
		id := t.GetID(i)
		if id == 0 {
			continue
		}
		vars := make([]interface{}, 0, len(t.fields)-1)
		var toGet []interface{}
		for pos, f := range t.fields {
			if pos == t.primary {
				continue
			}
			if f.isStruct {
				ni := getFieldPointer(i, f.pos)
				nt := s.types[typeName(ni)]
				toGet = append(toGet, ni)
				vars = append(vars, getFieldPointer(ni, nt.fields[nt.primary].pos))
			} else {
				vars = append(vars, getFieldPointer(i, f.pos))
			}
		}
		row := t.statements[get].QueryRow(id)
		err := row.Scan(vars...)
		if err == sql.ErrNoRows {
			t.SetID(i, 0)
		} else if err != nil {
			return err
		} else if len(toGet) > 0 {
			if err = s.get(toGet...); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *Store) GetPage(is []interface{}, offset int) (int, error) {
	if len(is) == 0 {
		return 0, nil
	}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	t := s.types[typeName(is[0])]
	rows, err := t.statements[getPage].Query(len(is), offset)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	return s.getPage(is, rows)
}

func (s *Store) getPage(is []interface{}, rows *sql.Rows) (int, error) {
	t := s.types[typeName(is[0])]
	n := 0
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err != nil {
			return 0, err
		}
		t.SetID(is[n], id)
		n++
	}
	is = is[:n]
	if err := rows.Err(); err == sql.ErrNoRows {
		return 0, nil
	} else if err != nil {
		return 0, err
	} else if err = s.get(is...); err != nil {
		return 0, err
	}
	return n, nil
}

func (s *Store) Remove(is ...interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	for _, i := range is {
		t, ok := s.types[typeName(i)]
		if !ok {
			return ErrUnregisteredType
		}
		_, err := t.statements[remove].Exec(t.GetID(i))
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) Count(i interface{}) (int, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if !isPointerStruct(i) {
		return 0, ErrNoPointerStruct
	}
	name := typeName(i)
	stmt := s.types[name].statements[count]
	res, err := stmt.Query()
	if err != nil {
		return 0, err
	}
	num := 0
	err = res.Scan(&num)
	return num, err
}

// Errors

var (
	ErrDBClosed         = errors.New("database already closed")
	ErrNoPointerStruct  = errors.New("given variable is not a pointer to a struct")
	ErrNoKey            = errors.New("could not determine key")
	ErrDuplicateColumn  = errors.New("duplicate column name found")
	ErrUnregisteredType = errors.New("type not registered")
	ErrInvalidType      = errors.New("invalid type")
)
<|endoftext|>


<commit_before>// Copyright 2014 Kelsey Hightower. All rights reserved.
// Use of this source code is governed by a BSD-style
// license found in the LICENSE file.

// Package memkv implements an in-memory key/value store.
package memkv

import (
	"errors"
	"path"
	"sort"
	"strings"
	"sync"
)

var ErrNotExist = errors.New("key does not exist")
var ErrNoMatch = errors.New("no keys match")

// A Store represents an in-memory key-value store safe for
// concurrent access.
type Store struct {
	FuncMap map[string]interface{}
	sync.RWMutex
	m map[string]KVPair
}

// New creates and initializes a new Store.
func New() Store {
	s := Store{m: make(map[string]KVPair)}
	s.FuncMap = map[string]interface{}{
		"exists":    s.Exists,
		"ls":        s.List,
		"lsdir":     s.ListDir,
		"get":       s.Get,
		"gets":      s.GetAll,
		"getallkvs": s.GetAllKVs,
		"getv":      s.GetValue,
		"getvs":     s.GetAllValues,
	}
	return s
}

// Delete deletes the KVPair associated with key.
func (s Store) Del(key string) {
	s.Lock()
	delete(s.m, key)
	s.Unlock()
}

// Exists checks for the existence of key in the store.
func (s Store) Exists(key string) bool {
	kv := s.Get(key)
	if kv.Value == "" {
		return false
	}
	return true
}

// Get gets the KVPair associated with key. If there is no KVPair
// associated with key, Get returns KVPair{}.
func (s Store) Get(key string) KVPair {
	s.RLock()
	kv := s.m[key]
	s.RUnlock()
	return kv
}

// GetAll returns a KVPair for all nodes with keys matching pattern.
// The syntax of patterns is the same as in path.Match.
func (s Store) GetAll(pattern string) KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		m, err := path.Match(pattern, kv.Key)
		if err != nil {
			return nil
		}
		if m {
			ks = append(ks, kv)
		}
	}
	if len(ks) == 0 {
		return nil
	}
	sort.Sort(ks)
	return ks
}

// GetAllKVs returns all KV-Pairs
func (s Store) GetAllKVs() KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		ks = append(ks, kv)
	}
	return ks
}

// GetValue gets the value associated with key. If there are no values
// associated with key, GetValue returns "".
func (s Store) GetValue(key string, v ...string) string {
	defaultValue := ""
	if len(v) > 0 {
		defaultValue = v[0]
	}

	kv := s.Get(key)
	if kv.Key == "" {
		return defaultValue
	}
	return kv.Value
}

func (s Store) GetAllValues(pattern string) []string {
	vs := make([]string, 0)
	for _, kv := range s.GetAll(pattern) {
		vs = append(vs, kv.Value)
	}
	sort.Strings(vs)
	return vs
}

func (s Store) List(filePath string) []string {
	vs := make([]string, 0)
	m := make(map[string]bool)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		if kv.Key == filePath {
			m[path.Base(kv.Key)] = true
			continue
		}
		if strings.HasPrefix(kv.Key, filePath) {
			m[strings.Split(stripKey(kv.Key, filePath), "/")[0]] = true
		}
	}
	for k := range m {
		vs = append(vs, k)
	}
	sort.Strings(vs)
	return vs
}

func (s Store) ListDir(filePath string) []string {
	vs := make([]string, 0)
	m := make(map[string]bool)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		if strings.HasPrefix(kv.Key, filePath) {
			items := strings.Split(stripKey(kv.Key, filePath), "/")
			if len(items) < 2 {
				continue
			}
			m[items[0]] = true
		}
	}
	for k := range m {
		vs = append(vs, k)
	}
	sort.Strings(vs)
	return vs
}

// Set sets the KVPair entry associated with key to value.
func (s Store) Set(key string, value string) {
	s.Lock()
	s.m[key] = KVPair{key, value}
	s.Unlock()
}

func (s Store) Purge() {
	s.Lock()
	for k := range s.m {
		delete(s.m, k)
	}
	s.Unlock()
}

func stripKey(key, prefix string) string {
	return strings.TrimPrefix(strings.TrimPrefix(key, prefix), "/")
}
<commit_msg>return the sorted value<commit_after>// Copyright 2014 Kelsey Hightower. All rights reserved.
// Use of this source code is governed by a BSD-style
// license found in the LICENSE file.

// Package memkv implements an in-memory key/value store.
package memkv

import (
	"errors"
	"path"
	"sort"
	"strings"
	"sync"
)

var ErrNotExist = errors.New("key does not exist")
var ErrNoMatch = errors.New("no keys match")

// A Store represents an in-memory key-value store safe for
// concurrent access.
type Store struct {
	FuncMap map[string]interface{}
	sync.RWMutex
	m map[string]KVPair
}

// New creates and initializes a new Store.
func New() Store {
	s := Store{m: make(map[string]KVPair)}
	s.FuncMap = map[string]interface{}{
		"exists":    s.Exists,
		"ls":        s.List,
		"lsdir":     s.ListDir,
		"get":       s.Get,
		"gets":      s.GetAll,
		"getallkvs": s.GetAllKVs,
		"getv":      s.GetValue,
		"getvs":     s.GetAllValues,
	}
	return s
}

// Delete deletes the KVPair associated with key.
func (s Store) Del(key string) {
	s.Lock()
	delete(s.m, key)
	s.Unlock()
}

// Exists checks for the existence of key in the store.
func (s Store) Exists(key string) bool {
	kv := s.Get(key)
	if kv.Value == "" {
		return false
	}
	return true
}

// Get gets the KVPair associated with key. If there is no KVPair
// associated with key, Get returns KVPair{}.
func (s Store) Get(key string) KVPair {
	s.RLock()
	kv := s.m[key]
	s.RUnlock()
	return kv
}

// GetAll returns a KVPair for all nodes with keys matching pattern.
// The syntax of patterns is the same as in path.Match.
func (s Store) GetAll(pattern string) KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		m, err := path.Match(pattern, kv.Key)
		if err != nil {
			return nil
		}
		if m {
			ks = append(ks, kv)
		}
	}
	if len(ks) == 0 {
		return nil
	}
	sort.Sort(ks)
	return ks
}

// GetAllKVs returns all KV-Pairs
func (s Store) GetAllKVs() KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		ks = append(ks, kv)
	}
	sort.Sort(ks)
	return ks
}

// GetValue gets the value associated with key. If there are no values
// associated with key, GetValue returns "".
func (s Store) GetValue(key string, v ...string) string {
	defaultValue := ""
	if len(v) > 0 {
		defaultValue = v[0]
	}

	kv := s.Get(key)
	if kv.Key == "" {
		return defaultValue
	}
	return kv.Value
}

func (s Store) GetAllValues(pattern string) []string {
	vs := make([]string, 0)
	for _, kv := range s.GetAll(pattern) {
		vs = append(vs, kv.Value)
	}
	sort.Strings(vs)
	return vs
}

func (s Store) List(filePath string) []string {
	vs := make([]string, 0)
	m := make(map[string]bool)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		if kv.Key == filePath {
			m[path.Base(kv.Key)] = true
			continue
		}
		if strings.HasPrefix(kv.Key, filePath) {
			m[strings.Split(stripKey(kv.Key, filePath), "/")[0]] = true
		}
	}
	for k := range m {
		vs = append(vs, k)
	}
	sort.Strings(vs)
	return vs
}

func (s Store) ListDir(filePath string) []string {
	vs := make([]string, 0)
	m := make(map[string]bool)
	s.RLock()
	defer s.RUnlock()
	for _, kv := range s.m {
		if strings.HasPrefix(kv.Key, filePath) {
			items := strings.Split(stripKey(kv.Key, filePath), "/")
			if len(items) < 2 {
				continue
			}
			m[items[0]] = true
		}
	}
	for k := range m {
		vs = append(vs, k)
	}
	sort.Strings(vs)
	return vs
}

// Set sets the KVPair entry associated with key to value.
func (s Store) Set(key string, value string) {
	s.Lock()
	s.m[key] = KVPair{key, value}
	s.Unlock()
}

func (s Store) Purge() {
	s.Lock()
	for k := range s.m {
		delete(s.m, k)
	}
	s.Unlock()
}

func stripKey(key, prefix string) string {
	return strings.TrimPrefix(strings.TrimPrefix(key, prefix), "/")
}
<|endoftext|>


<commit_before>package gotel

import (
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// Endpoint holds the reference to our DB connection
type Endpoint struct {
	Db *sql.DB
}

// reservation is when an app first registers that it will be checking into our gotel
type reservation struct {
	JobID                int    `json:"job_id"`
	Owner                string `json:"owner"`
	Notify               string `json:"notify"`
	App                  string `json:"app"`
	Component            string `json:"component"`
	Frequency            int    `json:"frequency"`
	TimeUnits            string `json:"time_units"`
	LastCheckin          int64  `json:"last_checkin"`
	LastCheckinStr       string `json:"last_checkin_str"` // human readable time
	TimeSinceLastCheckin string `json:"time_sine_last_checkin"`
	FailingSLA           bool   `json:"failing_sla"`
	NumCheckins          int    `json:"number_of_checkins"`
}

// checkin holds a struct that is populated when an app checks in as still alive
type checkin struct {
	App       string `json:"app"`
	Component string `json:"component"`
	Notes     string `json:"notes"`
}

// checkOut is for removing reservations
type checkOut struct {
	App       string `json:"app"`
	Component string `json:"component"`
}

// snooze holds a struct for when users want to pause an alert for maintenance
type snooze struct {
	App       string `json:"app"`
	Component string `json:"component"`
	Duration  int    `json:"duration"`
	TimeUnits string `json:"time_units"`
}

// InitDb initialzes and then bootstraps the database
func InitDb(host, user, pass string, conf config) *sql.DB {
	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:3306)/gotel", user, pass, host))
	if err != nil {
		panic(err)
	}
	err = db.Ping()
	if err != nil {
		panic(fmt.Sprintf("Unable to ping the DB at host [%s] user [%s]", host, user))
	}
	bootstrapDb(db, conf)
	return db
}

func storeReservation(db *sql.DB, r *reservation) (bool, error) {

	// get current unix time one day into the future as the initial insert data, give it one day to bake
	tomorrow := time.Now().Add(24 * time.Hour).UTC().Unix()
	now := time.Now().UTC().Unix()

	seconds := getSecondsFromUnits(r.Frequency, r.TimeUnits)
	if seconds < 10 {
		return false, errors.New("Unable to store reservations for less than 10 seconds at this time, for no real reason.")
	}

	stmt, err := db.Prepare(`INSERT INTO reservations(app, component, owner, notify, frequency, time_units, inserted_timestamp, last_checkin_timestamp)
		VALUES(?,?,?,?,?,?,?,?)
		ON DUPLICATE KEY UPDATE notify=?, frequency=?, time_units=?
		`)

	if err != nil {
		l.warn("unable to prepare statement %s", err)
		return false, errors.New("Unable to save record")
	}
	defer stmt.Close()

	res, err := stmt.Exec(r.App, r.Component, r.Owner, r.Notify, r.Frequency, r.TimeUnits, now, tomorrow, r.Notify, r.Frequency, r.TimeUnits)
	if err != nil {
		l.warn("Unable to insert record %s", err)
		return false, errors.New("Unable to save record")
	}

	rowCnt, err := res.RowsAffected()
	if err != nil {
		return false, errors.New("Unable to save record")
	}

	l.info("Insertedaffected = %d\n", rowCnt)
	return true, nil

}

func logHouseKeeping(db *sql.DB, c checkin, now int64) (bool, error) {

	//Insert
	stmt, err := db.Prepare("insert into housekeeping(app, component, notes, last_checkin_timestamp) values(?, ?, ?, ?)")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to store checkin")
	}
	defer stmt.Close()
	_, err = stmt.Exec(c.App, c.Component, c.Notes, now)
	if err != nil {
		l.warn("Unable to insert record %s", err)
		return false, errors.New("Unable to store checkin")
	}

	return true, nil
}

func storeCheckin(db *sql.DB, c checkin, now int64) (bool, error) {

	stmt, err := db.Prepare("UPDATE reservations SET last_checkin_timestamp = ?, num_checkins = num_checkins + 1 WHERE app=? AND component=?")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to prepare checkin")
	}
	defer stmt.Close()
	_, err = stmt.Exec(now, c.App, c.Component)
	if err != nil {
		l.warn("Unable to update reservation %s", err)
		return false, errors.New("Unable to store checkin")
	}

	return true, nil
}

func storeCheckOut(db *sql.DB, c *checkOut) (bool, error) {

	stmt, err := db.Prepare("DELETE FROM reservations WHERE app=? AND component=?")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to prepare checkin")
	}
	defer stmt.Close()
	_, err = stmt.Exec(c.App, c.Component)
	if err != nil {
		l.warn("Unable to update reservation %s", err)
		return false, errors.New("Unable to store checkin")
	}
	return true, nil
}

func storeSnooze(db *sql.DB, p *snooze) (bool, error) {
	futureSeconds := getSecondsFromUnits(p.Duration, p.TimeUnits)

	pausedTime := time.Now().Add(time.Duration(futureSeconds) * time.Second).UTC().Unix()

	stmt, err := db.Prepare("UPDATE reservations SET last_checkin_timestamp = ? WHERE app=? AND component=?")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to prepare snooze")
	}
	defer stmt.Close()

	_, err = stmt.Exec(pausedTime, p.App, p.Component)
	if err != nil {
		l.warn("Unable to update snooze %s", err)
		return false, errors.New("Unable to store snooze")
	}

	return true, nil
}

func getSecondsFromUnits(freq int, units string) int {
	var seconds int
	if units == "seconds" {
		seconds = freq
	} else if units == "minutes" {
		seconds = 60 * freq
	} else if units == "hours" {
		seconds = 60 * 60 * freq
	} else {
		seconds = 60 * 60 * 24 * freq
	}
	return seconds
}

func bootstrapDb(db *sql.DB, conf config) {

	l.info("Bootstrapping GoTel DB tables")

	alertSQL := `
			CREATE TABLE IF NOT EXISTS  alerts (
			  id int(11) unsigned NOT NULL AUTO_INCREMENT,
			  app varchar(30) DEFAULT NULL,
			  component varchar(30) DEFAULT NULL,
			  alert_time int(11) DEFAULT NULL,
			  alerters text DEFAULT NULL,
			  insert_time timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
			  PRIMARY KEY (id)
			) ENGINE=InnoDB DEFAULT CHARSET=utf8;`

	_, err := db.Exec(alertSQL)
	if err != nil {
		panic(err)
		return
	}

	sql := `
		CREATE TABLE IF NOT EXISTS reservations (
		  id int(11) NOT NULL AUTO_INCREMENT,
		  app varchar(150) DEFAULT NULL,
		  component varchar(150) DEFAULT NULL,
			owner text DEFAULT NULL,
		  notify text DEFAULT NULL,
		  frequency int(11) DEFAULT NULL,
		  time_units varchar(30) DEFAULT NULL,
		  inserted_timestamp int(11) DEFAULT NULL,
		  num_checkins int(11) DEFAULT '0',
		  last_alert_timestamp int(11) DEFAULT NULL,
		  last_checkin_timestamp int(11) DEFAULT NULL,
		  PRIMARY KEY (id),
		  UNIQUE KEY uniq_app (app,component)
		) ENGINE=InnoDB AUTO_INCREMENT=2 DEFAULT CHARSET=utf8;
			`
	_, err = db.Exec(sql)
	if err != nil {
		panic(err)
		return
	}

	housekeeping := `CREATE TABLE IF NOT EXISTS housekeeping (
		  id int(11) NOT NULL AUTO_INCREMENT,
		  app varchar(30) DEFAULT NULL,
		  component varchar(30) DEFAULT NULL,
		  notes text,
		  last_checkin_timestamp int(11) DEFAULT NULL,
		  insert_time timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
		  PRIMARY KEY (id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8;`
	_, err = db.Exec(housekeeping)
	if err != nil {
		panic(err)
		return
	}

	nodes := `CREATE TABLE IF NOT EXISTS nodes (
			  id int(11) unsigned NOT NULL AUTO_INCREMENT,
			  ip_address varchar(20) DEFAULT NULL,
			  node_id int(30) DEFAULT NULL,
			  PRIMARY KEY (id),
			  UNIQUE KEY uniq_ip (ip_address)
			) ENGINE=InnoDB AUTO_INCREMENT=7 DEFAULT CHARSET=utf8;`
	_, err = db.Exec(nodes)
	if err != nil {
		panic(err)
		return
	}

	// store gotel as the initial application to monitor
	l.info("Starting to bootstrap worker/coordinator reservations...")
	tomorrow := time.Now().Add(24 * time.Hour).UTC().Unix()
	now := time.Now().UTC().Unix()
	gotelApp := `INSERT INTO reservations (app, component, owner, notify, frequency, time_units, inserted_timestamp, last_checkin_timestamp)
		VALUES ('gotel', 'coordinator', ?, ?, 5, 'minutes', ?, ?) ON DUPLICATE KEY UPDATE owner=?`
	_, err = db.Exec(gotelApp, conf.Main.GotelOwnerEmail, conf.Main.GotelOwnerEmail, now, tomorrow, conf.Main.GotelOwnerEmail)
	if err != nil {
		l.warn("storing gotel/coordinator as initial app [%q]", err)
	} else {
		l.info("Inserted gotel/coordinator as first app to monitor")
	}

	gotelAppWorker := `INSERT INTO reservations (app, component, owner, notify, frequency, time_units, inserted_timestamp, last_checkin_timestamp)
		VALUES ('gotel', 'worker', ?, ?, 5, 'minutes', ?, ?) ON DUPLICATE KEY UPDATE owner=?`
	_, err = db.Exec(gotelAppWorker, conf.Main.GotelOwnerEmail, conf.Main.GotelOwnerEmail, now, tomorrow, conf.Main.GotelOwnerEmail)
	if err != nil {
		l.warn("storing gotel/worker as initial app [%v]", err)
	} else {
		l.info("Inserted gotel/worker as first app to monitor")
	}
}
<commit_msg>go vet fixes<commit_after>package gotel

import (
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// Endpoint holds the reference to our DB connection
type Endpoint struct {
	Db *sql.DB
}

// reservation is when an app first registers that it will be checking into our gotel
type reservation struct {
	JobID                int    `json:"job_id"`
	Owner                string `json:"owner"`
	Notify               string `json:"notify"`
	App                  string `json:"app"`
	Component            string `json:"component"`
	Frequency            int    `json:"frequency"`
	TimeUnits            string `json:"time_units"`
	LastCheckin          int64  `json:"last_checkin"`
	LastCheckinStr       string `json:"last_checkin_str"` // human readable time
	TimeSinceLastCheckin string `json:"time_sine_last_checkin"`
	FailingSLA           bool   `json:"failing_sla"`
	NumCheckins          int    `json:"number_of_checkins"`
}

// checkin holds a struct that is populated when an app checks in as still alive
type checkin struct {
	App       string `json:"app"`
	Component string `json:"component"`
	Notes     string `json:"notes"`
}

// checkOut is for removing reservations
type checkOut struct {
	App       string `json:"app"`
	Component string `json:"component"`
}

// snooze holds a struct for when users want to pause an alert for maintenance
type snooze struct {
	App       string `json:"app"`
	Component string `json:"component"`
	Duration  int    `json:"duration"`
	TimeUnits string `json:"time_units"`
}

// InitDb initialzes and then bootstraps the database
func InitDb(host, user, pass string, conf config) *sql.DB {
	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:3306)/gotel", user, pass, host))
	if err != nil {
		panic(err)
	}
	err = db.Ping()
	if err != nil {
		panic(fmt.Sprintf("Unable to ping the DB at host [%s] user [%s]", host, user))
	}
	bootstrapDb(db, conf)
	return db
}

func storeReservation(db *sql.DB, r *reservation) (bool, error) {

	// get current unix time one day into the future as the initial insert data, give it one day to bake
	tomorrow := time.Now().Add(24 * time.Hour).UTC().Unix()
	now := time.Now().UTC().Unix()

	seconds := getSecondsFromUnits(r.Frequency, r.TimeUnits)
	if seconds < 10 {
		return false, errors.New("Unable to store reservations for less than 10 seconds at this time, for no real reason.")
	}

	stmt, err := db.Prepare(`INSERT INTO reservations(app, component, owner, notify, frequency, time_units, inserted_timestamp, last_checkin_timestamp)
		VALUES(?,?,?,?,?,?,?,?)
		ON DUPLICATE KEY UPDATE notify=?, frequency=?, time_units=?
		`)

	if err != nil {
		l.warn("unable to prepare statement %s", err)
		return false, errors.New("Unable to save record")
	}
	defer stmt.Close()

	res, err := stmt.Exec(r.App, r.Component, r.Owner, r.Notify, r.Frequency, r.TimeUnits, now, tomorrow, r.Notify, r.Frequency, r.TimeUnits)
	if err != nil {
		l.warn("Unable to insert record %s", err)
		return false, errors.New("Unable to save record")
	}

	rowCnt, err := res.RowsAffected()
	if err != nil {
		return false, errors.New("Unable to save record")
	}

	l.info("Insertedaffected = %d\n", rowCnt)
	return true, nil

}

func logHouseKeeping(db *sql.DB, c checkin, now int64) (bool, error) {

	//Insert
	stmt, err := db.Prepare("insert into housekeeping(app, component, notes, last_checkin_timestamp) values(?, ?, ?, ?)")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to store checkin")
	}
	defer stmt.Close()
	_, err = stmt.Exec(c.App, c.Component, c.Notes, now)
	if err != nil {
		l.warn("Unable to insert record %s", err)
		return false, errors.New("Unable to store checkin")
	}

	return true, nil
}

func storeCheckin(db *sql.DB, c checkin, now int64) (bool, error) {

	stmt, err := db.Prepare("UPDATE reservations SET last_checkin_timestamp = ?, num_checkins = num_checkins + 1 WHERE app=? AND component=?")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to prepare checkin")
	}
	defer stmt.Close()
	_, err = stmt.Exec(now, c.App, c.Component)
	if err != nil {
		l.warn("Unable to update reservation %s", err)
		return false, errors.New("Unable to store checkin")
	}

	return true, nil
}

func storeCheckOut(db *sql.DB, c *checkOut) (bool, error) {

	stmt, err := db.Prepare("DELETE FROM reservations WHERE app=? AND component=?")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to prepare checkin")
	}
	defer stmt.Close()
	_, err = stmt.Exec(c.App, c.Component)
	if err != nil {
		l.warn("Unable to update reservation %s", err)
		return false, errors.New("Unable to store checkin")
	}
	return true, nil
}

func storeSnooze(db *sql.DB, p *snooze) (bool, error) {
	futureSeconds := getSecondsFromUnits(p.Duration, p.TimeUnits)

	pausedTime := time.Now().Add(time.Duration(futureSeconds) * time.Second).UTC().Unix()

	stmt, err := db.Prepare("UPDATE reservations SET last_checkin_timestamp = ? WHERE app=? AND component=?")
	if err != nil {
		l.warn("Unable to prepare record %s", err)
		return false, errors.New("Unable to prepare snooze")
	}
	defer stmt.Close()

	_, err = stmt.Exec(pausedTime, p.App, p.Component)
	if err != nil {
		l.warn("Unable to update snooze %s", err)
		return false, errors.New("Unable to store snooze")
	}

	return true, nil
}

func getSecondsFromUnits(freq int, units string) int {
	var seconds int
	if units == "seconds" {
		seconds = freq
	} else if units == "minutes" {
		seconds = 60 * freq
	} else if units == "hours" {
		seconds = 60 * 60 * freq
	} else {
		seconds = 60 * 60 * 24 * freq
	}
	return seconds
}

func bootstrapDb(db *sql.DB, conf config) {

	l.info("Bootstrapping GoTel DB tables")

	alertSQL := `
			CREATE TABLE IF NOT EXISTS  alerts (
			  id int(11) unsigned NOT NULL AUTO_INCREMENT,
			  app varchar(30) DEFAULT NULL,
			  component varchar(30) DEFAULT NULL,
			  alert_time int(11) DEFAULT NULL,
			  alerters text DEFAULT NULL,
			  insert_time timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
			  PRIMARY KEY (id)
			) ENGINE=InnoDB DEFAULT CHARSET=utf8;`

	_, err := db.Exec(alertSQL)
	if err != nil {
		panic(err)
	}

	sql := `
		CREATE TABLE IF NOT EXISTS reservations (
		  id int(11) NOT NULL AUTO_INCREMENT,
		  app varchar(150) DEFAULT NULL,
		  component varchar(150) DEFAULT NULL,
			owner text DEFAULT NULL,
		  notify text DEFAULT NULL,
		  frequency int(11) DEFAULT NULL,
		  time_units varchar(30) DEFAULT NULL,
		  inserted_timestamp int(11) DEFAULT NULL,
		  num_checkins int(11) DEFAULT '0',
		  last_alert_timestamp int(11) DEFAULT NULL,
		  last_checkin_timestamp int(11) DEFAULT NULL,
		  PRIMARY KEY (id),
		  UNIQUE KEY uniq_app (app,component)
		) ENGINE=InnoDB AUTO_INCREMENT=2 DEFAULT CHARSET=utf8;
			`
	_, err = db.Exec(sql)
	if err != nil {
		panic(err)
	}

	housekeeping := `CREATE TABLE IF NOT EXISTS housekeeping (
		  id int(11) NOT NULL AUTO_INCREMENT,
		  app varchar(30) DEFAULT NULL,
		  component varchar(30) DEFAULT NULL,
		  notes text,
		  last_checkin_timestamp int(11) DEFAULT NULL,
		  insert_time timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
		  PRIMARY KEY (id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8;`
	_, err = db.Exec(housekeeping)
	if err != nil {
		panic(err)
	}

	nodes := `CREATE TABLE IF NOT EXISTS nodes (
			  id int(11) unsigned NOT NULL AUTO_INCREMENT,
			  ip_address varchar(20) DEFAULT NULL,
			  node_id int(30) DEFAULT NULL,
			  PRIMARY KEY (id),
			  UNIQUE KEY uniq_ip (ip_address)
			) ENGINE=InnoDB AUTO_INCREMENT=7 DEFAULT CHARSET=utf8;`
	_, err = db.Exec(nodes)
	if err != nil {
		panic(err)
	}

	// store gotel as the initial application to monitor
	l.info("Starting to bootstrap worker/coordinator reservations...")
	tomorrow := time.Now().Add(24 * time.Hour).UTC().Unix()
	now := time.Now().UTC().Unix()
	gotelApp := `INSERT INTO reservations (app, component, owner, notify, frequency, time_units, inserted_timestamp, last_checkin_timestamp)
		VALUES ('gotel', 'coordinator', ?, ?, 5, 'minutes', ?, ?) ON DUPLICATE KEY UPDATE owner=?`
	_, err = db.Exec(gotelApp, conf.Main.GotelOwnerEmail, conf.Main.GotelOwnerEmail, now, tomorrow, conf.Main.GotelOwnerEmail)
	if err != nil {
		l.warn("storing gotel/coordinator as initial app [%q]", err)
	} else {
		l.info("Inserted gotel/coordinator as first app to monitor")
	}

	gotelAppWorker := `INSERT INTO reservations (app, component, owner, notify, frequency, time_units, inserted_timestamp, last_checkin_timestamp)
		VALUES ('gotel', 'worker', ?, ?, 5, 'minutes', ?, ?) ON DUPLICATE KEY UPDATE owner=?`
	_, err = db.Exec(gotelAppWorker, conf.Main.GotelOwnerEmail, conf.Main.GotelOwnerEmail, now, tomorrow, conf.Main.GotelOwnerEmail)
	if err != nil {
		l.warn("storing gotel/worker as initial app [%v]", err)
	} else {
		l.info("Inserted gotel/worker as first app to monitor")
	}
}
<|endoftext|>


<commit_before>package tally

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/textproto"
	"os"
	"os/signal"
	"reflect"
	"regexp"
	"strings"
	"time"
)

// The Action struct represents a possible reaction to a given line.
// Regexps are attached to the Action so that they can be compiled once
// when the bot is initialized instead of each time a line needs to be parsed.

type Action struct {
	re    *regexp.Regexp
	Parse func(*regexp.Regexp, string) interface{}
	Run   func(*Bot, interface{})
}

type Bot struct {
	Server   string
	Port     string
	Nick     string
	Channel  string
	Trac_URL string
	Trac_RSS string
	Interval time.Duration
	Ignore   []string
	conn     net.Conn

	actions []*Action
}

// Returns a new bot that has been configured according to
// config.json

func NewBot() *Bot {
	bot := new(Bot)
	f, err := ioutil.ReadFile("./config.json")
	if err != nil {
		log.Fatal("Error loading config.json: %v\n\n", err)
	}
	err = json.Unmarshal(f, bot)
	if err != nil {
		log.Fatal("Error unmarshalling config.json: %v\n\n", err)
	}
	return bot
}

// Establishes a connection to the server and joins a channel

func (bot *Bot) Connect() {
	dest := bot.Server + ":" + bot.Port
	conn, err := net.Dial("tcp", dest)
	if err != nil {
		log.Fatalf("Unable to connect to %s\nError: %v\n\n", dest, err)
	}
	log.Printf("Successfully connected to %s\n", dest)
	bot.conn = conn
	bot.Send("USER " + bot.Nick + " 8 * :" + bot.Nick + "\n")
	bot.Send("NICK " + bot.Nick + "\n")
	bot.Send("JOIN " + bot.Channel + "\n")
}

// Sends a string to the server. Strings should end with a
// newline character.

func (bot *Bot) Send(str string) {
	msg := []byte(str)
	_, err := bot.conn.Write(msg)
	if err != nil {
		log.Printf("Error sending: %s", msg)
		log.Printf("Error: %v", err)
	} else {
		log.Printf("Successfully sent: %s", msg)
	}
}

// Sends a message to the channel

func (bot *Bot) MsgChannel(line string) {
	bot.Send("PRIVMSG " + bot.Channel + " :" + line + "\n")
}

func (bot *Bot) PrivateMsg(user string, line string) {
	bot.Send("PRIVMSG " + user + " :" + line + "\n")
}

func (bot *Bot) parse(line string) {
	for i := range bot.Ignore {
		if strings.Contains(line, bot.Ignore[i]) {
			return
		}
	}
	for i := range bot.actions {
		action := bot.actions[i]
		resp := action.Parse(action.re, line)
		if resp != nil {
			go action.Run(bot, resp)
		}
	}
}

func (bot *Bot) AddAction(re string, parse func(*regexp.Regexp, string) interface{},
	run func(*Bot, interface{})) {
	action := new(Action)
	action.re = regexp.MustCompile(re)
	action.Parse = parse
	action.Run = run
	bot.actions = append(bot.actions, action)
}

func signalHandling(bot *Bot) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	for sig := range c {
		if sig == os.Interrupt {
			log.Printf("Bot received os.Interrupt, exiting normally.\n\n")
			bot.Send("QUIT :\n")
			bot.conn.Close()
			os.Exit(0)
		}
	}
}

func (bot *Bot) Run() {
	go signalHandling(bot)
	bot.SetActions()
	t := bot.NewTimelineUpdater(bot.Trac_RSS, bot.Interval)
	go t.Run()
	reader := bufio.NewReader(bot.conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()
		if err != nil {
			log.Printf("Error reading line: %s\n", line)
			log.Fatalf("Error: %v\n", err)
		} else {
			bot.parse(line)
		}
	}
}

func ParseTicket(re *regexp.Regexp, line string) interface{} {
	//Regex: #(\d+)([:alpha:])*
	matches := re.FindAllStringSubmatch(line, 10)
	var ticket_nums []string
	for i := range matches {
		if matches[i][2] == "" {
			ticket_nums = append(ticket_nums, matches[i][1])
		}
	}
	if len(ticket_nums) == 0 {
		return nil
	}
	return ticket_nums
}

func FetchTickets(bot *Bot, tickets interface{}) {
	url := bot.Trac_URL + "ticket/"
	used := make(map[string]bool)
	ticket_nums := reflect.ValueOf(tickets)
	for i := 0; i < ticket_nums.Len(); i++ {
		num := ticket_nums.Index(i).String()
		if used[num] {
			continue
		} else {
			used[num] = true
		}
		dest := url + num
		resp, err := http.Get(dest)
		if err != nil {
			fmt.Printf("%v\n", err)
		}
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			fmt.Printf("%v\n", err)
		}
		s := string(body)
		if strings.Contains(s, "
Error:") {

		} else {
			s = s[:300]
			a := strings.Split(s, "\n")
			title := strings.TrimLeft(a[7], " ")
			bot.MsgChannel(title)
			bot.MsgChannel(dest)
		}
	}
}

func ParsePing(re *regexp.Regexp, line string) interface{} {
	//Regex: ^PING
	str := re.FindString(line)
	if str == "" {
		return nil
	}
	return line
}

func SendPong(bot *Bot, str interface{}) {
	line := reflect.ValueOf(str).String()
	log.Printf("Received: %s\n", line)
	strs := strings.Split(line, " ")
	bot.Send("PONG " + strs[1] + "\n")
}

func (bot *Bot) SetActions() {
	if bot.Trac_URL != "" {
		bot.AddAction(`#(\d+)([:alpha:])*`, ParseTicket, FetchTickets)
	}
	bot.AddAction(`^PING`, ParsePing, SendPong)
}

func InitLogging() {
	logf, err := os.OpenFile("bot_log.txt", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0667)
	if err != nil {
		log.Fatal("Cannot open/create log file.\nError: %v\n\n", err)
	}
	log.SetOutput(logf)
}
<commit_msg>Tally won't repeat tickets until 5 mins later.<commit_after>package tally

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/textproto"
	"os"
	"os/signal"
	"reflect"
	"regexp"
	"strings"
	"time"
)

// The Action struct represents a possible reaction to a given line.
// Regexps are attached to the Action so that they can be compiled once
// when the bot is initialized instead of each time a line needs to be parsed.

type Action struct {
	re    *regexp.Regexp
	Parse func(*regexp.Regexp, string) interface{}
	Run   func(*Bot, interface{})
}

type Bot struct {
	Server   string
	Port     string
	Nick     string
	Channel  string
	Trac_URL string
	Tickets  map[string]bool
	Trac_RSS string
	Interval time.Duration
	Ignore   []string
	conn     net.Conn

	actions []*Action
}

// Returns a new bot that has been configured according to
// config.json

func NewBot() *Bot {
	bot := new(Bot)
	f, err := ioutil.ReadFile("./config.json")
	if err != nil {
		log.Fatal("Error loading config.json: %v\n\n", err)
	}
	err = json.Unmarshal(f, bot)
	if err != nil {
		log.Fatal("Error unmarshalling config.json: %v\n\n", err)
	}
	bot.Tickets = make(map[string]bool)
	return bot
}

// Establishes a connection to the server and joins a channel

func (bot *Bot) Connect() {
	dest := bot.Server + ":" + bot.Port
	conn, err := net.Dial("tcp", dest)
	if err != nil {
		log.Fatalf("Unable to connect to %s\nError: %v\n\n", dest, err)
	}
	log.Printf("Successfully connected to %s\n", dest)
	bot.conn = conn
	bot.Send("USER " + bot.Nick + " 8 * :" + bot.Nick + "\n")
	bot.Send("NICK " + bot.Nick + "\n")
	bot.Send("JOIN " + bot.Channel + "\n")
}

// Sends a string to the server. Strings should end with a
// newline character.

func (bot *Bot) Send(str string) {
	msg := []byte(str)
	_, err := bot.conn.Write(msg)
	if err != nil {
		log.Printf("Error sending: %s", msg)
		log.Printf("Error: %v", err)
	} else {
		log.Printf("Successfully sent: %s", msg)
	}
}

// Sends a message to the channel

func (bot *Bot) MsgChannel(line string) {
	bot.Send("PRIVMSG " + bot.Channel + " :" + line + "\n")
}

func (bot *Bot) PrivateMsg(user string, line string) {
	bot.Send("PRIVMSG " + user + " :" + line + "\n")
}

func (bot *Bot) parse(line string) {
	for i := range bot.Ignore {
		if strings.Contains(line, bot.Ignore[i]) {
			return
		}
	}
	for i := range bot.actions {
		action := bot.actions[i]
		resp := action.Parse(action.re, line)
		if resp != nil {
			go action.Run(bot, resp)
		}
	}
}

func (bot *Bot) AddAction(re string, parse func(*regexp.Regexp, string) interface{},
	run func(*Bot, interface{})) {
	action := new(Action)
	action.re = regexp.MustCompile(re)
	action.Parse = parse
	action.Run = run
	bot.actions = append(bot.actions, action)
}

func signalHandling(bot *Bot) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	for sig := range c {
		if sig == os.Interrupt {
			log.Printf("Bot received os.Interrupt, exiting normally.\n\n")
			bot.Send("QUIT :\n")
			bot.conn.Close()
			os.Exit(0)
		}
	}
}

func (bot *Bot) Run() {
	go signalHandling(bot)
	bot.SetActions()
	t := bot.NewTimelineUpdater(bot.Trac_RSS, bot.Interval)
	go t.Run()
	reader := bufio.NewReader(bot.conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()
		if err != nil {
			log.Printf("Error reading line: %s\n", line)
			log.Fatalf("Error: %v\n", err)
		} else {
			bot.parse(line)
		}
	}
}

func ParseTicket(re *regexp.Regexp, line string) interface{} {
	//Regex: #(\d+)([:alpha:])*
	matches := re.FindAllStringSubmatch(line, 10)
	var ticket_nums []string
	for i := range matches {
		if matches[i][2] == "" {
			ticket_nums = append(ticket_nums, matches[i][1])
		}
	}
	if len(ticket_nums) == 0 {
		return nil
	}
	return ticket_nums
}

func removeTicket(bot *Bot, num string) {
	time.Sleep(5 * time.Minute)
	bot.Tickets[num] = false
}

func FetchTickets(bot *Bot, tickets interface{}) {
	url := bot.Trac_URL + "ticket/"
	ticket_nums := reflect.ValueOf(tickets)
	for i := 0; i < ticket_nums.Len(); i++ {
		num := ticket_nums.Index(i).String()

		// Check to see if we gave a link to this ticket recently.
		if bot.Tickets[num] == false {
			bot.Tickets[num] = true
			go removeTicket(bot, num)
		} else {
			continue
		}

		dest := url + num
		resp, err := http.Get(dest)
		if err != nil {
			fmt.Printf("%v\n", err)
		}
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			fmt.Printf("%v\n", err)
		}
		s := string(body)
		if strings.Contains(s, "
Error:") {

		} else {
			s = s[:300]
			a := strings.Split(s, "\n")
			title := strings.TrimLeft(a[7], " ")
			bot.MsgChannel(title)
			bot.MsgChannel(dest)
		}
	}
}

func ParsePing(re *regexp.Regexp, line string) interface{} {
	//Regex: ^PING
	str := re.FindString(line)
	if str == "" {
		return nil
	}
	return line
}

func SendPong(bot *Bot, str interface{}) {
	line := reflect.ValueOf(str).String()
	log.Printf("Received: %s\n", line)
	strs := strings.Split(line, " ")
	bot.Send("PONG " + strs[1] + "\n")
}

func (bot *Bot) SetActions() {
	if bot.Trac_URL != "" {
		bot.AddAction(`#(\d+)([:alpha:])*`, ParseTicket, FetchTickets)
	}
	bot.AddAction(`^PING`, ParsePing, SendPong)
}

func InitLogging() {
	logf, err := os.OpenFile("bot_log.txt", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0667)
	if err != nil {
		log.Fatal("Cannot open/create log file.\nError: %v\n\n", err)
	}
	log.SetOutput(logf)
}
<|endoftext|>


<commit_before>// Copyright (C) 2013 Space Monkey, Inc.

package client

import (
	"flag"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"code.spacemonkey.com/go/errors"
	space_time "code.spacemonkey.com/go/space/time"
)

var (
	maxErrorLength = flag.Int("monitor.max_error_length", 40,
		"the max length for an error name")
)

type TaskMonitor struct {
	mtx             sync.Mutex
	current         uint64
	highwater       uint64
	total_started   uint64
	total_completed uint64
	success         uint64
	timing          *ValueMonitor
	errors          map[string]uint64
}

func NewTaskMonitor() *TaskMonitor {
	return &TaskMonitor{
		errors: make(map[string]uint64),
		timing: NewValueMonitor()}
}

type TaskCtx struct {
	start   time.Duration
	monitor *TaskMonitor
}

func (t *TaskMonitor) Start() *TaskCtx {
	t.mtx.Lock()
	t.current += 1
	t.total_started += 1
	if t.current > t.highwater {
		t.highwater = t.current
	}
	t.mtx.Unlock()
	return &TaskCtx{start: space_time.Monotonic(), monitor: t}
}

func (t *TaskMonitor) Stats(cb func(name string, val float64)) {
	t.mtx.Lock()
	current := t.current
	highwater := t.highwater
	total_started := t.total_started
	total_completed := t.total_completed
	success := t.success
	error_counts := make(map[string]uint64, len(t.errors))
	for error, count := range t.errors {
		error_counts[error] = count
	}
	t.mtx.Unlock()

	errors := make([]string, 0, len(error_counts))
	for error := range error_counts {
		errors = append(errors, error)
	}
	sort.Strings(errors)

	cb("current", float64(current))
	for _, error := range errors {
		cb(fmt.Sprintf("error_%s", error), float64(error_counts[error]))
	}
	cb("highwater", float64(highwater))
	cb("success", float64(success))
	t.timing.Stats(func(name string, val float64) {
		if name != "count" {
			cb(fmt.Sprintf("time_%s", name), val)
		}
	})
	cb("total_completed", float64(total_completed))
	cb("total_started", float64(total_started))
}

func (c *TaskCtx) Finish(err_ref *error) {
	duration := space_time.Monotonic() - c.start
	var error_name string
	var err error
	if err_ref != nil {
		err = *err_ref
	}
	rec := recover()
	if rec != nil {
		var ok bool
		err, ok = rec.(error)
		if !ok {
			err = errors.PanicError.New("%v", rec)
		}
	}
	if err != nil {
		error_name = errors.GetClass(err).String()
		if len(error_name) > *maxErrorLength {
			error_name = error_name[:*maxErrorLength]
		}
		error_name = SanitizeName(error_name)
	}

	c.monitor.mtx.Lock()
	c.monitor.current -= 1
	c.monitor.total_completed += 1
	if err != nil {
		c.monitor.errors[error_name] += 1
	} else {
		c.monitor.success += 1
	}
	c.monitor.mtx.Unlock()
	c.monitor.timing.Add(duration.Seconds())

	// doh, we didn't actually want to stop the panic codepath.
	// we have to repanic
	if rec != nil {
		panic(rec)
	}
}

func (self *MonitorGroup) Task() func(*error) {
	caller_name := CallerName()
	idx := strings.LastIndex(caller_name, "/")
	if idx >= 0 {
		caller_name = caller_name[idx+1:]
	}
	idx = strings.Index(caller_name, ".")
	if idx >= 0 {
		caller_name = caller_name[idx+1:]
	}
	return self.TaskNamed(caller_name)
}

func (self *MonitorGroup) TaskNamed(name string) func(*error) {
	name = SanitizeName(name)
	monitor, err := self.monitors.Get(name, func(_ interface{}) (interface{}, error) {
		return NewTaskMonitor(), nil
	})
	if err != nil {
		handleError(err)
		return func(*error) {}
	}
	task_monitor, ok := monitor.(*TaskMonitor)
	if !ok {
		handleError(errors.ProgrammerError.New(
			"monitor already exists with different type for name %s", name))
		return func(*error) {}
	}
	ctx := task_monitor.Start()
	return ctx.Finish
}
<commit_msg>space monkey internal commit export<commit_after>// Copyright (C) 2013 Space Monkey, Inc.

package client

import (
	"flag"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"code.spacemonkey.com/go/errors"
	space_time "code.spacemonkey.com/go/space/time"
)

var (
	maxErrorLength = flag.Int("monitor.max_error_length", 40,
		"the max length for an error name")
)

type TaskMonitor struct {
	mtx             sync.Mutex
	current         uint64
	highwater       uint64
	total_started   uint64
	total_completed uint64
	success         uint64
	timing          *ValueMonitor
	errors          map[string]uint64
	panics          uint64
}

func NewTaskMonitor() *TaskMonitor {
	return &TaskMonitor{
		errors: make(map[string]uint64),
		timing: NewValueMonitor()}
}

type TaskCtx struct {
	start   time.Duration
	monitor *TaskMonitor
}

func (t *TaskMonitor) Start() *TaskCtx {
	t.mtx.Lock()
	t.current += 1
	t.total_started += 1
	if t.current > t.highwater {
		t.highwater = t.current
	}
	t.mtx.Unlock()
	return &TaskCtx{start: space_time.Monotonic(), monitor: t}
}

func (t *TaskMonitor) Stats(cb func(name string, val float64)) {
	t.mtx.Lock()
	current := t.current
	highwater := t.highwater
	total_started := t.total_started
	total_completed := t.total_completed
	success := t.success
	panics := t.panics
	error_counts := make(map[string]uint64, len(t.errors))
	for error, count := range t.errors {
		error_counts[error] = count
	}
	t.mtx.Unlock()

	errors := make([]string, 0, len(error_counts))
	for error := range error_counts {
		errors = append(errors, error)
	}
	sort.Strings(errors)

	cb("current", float64(current))
	for _, error := range errors {
		cb(fmt.Sprintf("error_%s", error), float64(error_counts[error]))
	}
	cb("highwater", float64(highwater))
	cb("panics", float64(panics))
	cb("success", float64(success))
	t.timing.Stats(func(name string, val float64) {
		if name != "count" {
			cb(fmt.Sprintf("time_%s", name), val)
		}
	})
	cb("total_completed", float64(total_completed))
	cb("total_started", float64(total_started))
}

func (c *TaskCtx) Finish(err_ref *error) {
	duration := space_time.Monotonic() - c.start
	var error_name string
	var err error
	if err_ref != nil {
		err = *err_ref
	}
	rec := recover()
	if rec != nil {
		var ok bool
		err, ok = rec.(error)
		if !ok || err == nil {
			err = errors.PanicError.New("%v", rec)
		}
	}
	if err != nil {
		error_name = errors.GetClass(err).String()
		if len(error_name) > *maxErrorLength {
			error_name = error_name[:*maxErrorLength]
		}
		error_name = SanitizeName(error_name)
	}

	c.monitor.mtx.Lock()
	c.monitor.current -= 1
	c.monitor.total_completed += 1
	if err != nil {
		c.monitor.errors[error_name] += 1
		if rec != nil {
			c.monitor.panics += 1
		}
	} else {
		c.monitor.success += 1
	}
	c.monitor.mtx.Unlock()
	c.monitor.timing.Add(duration.Seconds())

	// doh, we didn't actually want to stop the panic codepath.
	// we have to repanic
	if rec != nil {
		panic(rec)
	}
}

func (self *MonitorGroup) Task() func(*error) {
	caller_name := CallerName()
	idx := strings.LastIndex(caller_name, "/")
	if idx >= 0 {
		caller_name = caller_name[idx+1:]
	}
	idx = strings.Index(caller_name, ".")
	if idx >= 0 {
		caller_name = caller_name[idx+1:]
	}
	return self.TaskNamed(caller_name)
}

func (self *MonitorGroup) TaskNamed(name string) func(*error) {
	name = SanitizeName(name)
	monitor, err := self.monitors.Get(name, func(_ interface{}) (interface{}, error) {
		return NewTaskMonitor(), nil
	})
	if err != nil {
		handleError(err)
		return func(*error) {}
	}
	task_monitor, ok := monitor.(*TaskMonitor)
	if !ok {
		handleError(errors.ProgrammerError.New(
			"monitor already exists with different type for name %s", name))
		return func(*error) {}
	}
	ctx := task_monitor.Start()
	return ctx.Finish
}
<|endoftext|>


<commit_before>//
//  telgo
//
//
//  Copyright (C) 2015 Christian Pointner <equinox@helsinki.at>
//
//  This file is part of telgo.
//
//  telgo is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  any later version.
//
//  telgo is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with telgo. If not, see <http://www.gnu.org/licenses/>.
//

// Package telgo contains a simple telnet server which can be used as a
// control/debug interface for applications.
// The telgo telnet server does all the client handling and runs configurable
// commands as go routines. It also supports handling of basic inline telnet
// commands used by variaus telnet clients to configure the client connection.
// For now every negotiable telnet option will be discarded but the telnet
// command IP (interrupt process) is understood and can be used to terminate
// long running user commands.
package telgo

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
)

var (
	tl = log.New(os.Stderr, "[telnet]\t", log.LstdFlags)
)

const (
	EOT  = byte(4)
	IP   = byte(244)
	WILL = byte(251)
	WONT = byte(252)
	DO   = byte(253)
	DONT = byte(254)
	IAC  = byte(255)
)

// This is the signature of telgo command functions. It receives a pointer to
// the telgo client struct and a slice of strings containing the arguments the
// user has supplied. The first argument is always the command name itself.
// If this function returns true the client connection will be terminated.
type TelgoCmd func(c *TelnetClient, args []string) bool
type TelgoCmdList map[string]TelgoCmd

// This struct is used to export the raw tcp connection to the client as well as
// the UserData which got supplied to NewTelnetServer.
// The Cancel channel will get ready for reading when the user hits Ctrl-C or
// the connection got terminated. This can be used for long running telgo commands
// to be aborted.
type TelnetClient struct {
	Conn     net.Conn
	UserData interface{}
	Cancel   chan bool
	scanner  *bufio.Scanner
	writer   *bufio.Writer
	prompt   string
	commands *TelgoCmdList
	iacout   chan []byte
	stdout   chan []byte
}

func newTelnetClient(conn net.Conn, prompt string, commands *TelgoCmdList, userdata interface{}) (c *TelnetClient) {
	tl.Println("telgo: new client from:", conn.RemoteAddr())
	c = &TelnetClient{}
	c.Conn = conn
	c.scanner = bufio.NewScanner(conn)
	c.writer = bufio.NewWriter(conn)
	c.prompt = prompt
	c.commands = commands
	c.UserData = userdata
	c.stdout = make(chan []byte)
	c.Cancel = make(chan bool, 1)
	// the telnet split function needs some closures to handle OOB telnet commands
	c.iacout = make(chan []byte)
	lastiiac := 0
	c.scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		return scanLines(data, atEOF, c.iacout, &lastiiac)
	})
	return c
}

// This writes a 'raw' string to the client. For the most part the usage of Say
// and Sayln is recommended. WriteString will take care of escaping IAC bytes
// inside your string.
func (c *TelnetClient) WriteString(text string) {
	c.stdout <- bytes.Replace([]byte(text), []byte{IAC}, []byte{IAC, IAC}, -1)
}

// This is a simple Printf like interface which sends responses to the client.
func (c *TelnetClient) Say(format string, a ...interface{}) {
	c.WriteString(fmt.Sprintf(format, a...))
}

// This is the same as Say but also adds a new-line at the end of the string.
func (c *TelnetClient) Sayln(format string, a ...interface{}) {
	c.WriteString(fmt.Sprintf(format, a...) + "\r\n")
}

// TODO: fix split function to respect "" and ''
func (c *TelnetClient) handleCmd(cmdstr string, done chan<- bool) {
	quit := false
	defer func() { done <- quit }()

	cmdslice := strings.Fields(cmdstr)
	if len(cmdslice) == 0 || cmdslice[0] == "" {
		return
	}

	select {
	case <-c.Cancel: // consume potentially pending cancel request
	default:
	}
	for cmd, cmdfunc := range *c.commands {
		if cmdslice[0] == cmd {
			quit = cmdfunc(c, cmdslice)
			return
		}
	}
	c.Sayln("unknown command '%s'", cmdslice[0])
}

func handleIac(iac []byte, iacout chan<- []byte) {
	switch iac[1] {
	case WILL, WONT: // Don't accept any proposed options
		iac[1] = DONT
	case DO, DONT:
		iac[1] = WONT
	case IP:
		// pass this through to client.handle which will cancel the process
	default:
		tl.Printf("ignoring unimplemented telnet command: %X", iac[1])
		return
	}
	iacout <- iac
}

func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

func dropIAC(data []byte) []byte {
	var token []byte
	iiac := 0
	for {
		niiac := bytes.IndexByte(data[iiac:], IAC)
		if niiac >= 0 {
			token = append(token, data[iiac:iiac+niiac]...)
			iiac += niiac
			if (len(data) - iiac) < 2 {
				return token
			}
			switch data[iiac+1] {
			case DONT, DO, WONT, WILL:
				if (len(data) - iiac) < 3 {
					return token
				}
				iiac += 3
			case IAC:
				token = append(token, IAC)
				fallthrough
			default:
				iiac += 2
			}
		} else {
			token = append(token, data[iiac:]...)
			break
		}
	}
	return token
}

func compareIdx(a, b int) int {
	if a < 0 {
		a = int(^uint(0) >> 1)
	}
	if b < 0 {
		b = int(^uint(0) >> 1)
	}
	return a - b
}

func scanLines(data []byte, atEOF bool, iacout chan<- []byte, lastiiac *int) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}

	inl := bytes.IndexByte(data, '\n') // index of first newline character
	ieot := bytes.IndexByte(data, EOT) // index of first End of Transmission

	iiac := *lastiiac
	for {
		niiac := bytes.IndexByte(data[iiac:], IAC) // index of first/next telnet IAC
		if niiac >= 0 {
			iiac += niiac
		} else {
			iiac = niiac
		}

		if inl >= 0 && compareIdx(inl, ieot) < 0 && compareIdx(inl, iiac) < 0 {
			*lastiiac = 0
			return inl + 1, dropCR(dropIAC(data[0:inl])), nil // found a complete line
		}
		if ieot >= 0 && compareIdx(ieot, iiac) < 0 {
			*lastiiac = 0
			return ieot + 1, data[ieot : ieot+1], nil // found a EOT (aka Ctrl-D)
		}
		if iiac >= 0 {
			l := 2
			if (len(data) - iiac) < 2 {
				return 0, nil, nil // data does not yet contain the telnet command code -> need more data
			}
			switch data[iiac+1] {
			case DONT, DO, WONT, WILL:
				if (len(data) - iiac) < 3 {
					return 0, nil, nil // this is a 3-byte command and data does not yet contain the option code -> need more data
				}
				l = 3
			case IAC:
				iiac += 2
				continue
			}
			handleIac(data[iiac:iiac+l], iacout)
			iiac += l
			*lastiiac = iiac
		} else {
			break
		}
	}
	if atEOF {
		return len(data), dropCR(data), nil // allow last line to have no new line
	}
	return 0, nil, nil // we have found none of the escape codes -> need more data
}

func (c *TelnetClient) recv(in chan<- string) {
	defer close(in)

	for c.scanner.Scan() {
		b := c.scanner.Bytes()
		if len(b) > 0 && b[0] == EOT {
			tl.Printf("telgo(%s): Ctrl-D received, closing", c.Conn.RemoteAddr())
			return
		}
		in <- string(b)
	}
	if err := c.scanner.Err(); err != nil {
		tl.Printf("telgo(%s): recv() error: %s", c.Conn.RemoteAddr(), err)
	} else {
		tl.Printf("telgo(%s): Connection closed by foreign host", c.Conn.RemoteAddr())
	}
}

func (c *TelnetClient) cancel() {
	select {
	case c.Cancel <- true:
	default: // process got canceled already
	}
}

func (c *TelnetClient) send(quit <-chan bool) {
	for {
		select {
		case <-quit:
			return
		case iac := <-c.iacout:
			if iac[1] == IP {
				c.cancel()
			} else {
				c.writer.Write(iac)
				c.writer.Flush()
			}
		case data := <-c.stdout:
			c.writer.Write(data)
			c.writer.Flush()
		}
	}
}

func (c *TelnetClient) handle() {
	defer c.Conn.Close()

	in := make(chan string)
	go c.recv(in)

	quit_send := make(chan bool)
	go c.send(quit_send)
	defer func() { quit_send <- true }()

	defer c.cancel() // make sure to cancel possible running job when closing connection

	var cmd_backlog []string
	done := make(chan bool)
	busy := false
	c.WriteString(c.prompt)
	for {
		select {
		case cmd, ok := <-in:
			if !ok {
				return
			}
			if len(cmd) > 0 {
				if !busy {
					go c.handleCmd(cmd, done)
					busy = true
				} else {
					cmd_backlog = append(cmd_backlog, cmd)
				}
			} else {
				c.WriteString(c.prompt)
			}
		case exit := <-done:
			if exit {
				return
			}
			c.WriteString(c.prompt)
			if len(cmd_backlog) > 0 {
				go c.handleCmd(cmd_backlog[0], done)
				cmd_backlog = cmd_backlog[1:]
			} else {
				busy = false
			}
		}
	}
}

type TelnetServer struct {
	addr     string
	prompt   string
	commands TelgoCmdList
	userdata interface{}
}

// This creates a new telnet server. addr is the address to bind/listen to on and will be passed through
// to net.Listen(). The prompt will be sent to the client whenever the telgo server is ready for new command
// TelgoCmdList contains a list of available commands and userdata will be made available to called telgo
// commands through the client struct.
func NewTelnetServer(addr, prompt string, commands TelgoCmdList, userdata interface{}) (s *TelnetServer) {
	s = &TelnetServer{}
	s.addr = addr
	s.prompt = prompt
	s.commands = commands
	s.userdata = userdata
	return s
}

// This runs the telnet server and spawns go routines for every connecting client.
func (self *TelnetServer) Run() error {
	tl.Println("telgo: listening on", self.addr)

	server, err := net.Listen("tcp", self.addr)
	if err != nil {
		tl.Println("telgo: Listen() Error:", err)
		return err
	}

	for {
		conn, err := server.Accept()
		if err != nil {
			tl.Println("telgo: Accept() Error:", err)
			return err
		}

		c := newTelnetClient(conn, self.prompt, &self.commands, self.userdata)
		go c.handle()
	}
}
<commit_msg>fixed docs<commit_after>//
//  telgo
//
//
//  Copyright (C) 2015 Christian Pointner <equinox@helsinki.at>
//
//  This file is part of telgo.
//
//  telgo is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  any later version.
//
//  telgo is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with telgo. If not, see <http://www.gnu.org/licenses/>.
//

// Package telgo contains a simple telnet server which can be used as a
// control/debug interface for applications.
// The telgo telnet server does all the client handling and runs configurable
// commands as go routines. It also supports handling of basic inline telnet
// commands used by variaus telnet clients to configure the client connection.
// For now every negotiable telnet option will be discarded but the telnet
// command IP (interrupt process) is understood and can be used to terminate
// long running user commands.
package telgo

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
)

var (
	tl = log.New(os.Stderr, "[telnet]\t", log.LstdFlags)
)

const (
	EOT  = byte(4)
	IP   = byte(244)
	WILL = byte(251)
	WONT = byte(252)
	DO   = byte(253)
	DONT = byte(254)
	IAC  = byte(255)
)

// This is the signature of telgo command functions. It receives a pointer to
// the telgo client struct and a slice of strings containing the arguments the
// user has supplied. The first argument is always the command name itself.
// If this function returns true the client connection will be terminated.
type TelgoCmd func(c *TelnetClient, args []string) bool
type TelgoCmdList map[string]TelgoCmd

// This struct is used to export the raw tcp connection to the client as well as
// the UserData which got supplied to NewTelnetServer.
// The Cancel channel will get ready for reading when the user hits Ctrl-C or
// the connection got terminated. This can be used for long running telgo commands
// to be aborted.
type TelnetClient struct {
	Conn     net.Conn
	UserData interface{}
	Cancel   chan bool
	scanner  *bufio.Scanner
	writer   *bufio.Writer
	prompt   string
	commands *TelgoCmdList
	iacout   chan []byte
	stdout   chan []byte
}

func newTelnetClient(conn net.Conn, prompt string, commands *TelgoCmdList, userdata interface{}) (c *TelnetClient) {
	tl.Println("telgo: new client from:", conn.RemoteAddr())
	c = &TelnetClient{}
	c.Conn = conn
	c.scanner = bufio.NewScanner(conn)
	c.writer = bufio.NewWriter(conn)
	c.prompt = prompt
	c.commands = commands
	c.UserData = userdata
	c.stdout = make(chan []byte)
	c.Cancel = make(chan bool, 1)
	// the telnet split function needs some closures to handle OOB telnet commands
	c.iacout = make(chan []byte)
	lastiiac := 0
	c.scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		return scanLines(data, atEOF, c.iacout, &lastiiac)
	})
	return c
}

// This writes a 'raw' string to the client. For the most part the usage of Say
// and Sayln is recommended. WriteString will take care of escaping IAC bytes
// inside your string.
func (c *TelnetClient) WriteString(text string) {
	c.stdout <- bytes.Replace([]byte(text), []byte{IAC}, []byte{IAC, IAC}, -1)
}

// This is a simple Printf like interface which sends responses to the client.
func (c *TelnetClient) Say(format string, a ...interface{}) {
	c.WriteString(fmt.Sprintf(format, a...))
}

// This is the same as Say but also adds a new-line at the end of the string.
func (c *TelnetClient) Sayln(format string, a ...interface{}) {
	c.WriteString(fmt.Sprintf(format, a...) + "\r\n")
}

// TODO: fix split function to respect "" and ''
func (c *TelnetClient) handleCmd(cmdstr string, done chan<- bool) {
	quit := false
	defer func() { done <- quit }()

	cmdslice := strings.Fields(cmdstr)
	if len(cmdslice) == 0 || cmdslice[0] == "" {
		return
	}

	select {
	case <-c.Cancel: // consume potentially pending cancel request
	default:
	}
	for cmd, cmdfunc := range *c.commands {
		if cmdslice[0] == cmd {
			quit = cmdfunc(c, cmdslice)
			return
		}
	}
	c.Sayln("unknown command '%s'", cmdslice[0])
}

func handleIac(iac []byte, iacout chan<- []byte) {
	switch iac[1] {
	case WILL, WONT: // Don't accept any proposed options
		iac[1] = DONT
	case DO, DONT:
		iac[1] = WONT
	case IP:
		// pass this through to client.handle which will cancel the process
	default:
		tl.Printf("ignoring unimplemented telnet command: %X", iac[1])
		return
	}
	iacout <- iac
}

func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

func dropIAC(data []byte) []byte {
	var token []byte
	iiac := 0
	for {
		niiac := bytes.IndexByte(data[iiac:], IAC)
		if niiac >= 0 {
			token = append(token, data[iiac:iiac+niiac]...)
			iiac += niiac
			if (len(data) - iiac) < 2 {
				return token
			}
			switch data[iiac+1] {
			case DONT, DO, WONT, WILL:
				if (len(data) - iiac) < 3 {
					return token
				}
				iiac += 3
			case IAC:
				token = append(token, IAC)
				fallthrough
			default:
				iiac += 2
			}
		} else {
			token = append(token, data[iiac:]...)
			break
		}
	}
	return token
}

func compareIdx(a, b int) int {
	if a < 0 {
		a = int(^uint(0) >> 1)
	}
	if b < 0 {
		b = int(^uint(0) >> 1)
	}
	return a - b
}

func scanLines(data []byte, atEOF bool, iacout chan<- []byte, lastiiac *int) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}

	inl := bytes.IndexByte(data, '\n') // index of first newline character
	ieot := bytes.IndexByte(data, EOT) // index of first End of Transmission

	iiac := *lastiiac
	for {
		niiac := bytes.IndexByte(data[iiac:], IAC) // index of first/next telnet IAC
		if niiac >= 0 {
			iiac += niiac
		} else {
			iiac = niiac
		}

		if inl >= 0 && compareIdx(inl, ieot) < 0 && compareIdx(inl, iiac) < 0 {
			*lastiiac = 0
			return inl + 1, dropCR(dropIAC(data[0:inl])), nil // found a complete line
		}
		if ieot >= 0 && compareIdx(ieot, iiac) < 0 {
			*lastiiac = 0
			return ieot + 1, data[ieot : ieot+1], nil // found a EOT (aka Ctrl-D)
		}
		if iiac >= 0 {
			l := 2
			if (len(data) - iiac) < 2 {
				return 0, nil, nil // data does not yet contain the telnet command code -> need more data
			}
			switch data[iiac+1] {
			case DONT, DO, WONT, WILL:
				if (len(data) - iiac) < 3 {
					return 0, nil, nil // this is a 3-byte command and data does not yet contain the option code -> need more data
				}
				l = 3
			case IAC:
				iiac += 2
				continue
			}
			handleIac(data[iiac:iiac+l], iacout)
			iiac += l
			*lastiiac = iiac
		} else {
			break
		}
	}
	if atEOF {
		return len(data), dropCR(data), nil // allow last line to have no new line
	}
	return 0, nil, nil // we have found none of the escape codes -> need more data
}

func (c *TelnetClient) recv(in chan<- string) {
	defer close(in)

	for c.scanner.Scan() {
		b := c.scanner.Bytes()
		if len(b) > 0 && b[0] == EOT {
			tl.Printf("telgo(%s): Ctrl-D received, closing", c.Conn.RemoteAddr())
			return
		}
		in <- string(b)
	}
	if err := c.scanner.Err(); err != nil {
		tl.Printf("telgo(%s): recv() error: %s", c.Conn.RemoteAddr(), err)
	} else {
		tl.Printf("telgo(%s): Connection closed by foreign host", c.Conn.RemoteAddr())
	}
}

func (c *TelnetClient) cancel() {
	select {
	case c.Cancel <- true:
	default: // process got canceled already
	}
}

func (c *TelnetClient) send(quit <-chan bool) {
	for {
		select {
		case <-quit:
			return
		case iac := <-c.iacout:
			if iac[1] == IP {
				c.cancel()
			} else {
				c.writer.Write(iac)
				c.writer.Flush()
			}
		case data := <-c.stdout:
			c.writer.Write(data)
			c.writer.Flush()
		}
	}
}

func (c *TelnetClient) handle() {
	defer c.Conn.Close()

	in := make(chan string)
	go c.recv(in)

	quit_send := make(chan bool)
	go c.send(quit_send)
	defer func() { quit_send <- true }()

	defer c.cancel() // make sure to cancel possible running job when closing connection

	var cmd_backlog []string
	done := make(chan bool)
	busy := false
	c.WriteString(c.prompt)
	for {
		select {
		case cmd, ok := <-in:
			if !ok {
				return
			}
			if len(cmd) > 0 {
				if !busy {
					go c.handleCmd(cmd, done)
					busy = true
				} else {
					cmd_backlog = append(cmd_backlog, cmd)
				}
			} else {
				c.WriteString(c.prompt)
			}
		case exit := <-done:
			if exit {
				return
			}
			c.WriteString(c.prompt)
			if len(cmd_backlog) > 0 {
				go c.handleCmd(cmd_backlog[0], done)
				cmd_backlog = cmd_backlog[1:]
			} else {
				busy = false
			}
		}
	}
}

type TelnetServer struct {
	addr     string
	prompt   string
	commands TelgoCmdList
	userdata interface{}
}

// This creates a new telnet server. addr is the address to bind/listen to on and will be passed through
// to net.Listen(). The prompt will be sent to the client whenever the telgo server is ready for a new command.
// TelgoCmdList is a list of available commands and userdata will be made available to called telgo
// commands through the client struct.
func NewTelnetServer(addr, prompt string, commands TelgoCmdList, userdata interface{}) (s *TelnetServer) {
	s = &TelnetServer{}
	s.addr = addr
	s.prompt = prompt
	s.commands = commands
	s.userdata = userdata
	return s
}

// This runs the telnet server and spawns go routines for every connecting client.
func (self *TelnetServer) Run() error {
	tl.Println("telgo: listening on", self.addr)

	server, err := net.Listen("tcp", self.addr)
	if err != nil {
		tl.Println("telgo: Listen() Error:", err)
		return err
	}

	for {
		conn, err := server.Accept()
		if err != nil {
			tl.Println("telgo: Accept() Error:", err)
			return err
		}

		c := newTelnetClient(conn, self.prompt, &self.commands, self.userdata)
		go c.handle()
	}
}
<|endoftext|>