<commit_before>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package client

import (
	"html/template"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	"github.com/caixw/gitype/data"
	"github.com/caixw/gitype/helper"
	"github.com/caixw/gitype/vars"
	"github.com/issue9/logs"
	"github.com/issue9/utils"
)

const contentTypeKey = "Content-Type"

// 生成一个带编码的 content-type 报头内容
func buildContentTypeContent(mime string) string {
	return mime + ";charset=utf-8"
}

// 设置页面的编码，若已经存在，则不会受影响。
// 要强制指定，请直接使用 w.Header().Set()
func setContentType(w http.ResponseWriter, mime string) {
	h := w.Header()
	if len(h.Get(contentTypeKey)) == 0 {
		h.Set(contentTypeKey, buildContentTypeContent(mime))
	}
}

// 用于描述一个页面的所有无素
type page struct {
	client   *Client
	Info     *info
	template *template.Template // 用于当前页面渲染的模板
	response http.ResponseWriter
	request  *http.Request

	Title       string       // 文章标题，可以为空
	Subtitle    string       // 副标题
	Canonical   string       // 当前页的唯一链接
	Keywords    string       // meta.keywords 的值
	Description string       // meta.description 的值
	PrevPage    *data.Link   // 前一页
	NextPage    *data.Link   // 下一页
	Type        string       // 当前页面类型
	Author      *data.Author // 作者
	License     *data.Link   // 当前页的版本信息，可以为空
	Theme       *data.Theme

	// 以下内容，仅在对应的页面才会有内容
	Q        string          // 搜索关键字
	Tag      *data.Tag       // 标签详细页面，非标签详细页，则为空
	Posts    []*data.Post    // 文章列表，仅标签详情页和搜索页用到。
	Post     *data.Post      // 文章详细内容，仅文章页面用到。
	Archives []*data.Archive // 归档
}

// 页面的附加信息，除非重新加载数据，否则内容不会变。
type info struct {
	AppName    string // 程序名称
	AppURL     string // 程序官网
	AppVersion string // 当前程序的版本号
	GoVersion  string // 编译的 Go 版本号

	SiteName    string     // 网站名称
	URL         string     // 网站地址，若是一个子目录，则需要包含该子目录
	Icon        *data.Icon // 网站图标
	Language    string     // 页面语言
	PostSize    int        // 总文章数量
	Beian       string     // 备案号
	Uptime      time.Time  // 上线时间
	LastUpdated time.Time  // 最后更新时间
	RSS         *data.Link // RSS，NOTICE:指针方便模板判断其值是否为空
	Atom        *data.Link
	Opensearch  *data.Link
	Tags        []*data.Tag  // 标签列表
	Series      []*data.Tag  // 专题列表
	Links       []*data.Link // 友情链接
	Menus       []*data.Link // 导航菜单
}

func (client *Client) newInfo() *info {
	d := client.data

	info := &info{
		AppName:    vars.Name,
		AppURL:     vars.URL,
		AppVersion: vars.Version(),
		GoVersion:  runtime.Version(),

		SiteName:    d.Title,
		URL:         d.URL,
		Icon:        d.Icon,
		Language:    d.Language,
		PostSize:    len(d.Posts),
		Beian:       d.Beian,
		Uptime:      d.Uptime,
		LastUpdated: d.Created,
		Tags:        d.Tags,
		Series:      d.Series,
		Links:       d.Links,
		Menus:       d.Menus,
	}

	if d.RSS != nil {
		info.RSS = &data.Link{
			Title: d.RSS.Title,
			URL:   d.RSS.URL,
			Type:  d.RSS.Type,
		}
	}

	if d.Atom != nil {
		info.Atom = &data.Link{
			Title: d.Atom.Title,
			URL:   d.Atom.URL,
			Type:  d.Atom.Type,
		}
	}

	if d.Opensearch != nil {
		info.Opensearch = &data.Link{
			Title: d.Opensearch.Title,
			URL:   d.Opensearch.URL,
			Type:  d.Opensearch.Type,
		}
	}

	return info
}

func (client *Client) page(typ string, w http.ResponseWriter, r *http.Request) *page {
	theme := client.getRequestTheme(r)
	d := client.data

	return &page{
		client:   client,
		Info:     client.info,
		template: theme.Template,
		response: w,
		request:  r,

		Subtitle:    d.Subtitle,
		Keywords:    d.Keywords,
		Description: d.Description,
		Type:        typ,
		Author:      d.Author,
		License:     d.License,
		Theme:       theme,
	}
}

func (p *page) nextPage(url, text string) {
	if len(text) == 0 {
		text = vars.NextPageText
	}

	p.NextPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "next",
	}
}

func (p *page) prevPage(url, text string) {
	if len(text) == 0 {
		text = vars.PrevPageText
	}

	p.PrevPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "prev",
	}
}

// 输出当前内容到指定模板
func (p *page) render(name string) {
	setContentType(p.response, p.client.data.Type)

	cookie := &http.Cookie{
		Name:     vars.CookieKeyTheme,
		Value:    p.Theme.ID,
		HttpOnly: vars.CookieHTTPOnly,
	}
	if p.Theme.ID != p.client.data.Themes[0].ID {
		cookie.MaxAge = vars.CookieMaxAge
	} else {
		cookie.MaxAge = -1
	}
	cookie.Expires = time.Now().Add(time.Second * time.Duration(vars.CookieMaxAge))
	http.SetCookie(p.response, cookie)

	err := p.template.ExecuteTemplate(p.response, name, p)
	if err != nil {
		logs.Error(err)
		p.client.renderError(p.response, p.request, http.StatusInternalServerError)
		return
	}
}

// 从客户端获取主题内容
func (client *Client) getRequestTheme(r *http.Request) *data.Theme {
	// 获取主题名称
	name := r.FormValue(vars.CookieKeyTheme)
	if len(name) == 0 {
		cookie, err := r.Cookie(vars.CookieKeyTheme)
		if err != nil && err != http.ErrNoCookie { // 有记录错误，但不退出
			logs.Error(err)
		}

		if cookie != nil {
			name = cookie.Value
		}
	}

	// 查询对应名称的主题
	for _, t := range client.data.Themes {
		if name == t.ID {
			return t
		}
	}

	return client.data.Themes[0] // 不存在的情况下，返回默认主题
}

// 输出一个特定状态码下的错误页面。
// 若该页面模板不存在，则输出状态码对应的文本内容。
// 只查找当前主题目录下的相关文件。
// 只对状态码大于等于 400 的起作用。
func (client *Client) renderError(w http.ResponseWriter, r *http.Request, code int) {
	if code < 400 {
		return
	}
	logs.Debug("输出非正常状态码：", code)

	// 根据情况输出内容，若不存在模板，则直接输出最简单的状态码对应的文本。
	theme := client.getRequestTheme(r)
	filename := strconv.Itoa(code) + vars.TemplateExtension
	path := filepath.Join(client.path.ThemesDir, theme.ID, filename)
	if !utils.FileExists(path) {
		logs.Debugf("模板文件 %s 不存在\n", path)
		helper.StatusError(w, code)
		return
	}

	data, err := ioutil.ReadFile(path)
	if err != nil {
		logs.Errorf("读取模板文件 %s 时出现以下错误: %v\n", path, err)
		helper.StatusError(w, code)
		return
	}

	setContentType(w, client.data.Type)
	w.WriteHeader(code)
	w.Write(data)
}
<commit_msg>去掉不再使用的变量 client.page.template 和 client.page.Theme.Template 值同，使用 client.page.Theme.Template 代码， 去掉 client.page.template 的相关代码。<commit_after>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package client

import (
	"io/ioutil"
	"net/http"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	"github.com/caixw/gitype/data"
	"github.com/caixw/gitype/helper"
	"github.com/caixw/gitype/vars"
	"github.com/issue9/logs"
	"github.com/issue9/utils"
)

const contentTypeKey = "Content-Type"

// 生成一个带编码的 content-type 报头内容
func buildContentTypeContent(mime string) string {
	return mime + ";charset=utf-8"
}

// 设置页面的编码，若已经存在，则不会受影响。
// 要强制指定，请直接使用 w.Header().Set()
func setContentType(w http.ResponseWriter, mime string) {
	h := w.Header()
	if len(h.Get(contentTypeKey)) == 0 {
		h.Set(contentTypeKey, buildContentTypeContent(mime))
	}
}

// 用于描述一个页面的所有无素
type page struct {
	client   *Client
	Info     *info
	response http.ResponseWriter
	request  *http.Request

	Title       string       // 文章标题，可以为空
	Subtitle    string       // 副标题
	Canonical   string       // 当前页的唯一链接
	Keywords    string       // meta.keywords 的值
	Description string       // meta.description 的值
	PrevPage    *data.Link   // 前一页
	NextPage    *data.Link   // 下一页
	Type        string       // 当前页面类型
	Author      *data.Author // 作者
	License     *data.Link   // 当前页的版本信息，可以为空
	Theme       *data.Theme  // 当前页面所使用的主题

	// 以下内容，仅在对应的页面才会有内容
	Q        string          // 搜索关键字
	Tag      *data.Tag       // 标签详细页面，非标签详细页，则为空
	Posts    []*data.Post    // 文章列表，仅标签详情页和搜索页用到。
	Post     *data.Post      // 文章详细内容，仅文章页面用到。
	Archives []*data.Archive // 归档
}

// 页面的附加信息，除非重新加载数据，否则内容不会变。
type info struct {
	AppName    string // 程序名称
	AppURL     string // 程序官网
	AppVersion string // 当前程序的版本号
	GoVersion  string // 编译的 Go 版本号

	SiteName    string     // 网站名称
	URL         string     // 网站地址，若是一个子目录，则需要包含该子目录
	Icon        *data.Icon // 网站图标
	Language    string     // 页面语言
	PostSize    int        // 总文章数量
	Beian       string     // 备案号
	Uptime      time.Time  // 上线时间
	LastUpdated time.Time  // 最后更新时间
	RSS         *data.Link // RSS，NOTICE:指针方便模板判断其值是否为空
	Atom        *data.Link
	Opensearch  *data.Link
	Tags        []*data.Tag  // 标签列表
	Series      []*data.Tag  // 专题列表
	Links       []*data.Link // 友情链接
	Menus       []*data.Link // 导航菜单
}

func (client *Client) newInfo() *info {
	d := client.data

	info := &info{
		AppName:    vars.Name,
		AppURL:     vars.URL,
		AppVersion: vars.Version(),
		GoVersion:  runtime.Version(),

		SiteName:    d.Title,
		URL:         d.URL,
		Icon:        d.Icon,
		Language:    d.Language,
		PostSize:    len(d.Posts),
		Beian:       d.Beian,
		Uptime:      d.Uptime,
		LastUpdated: d.Created,
		Tags:        d.Tags,
		Series:      d.Series,
		Links:       d.Links,
		Menus:       d.Menus,
	}

	if d.RSS != nil {
		info.RSS = &data.Link{
			Title: d.RSS.Title,
			URL:   d.RSS.URL,
			Type:  d.RSS.Type,
		}
	}

	if d.Atom != nil {
		info.Atom = &data.Link{
			Title: d.Atom.Title,
			URL:   d.Atom.URL,
			Type:  d.Atom.Type,
		}
	}

	if d.Opensearch != nil {
		info.Opensearch = &data.Link{
			Title: d.Opensearch.Title,
			URL:   d.Opensearch.URL,
			Type:  d.Opensearch.Type,
		}
	}

	return info
}

func (client *Client) page(typ string, w http.ResponseWriter, r *http.Request) *page {
	theme := client.getRequestTheme(r)
	d := client.data

	return &page{
		client:   client,
		Info:     client.info,
		response: w,
		request:  r,

		Subtitle:    d.Subtitle,
		Keywords:    d.Keywords,
		Description: d.Description,
		Type:        typ,
		Author:      d.Author,
		License:     d.License,
		Theme:       theme,
	}
}

func (p *page) nextPage(url, text string) {
	if len(text) == 0 {
		text = vars.NextPageText
	}

	p.NextPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "next",
	}
}

func (p *page) prevPage(url, text string) {
	if len(text) == 0 {
		text = vars.PrevPageText
	}

	p.PrevPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "prev",
	}
}

// 输出当前内容到指定模板
func (p *page) render(name string) {
	setContentType(p.response, p.client.data.Type)

	cookie := &http.Cookie{
		Name:     vars.CookieKeyTheme,
		Value:    p.Theme.ID,
		HttpOnly: vars.CookieHTTPOnly,
	}
	if p.Theme.ID != p.client.data.Themes[0].ID {
		cookie.MaxAge = vars.CookieMaxAge
	} else {
		cookie.MaxAge = -1
	}
	cookie.Expires = time.Now().Add(time.Second * time.Duration(vars.CookieMaxAge))
	http.SetCookie(p.response, cookie)

	err := p.Theme.Template.ExecuteTemplate(p.response, name, p)
	if err != nil {
		logs.Error(err)
		p.client.renderError(p.response, p.request, http.StatusInternalServerError)
		return
	}
}

// 从客户端获取主题内容
func (client *Client) getRequestTheme(r *http.Request) *data.Theme {
	// 获取主题名称
	name := r.FormValue(vars.CookieKeyTheme)
	if len(name) == 0 {
		cookie, err := r.Cookie(vars.CookieKeyTheme)
		if err != nil && err != http.ErrNoCookie { // 有记录错误，但不退出
			logs.Error(err)
		}

		if cookie != nil {
			name = cookie.Value
		}
	}

	// 查询对应名称的主题
	for _, t := range client.data.Themes {
		if name == t.ID {
			return t
		}
	}

	return client.data.Themes[0] // 不存在的情况下，返回默认主题
}

// 输出一个特定状态码下的错误页面。
// 若该页面模板不存在，则输出状态码对应的文本内容。
// 只查找当前主题目录下的相关文件。
// 只对状态码大于等于 400 的起作用。
func (client *Client) renderError(w http.ResponseWriter, r *http.Request, code int) {
	if code < 400 {
		return
	}
	logs.Debug("输出非正常状态码：", code)

	// 根据情况输出内容，若不存在模板，则直接输出最简单的状态码对应的文本。
	theme := client.getRequestTheme(r)
	filename := strconv.Itoa(code) + vars.TemplateExtension
	path := filepath.Join(client.path.ThemesDir, theme.ID, filename)
	if !utils.FileExists(path) {
		logs.Debugf("模板文件 %s 不存在\n", path)
		helper.StatusError(w, code)
		return
	}

	data, err := ioutil.ReadFile(path)
	if err != nil {
		logs.Errorf("读取模板文件 %s 时出现以下错误: %v\n", path, err)
		helper.StatusError(w, code)
		return
	}

	setContentType(w, client.data.Type)
	w.WriteHeader(code)
	w.Write(data)
}
<|endoftext|>


<commit_before>package libp2pquic

import (
	"crypto/tls"
	"net"

	pstore "github.com/libp2p/go-libp2p-peerstore"
	tpt "github.com/libp2p/go-libp2p-transport"
	quicconn "github.com/marten-seemann/quic-conn"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

type listener struct {
	laddr        ma.Multiaddr
	quicListener net.Listener

	transport tpt.Transport
}

func newListener(laddr ma.Multiaddr, peers pstore.Peerstore, transport tpt.Transport) (*listener, error) {
	tlsConf := &tls.Config{}
	network, host, err := manet.DialArgs(laddr)
	if err != nil {
		return nil, err
	}
	qln, err := quicconn.Listen(network, host, tlsConf)
	if err != nil {
		return nil, err
	}
	return &listener{
		laddr:        laddr,
		quicListener: qln,
		transport:    transport,
	}, nil
}

func (l *listener) Accept() (tpt.Conn, error) {
	c, err := l.quicListener.Accept()
	if err != nil {
		return nil, err
	}

	mnc, err := manet.WrapNetConn(c)
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: mnc,
		Tpt:  l.transport,
	}, nil
}

func (l *listener) Close() error {
	return l.quicListener.Close()
}

func (l *listener) Addr() net.Addr {
	return l.quicListener.Addr()
}

func (l *listener) Multiaddr() ma.Multiaddr {
	return l.laddr
}

var _ tpt.Listener = &listener{}
<commit_msg>initialize the listener with a valid TLS config<commit_after>package libp2pquic

import (
	"net"

	pstore "github.com/libp2p/go-libp2p-peerstore"
	tpt "github.com/libp2p/go-libp2p-transport"
	testdata "github.com/lucas-clemente/quic-go/testdata"
	quicconn "github.com/marten-seemann/quic-conn"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

type listener struct {
	laddr        ma.Multiaddr
	quicListener net.Listener

	transport tpt.Transport
}

func newListener(laddr ma.Multiaddr, peers pstore.Peerstore, transport tpt.Transport) (*listener, error) {
	// we need to provide a certificate here
	// use the demo certificate from quic-go
	tlsConf := testdata.GetTLSConfig()
	network, host, err := manet.DialArgs(laddr)
	if err != nil {
		return nil, err
	}
	qln, err := quicconn.Listen(network, host, tlsConf)
	if err != nil {
		return nil, err
	}
	return &listener{
		laddr:        laddr,
		quicListener: qln,
		transport:    transport,
	}, nil
}

func (l *listener) Accept() (tpt.Conn, error) {
	c, err := l.quicListener.Accept()
	if err != nil {
		return nil, err
	}

	mnc, err := manet.WrapNetConn(c)
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: mnc,
		Tpt:  l.transport,
	}, nil
}

func (l *listener) Close() error {
	return l.quicListener.Close()
}

func (l *listener) Addr() net.Addr {
	return l.quicListener.Addr()
}

func (l *listener) Multiaddr() ma.Multiaddr {
	return l.laddr
}

var _ tpt.Listener = &listener{}
<|endoftext|>


<commit_before>package device

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/ScriptRock/crypto/ssh"
)

type CiscoDevice struct {
	Username  string
	Password  string
	Enable    string
	name      string
	Hostname  string
	stdin     io.WriteCloser
	stdout    io.Reader
	session   *ssh.Session
	Echo      bool
	EnableLog bool
	Logdir    string
	Log       *os.File
	Prompt    string
}

func (d *CiscoDevice) Connect() error {
	config := &ssh.ClientConfig{
		User: d.Username,
		Auth: []ssh.AuthMethod{
			ssh.Password(d.Password),
		},
		Config: ssh.Config{
			Ciphers: ssh.AllSupportedCiphers(),
		},
	}
	client, err := ssh.Dial("tcp", d.Hostname+":22", config)
	if err != nil {
		return err
	}
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	d.stdin, _ = session.StdinPipe()
	d.stdout, _ = session.StdoutPipe()
	d.Echo = true
	d.EnableLog = true
	modes := ssh.TerminalModes{
		ssh.ECHO:          0, // disable echoing
		ssh.OCRNL:         0,
		ssh.TTY_OP_ISPEED: 38400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 38400, // output speed = 14.4kbaud
	}
	session.RequestPty("vt100", 0, 2000, modes)
	session.Shell()
	if d.Logdir != "" {
		t := time.Now()
		d.Log, err = os.OpenFile(filepath.Join(d.Logdir, t.Format("200601021504")+"-"+d.Hostname), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			return err
		}
	}
	d.init()
	d.session = session
	return nil
}

func (d *CiscoDevice) Close() {
	d.session.Close()
}

func (d *CiscoDevice) Cmd(cmd string) (string, error) {
	bufstdout := bufio.NewReader(d.stdout)
	lines := strings.Split(cmd, "!")
	for _, line := range lines {
		io.WriteString(d.stdin, line+"\n")
		time.Sleep(time.Millisecond * 100)
	}
	output, err := d.readln(bufstdout)
	if err != nil {
		return "", err
	}
	output = strings.Replace(output, "\r", "", -1)
	if d.Echo == false {
		output = strings.Replace(output, lines[0], "", 1)
	}
	if d.Logdir != "" {
		return "", nil
	}
	return output, nil
}

func (d *CiscoDevice) init() {
	bufstdout := bufio.NewReader(d.stdout)
	io.WriteString(d.stdin, "enable\n")
	time.Sleep(time.Millisecond * 100)
	re := regexp.MustCompile("assword:")
	buf := make([]byte, 1000)
	loadStr := ""
	for {
		n, err := bufstdout.Read(buf)
		if err != nil {
			log.Fatal(err)
		}
		loadStr += string(buf[:n])
		if re.MatchString(loadStr) {
			io.WriteString(d.stdin, d.Enable+"\n")
			break
		} else {
			break
		}
	}
	d.Cmd("terminal length 0")
	d.Cmd("")
	prompt, _ := d.Cmd("")
	d.Prompt = strings.TrimSpace(prompt)
	d.Prompt = strings.Replace(d.Prompt, "#", "", -1)
}

func (d *CiscoDevice) readln(r *bufio.Reader) (string, error) {
	var re *regexp.Regexp
	if d.Prompt == "" {
		re = regexp.MustCompile("[[:alnum:]]#.?$")
	} else {
		re = regexp.MustCompile(d.Prompt + ".*?#$")
	}
	buf := make([]byte, 10000)
	loadStr := ""
	for {
		n, err := r.Read(buf)
		if err != nil {
			return "", err
		}
		loadStr += string(buf[:n])
		// logging to file if necessary
		if d.Logdir != "" {
			if d.EnableLog {
				fmt.Fprint(d.Log, string(buf[:n]))
			}
		}
		if re.MatchString(string(buf[:n])) {
			break
		}
	}
	return loadStr, nil
}
<commit_msg>Work around cisco config problem<commit_after>package device

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/ScriptRock/crypto/ssh"
)

type CiscoDevice struct {
	Username  string
	Password  string
	Enable    string
	name      string
	Hostname  string
	stdin     io.WriteCloser
	stdout    io.Reader
	session   *ssh.Session
	Echo      bool
	EnableLog bool
	Logdir    string
	Log       *os.File
	Prompt    string
}

func (d *CiscoDevice) Connect() error {
	config := &ssh.ClientConfig{
		User: d.Username,
		Auth: []ssh.AuthMethod{
			ssh.Password(d.Password),
		},
		Config: ssh.Config{
			Ciphers: ssh.AllSupportedCiphers(),
		},
	}
	client, err := ssh.Dial("tcp", d.Hostname+":22", config)
	if err != nil {
		return err
	}
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	d.stdin, _ = session.StdinPipe()
	d.stdout, _ = session.StdoutPipe()
	d.Echo = true
	d.EnableLog = true
	modes := ssh.TerminalModes{
		ssh.ECHO:          0, // disable echoing
		ssh.OCRNL:         0,
		ssh.TTY_OP_ISPEED: 38400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 38400, // output speed = 14.4kbaud
	}
	session.RequestPty("vt100", 0, 2000, modes)
	session.Shell()
	if d.Logdir != "" {
		t := time.Now()
		d.Log, err = os.OpenFile(filepath.Join(d.Logdir, t.Format("200601021504")+"-"+d.Hostname), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			return err
		}
	}
	d.init()
	d.session = session
	return nil
}

func (d *CiscoDevice) Close() {
	d.session.Close()
}

func (d *CiscoDevice) Cmd(cmd string) (string, error) {
	bufstdout := bufio.NewReader(d.stdout)
	lines := strings.Split(cmd, "!")
	for _, line := range lines {
		io.WriteString(d.stdin, line+"\n")
		time.Sleep(time.Millisecond * 100)
	}
	output, err := d.readln(bufstdout)
	if err != nil {
		return "", err
	}
	output = strings.Replace(output, "\r", "", -1)
	if d.Echo == false {
		output = strings.Replace(output, lines[0], "", 1)
	}
	if d.Logdir != "" {
		return "", nil
	}
	return output, nil
}

func (d *CiscoDevice) init() {
	bufstdout := bufio.NewReader(d.stdout)
	io.WriteString(d.stdin, "enable\n")
	time.Sleep(time.Millisecond * 100)
	re := regexp.MustCompile("assword:")
	buf := make([]byte, 1000)
	loadStr := ""
	for {
		n, err := bufstdout.Read(buf)
		if err != nil {
			log.Fatal(err)
		}
		loadStr += string(buf[:n])
		if re.MatchString(loadStr) {
			io.WriteString(d.stdin, d.Enable+"\n")
			break
		} else {
			break
		}
	}
	d.Cmd("terminal length 0")
	d.Cmd("")
	prompt, _ := d.Cmd("")
	d.Prompt = strings.TrimSpace(prompt)
	d.Prompt = strings.Replace(d.Prompt, "#", "", -1)
	// sometimes using conf t makes the (config-xx-something) so long that only 10 chars of
	// original prompt remain
	if len(d.Prompt) > 10 {
		d.Prompt = d.Prompt[:10]
	}
}

func (d *CiscoDevice) readln(r *bufio.Reader) (string, error) {
	var re *regexp.Regexp
	if d.Prompt == "" {
		re = regexp.MustCompile("[[:alnum:]]#.?$")
	} else {
		re = regexp.MustCompile(d.Prompt + ".*?#.?$")
	}
	buf := make([]byte, 10000)
	loadStr := ""
	for {
		n, err := r.Read(buf)
		if err != nil {
			return "", err
		}
		loadStr += string(buf[:n])
		// logging to file if necessary
		if d.Logdir != "" {
			if d.EnableLog {
				fmt.Fprint(d.Log, string(buf[:n]))
			}
		}
		if re.MatchString(string(buf[:n])) {
			break
		}
	}
	return loadStr, nil
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"

	"github.com/hashicorp/nomad/nomad/structs"
)

type allocTuple struct {
	exist, updated *structs.Allocation
}

// diffResult is used to return the sets that result from a diff
type diffResult struct {
	added   []*structs.Allocation
	removed []*structs.Allocation
	updated []allocTuple
	ignore  []*structs.Allocation
}

func (d *diffResult) GoString() string {
	return fmt.Sprintf("allocs: (added %d) (removed %d) (updated %d) (ignore %d)",
		len(d.added), len(d.removed), len(d.updated), len(d.ignore))
}

// diffAllocs is used to diff the existing and updated allocations
// to see what has happened.
func diffAllocs(existing []*structs.Allocation, allocs *allocUpdates) *diffResult {
	// Scan the existing allocations
	result := &diffResult{}
	existIdx := make(map[string]struct{})
	for _, exist := range existing {
		// Mark this as existing
		existIdx[exist.ID] = struct{}{}

		// Check if the alloc was updated or filtered because an update wasn't
		// needed.
		alloc, pulled := allocs.pulled[exist.ID]
		_, filtered := allocs.filtered[exist.ID]

		// If not updated or filtered, removed
		if !pulled && !filtered {
			result.removed = append(result.removed, exist)
			continue
		}

		// Check for an update
		if pulled && alloc.AllocModifyIndex > exist.AllocModifyIndex {
			result.updated = append(result.updated, allocTuple{exist, alloc})
			continue
		}

		// Ignore this
		result.ignore = append(result.ignore, exist)
	}

	// Scan the updated allocations for any that are new
	for id, pulled := range allocs.pulled {
		if _, ok := existIdx[id]; !ok {
			result.added = append(result.added, pulled)
		}
	}
	return result
}

// shuffleStrings randomly shuffles the list of strings
func shuffleStrings(list []string) {
	for i := range list {
		j := rand.Intn(i + 1)
		list[i], list[j] = list[j], list[i]
	}
}

// persistState is used to help with saving state
func persistState(path string, data interface{}) error {
	buf, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to encode state: %v", err)
	}
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return fmt.Errorf("failed to make dirs for %s: %v", path, err)
	}
	if err := ioutil.WriteFile(path, buf, 0600); err != nil {
		return fmt.Errorf("failed to save state: %v", err)
	}
	return nil
}

// restoreState is used to read back in the persisted state
func restoreState(path string, data interface{}) error {
	buf, err := ioutil.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return fmt.Errorf("failed to read state: %v", err)
	}
	if err := json.Unmarshal(buf, data); err != nil {
		return fmt.Errorf("failed to decode state: %v", err)
	}
	return nil
}
<commit_msg>write state to temp file and then rename<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"

	"github.com/hashicorp/nomad/nomad/structs"
)

type allocTuple struct {
	exist, updated *structs.Allocation
}

// diffResult is used to return the sets that result from a diff
type diffResult struct {
	added   []*structs.Allocation
	removed []*structs.Allocation
	updated []allocTuple
	ignore  []*structs.Allocation
}

func (d *diffResult) GoString() string {
	return fmt.Sprintf("allocs: (added %d) (removed %d) (updated %d) (ignore %d)",
		len(d.added), len(d.removed), len(d.updated), len(d.ignore))
}

// diffAllocs is used to diff the existing and updated allocations
// to see what has happened.
func diffAllocs(existing []*structs.Allocation, allocs *allocUpdates) *diffResult {
	// Scan the existing allocations
	result := &diffResult{}
	existIdx := make(map[string]struct{})
	for _, exist := range existing {
		// Mark this as existing
		existIdx[exist.ID] = struct{}{}

		// Check if the alloc was updated or filtered because an update wasn't
		// needed.
		alloc, pulled := allocs.pulled[exist.ID]
		_, filtered := allocs.filtered[exist.ID]

		// If not updated or filtered, removed
		if !pulled && !filtered {
			result.removed = append(result.removed, exist)
			continue
		}

		// Check for an update
		if pulled && alloc.AllocModifyIndex > exist.AllocModifyIndex {
			result.updated = append(result.updated, allocTuple{exist, alloc})
			continue
		}

		// Ignore this
		result.ignore = append(result.ignore, exist)
	}

	// Scan the updated allocations for any that are new
	for id, pulled := range allocs.pulled {
		if _, ok := existIdx[id]; !ok {
			result.added = append(result.added, pulled)
		}
	}
	return result
}

// shuffleStrings randomly shuffles the list of strings
func shuffleStrings(list []string) {
	for i := range list {
		j := rand.Intn(i + 1)
		list[i], list[j] = list[j], list[i]
	}
}

// persistState is used to help with saving state
func persistState(path string, data interface{}) error {
	buf, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to encode state: %v", err)
	}
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return fmt.Errorf("failed to make dirs for %s: %v", path, err)
	}
	tmpPath := path + ".tmp"
	if err := ioutil.WriteFile(tmpPath, buf, 0600); err != nil {
		return fmt.Errorf("failed to save state to tmp: %v", err)
	}
	if err := os.Rename(tmpPath, path); err != nil {
		return fmt.Errorf("failed to rename tmp to path: %v", err)
	}
	return nil
}

// restoreState is used to read back in the persisted state
func restoreState(path string, data interface{}) error {
	buf, err := ioutil.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return fmt.Errorf("failed to read state: %v", err)
	}
	if err := json.Unmarshal(buf, data); err != nil {
		return fmt.Errorf("failed to decode state: %v", err)
	}
	return nil
}
<|endoftext|>


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

package driver

import (
	"context"
	"path"
	"time"

	"github.com/arangodb/go-driver/util"
)

// NewClient creates a new Client based on the given config setting.
func NewClient(config ClientConfig) (Client, error) {
	if config.Connection == nil {
		return nil, WithStack(InvalidArgumentError{Message: "Connection is not set"})
	}
	conn := config.Connection
	if config.Authentication != nil {
		var err error
		conn, err = conn.SetAuthentication(config.Authentication)
		if err != nil {
			return nil, WithStack(err)
		}
	}
	c := &client{
		conn: conn,
	}
	if config.SynchronizeEndpointsInterval > 0 {
		go c.autoSynchronizeEndpoints(config.SynchronizeEndpointsInterval)
	}
	return c, nil
}

// client implements the Client interface.
type client struct {
	conn Connection
}

// Connection returns the connection used by this client
func (c *client) Connection() Connection {
	return c.conn
}

// SynchronizeEndpoints fetches all endpoints from an ArangoDB cluster and updates the
// connection to use those endpoints.
// When this client is connected to a single server, nothing happens.
// When this client is connected to a cluster of servers, the connection will be updated to reflect
// the layout of the cluster.
func (c *client) SynchronizeEndpoints(ctx context.Context) error {
	return c.SynchronizeEndpoints2(ctx, "")
}

// SynchronizeEndpoints2 fetches all endpoints from an ArangoDB cluster and updates the
// connection to use those endpoints.
// When this client is connected to a single server, nothing happens.
// When this client is connected to a cluster of servers, the connection will be updated to reflect
// the layout of the cluster.
// Compared to SynchronizeEndpoints, this function expects a database name as additional parameter.
// This database name is used to call `_db/<dbname>/_api/cluster/endpoints`. SynchronizeEndpoints uses
// the default database, i.e. `_system`. In the case the user does not have access to `_system`,
// SynchronizeEndpoints does not work with earlier versions of arangodb.
func (c *client) SynchronizeEndpoints2(ctx context.Context, dbname string) error {
	// Cluster mode, fetch endpoints
	cep, err := c.clusterEndpoints(ctx, dbname)
	if err != nil {
		// ignore Forbidden: automatic failover is not enabled errors
		if !IsArangoErrorWithErrorNum(err, 403, 0, 11) { // 3.2 returns no error code, thus check for 0
			return WithStack(err)
		}

		return nil
	}
	var endpoints []string
	for _, ep := range cep.Endpoints {
		endpoints = append(endpoints, util.FixupEndpointURLScheme(ep.Endpoint))
	}

	// Update connection
	if err := c.conn.UpdateEndpoints(endpoints); err != nil {
		return WithStack(err)
	}

	return nil
}

// autoSynchronizeEndpoints performs automatic endpoint synchronization.
func (c *client) autoSynchronizeEndpoints(interval time.Duration) {
	for {
		// SynchronizeEndpoints endpoints
		c.SynchronizeEndpoints(nil)

		// Wait a bit
		time.Sleep(interval)
	}
}

type clusterEndpointsResponse struct {
	Endpoints []clusterEndpoint `json:"endpoints,omitempty"`
}

type clusterEndpoint struct {
	Endpoint string `json:"endpoint,omitempty"`
}

// clusterEndpoints returns the endpoints of a cluster.
func (c *client) clusterEndpoints(ctx context.Context, dbname string) (clusterEndpointsResponse, error) {
	var url string
	if dbname == "" {
		url = "_api/cluster/endpoints"
	} else {
		url = path.Join("_db", pathEscape(dbname), "_api/cluster/endpoints")
	}
	req, err := c.conn.NewRequest("GET", url)
	if err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	applyContextSettings(ctx, req)
	resp, err := c.conn.Do(ctx, req)
	if err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	if err := resp.CheckStatus(200); err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	var data clusterEndpointsResponse
	if err := resp.ParseBody("", &data); err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	return data, nil
}
<commit_msg>Fix tests for current devel.<commit_after>//
// DISCLAIMER
//
// Copyright 2017 ArangoDB GmbH, Cologne, Germany
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Copyright holder is ArangoDB GmbH, Cologne, Germany
//
// Author Ewout Prangsma
//

package driver

import (
	"context"
	"path"
	"time"

	"github.com/arangodb/go-driver/util"
)

// NewClient creates a new Client based on the given config setting.
func NewClient(config ClientConfig) (Client, error) {
	if config.Connection == nil {
		return nil, WithStack(InvalidArgumentError{Message: "Connection is not set"})
	}
	conn := config.Connection
	if config.Authentication != nil {
		var err error
		conn, err = conn.SetAuthentication(config.Authentication)
		if err != nil {
			return nil, WithStack(err)
		}
	}
	c := &client{
		conn: conn,
	}
	if config.SynchronizeEndpointsInterval > 0 {
		go c.autoSynchronizeEndpoints(config.SynchronizeEndpointsInterval)
	}
	return c, nil
}

// client implements the Client interface.
type client struct {
	conn Connection
}

// Connection returns the connection used by this client
func (c *client) Connection() Connection {
	return c.conn
}

// SynchronizeEndpoints fetches all endpoints from an ArangoDB cluster and updates the
// connection to use those endpoints.
// When this client is connected to a single server, nothing happens.
// When this client is connected to a cluster of servers, the connection will be updated to reflect
// the layout of the cluster.
func (c *client) SynchronizeEndpoints(ctx context.Context) error {
	return c.SynchronizeEndpoints2(ctx, "")
}

// SynchronizeEndpoints2 fetches all endpoints from an ArangoDB cluster and updates the
// connection to use those endpoints.
// When this client is connected to a single server, nothing happens.
// When this client is connected to a cluster of servers, the connection will be updated to reflect
// the layout of the cluster.
// Compared to SynchronizeEndpoints, this function expects a database name as additional parameter.
// This database name is used to call `_db/<dbname>/_api/cluster/endpoints`. SynchronizeEndpoints uses
// the default database, i.e. `_system`. In the case the user does not have access to `_system`,
// SynchronizeEndpoints does not work with earlier versions of arangodb.
func (c *client) SynchronizeEndpoints2(ctx context.Context, dbname string) error {
	// Cluster mode, fetch endpoints
	cep, err := c.clusterEndpoints(ctx, dbname)
	if err != nil {
		// ignore Forbidden: automatic failover is not enabled errors
		if !IsArangoErrorWithErrorNum(err, 403, 501, 0, 11) { // 3.2 returns no error code, thus check for 0
			// 501 is in there since 3.7 for some time returned this in a single server
			return WithStack(err)
		}

		return nil
	}
	var endpoints []string
	for _, ep := range cep.Endpoints {
		endpoints = append(endpoints, util.FixupEndpointURLScheme(ep.Endpoint))
	}

	// Update connection
	if err := c.conn.UpdateEndpoints(endpoints); err != nil {
		return WithStack(err)
	}

	return nil
}

// autoSynchronizeEndpoints performs automatic endpoint synchronization.
func (c *client) autoSynchronizeEndpoints(interval time.Duration) {
	for {
		// SynchronizeEndpoints endpoints
		c.SynchronizeEndpoints(nil)

		// Wait a bit
		time.Sleep(interval)
	}
}

type clusterEndpointsResponse struct {
	Endpoints []clusterEndpoint `json:"endpoints,omitempty"`
}

type clusterEndpoint struct {
	Endpoint string `json:"endpoint,omitempty"`
}

// clusterEndpoints returns the endpoints of a cluster.
func (c *client) clusterEndpoints(ctx context.Context, dbname string) (clusterEndpointsResponse, error) {
	var url string
	if dbname == "" {
		url = "_api/cluster/endpoints"
	} else {
		url = path.Join("_db", pathEscape(dbname), "_api/cluster/endpoints")
	}
	req, err := c.conn.NewRequest("GET", url)
	if err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	applyContextSettings(ctx, req)
	resp, err := c.conn.Do(ctx, req)
	if err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	if err := resp.CheckStatus(200); err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	var data clusterEndpointsResponse
	if err := resp.ParseBody("", &data); err != nil {
		return clusterEndpointsResponse{}, WithStack(err)
	}
	return data, nil
}
<|endoftext|>


<commit_before>package gandalf

import (
	"bytes"
	"encoding/json"
	. "launchpad.net/gocheck"
	"net/http/httptest"
)

func (s *S) TestDoRequest(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	body := bytes.NewBufferString(`{"foo":"bar"}`)
	response, err := client.doRequest("POST", "/test", body)
	c.Assert(err, IsNil)
	c.Assert(response.StatusCode, Equals, 200)
	c.Assert(string(h.body), Equals, `{"foo":"bar"}`)
	c.Assert(h.url, Equals, "/test")
}

func (s *S) TestDoRequestShouldNotSetContentTypeToJsonWhenBodyIsNil(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	response, err := client.doRequest("DELETE", "/test", nil)
	c.Assert(err, IsNil)
	c.Assert(response.StatusCode, Equals, 200)
	c.Assert(h.header.Get("Content-Type"), Not(Equals), "application/json")
}

func (s *S) TestDoRequestConnectionError(c *C) {
	client := Client{Endpoint: "http://127.0.0.1:747399"}
	response, err := client.doRequest("GET", "/", nil)
	c.Assert(response, IsNil)
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "Failed to connect to Gandalf server, it's probably down.")
}

func (s *S) TestPost(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	r := repository{Name: "test", Users: []string{"samwan"}}
	err := client.post(r, "/repository")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository")
	c.Assert(h.method, Equals, "POST")
	c.Assert(string(h.body), Equals, `{"name":"test","users":["samwan"],"ispublic":false}`)
}

func (s *S) TestPostWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	r := repository{Name: "test", Users: []string{"samwan"}}
	err := client.post(r, "/repository")
	c.Assert(err, ErrorMatches, "^Error performing requested operation\n$")
}

func (s *S) TestDelete(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.delete(nil, "/user/someuser")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestDeleteWithConnectionError(c *C) {
	client := Client{Endpoint: "http://127.0.0.1:747399"}
	err := client.delete(nil, "/users/something")
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "Failed to connect to Gandalf server, it's probably down.")
}

func (s *S) TestDeleteWithResponseError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.delete(nil, "/user/someuser")
	c.Assert(err, ErrorMatches, "^Error performing requested operation\n$")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestDeleteWithBody(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.delete(map[string]string{"test": "foo"}, "/user/someuser")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, `{"test":"foo"}`)
}

func (s *S) TestGet(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.get("/user/someuser")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "GET")
}

func (s *S) TestGetWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.get("/user/someuser")
	c.Assert(err, ErrorMatches, "^Error performing requested operation\n$")
}

func (s *S) TestFormatBody(c *C) {
	b, err := (&Client{}).formatBody(map[string]string{"test": "foo"})
	c.Assert(err, IsNil)
	c.Assert(b.String(), Equals, `{"test":"foo"}`)
}

func (s *S) TestFormatBodyReturnJsonNullWithNilBody(c *C) {
	b, err := (&Client{}).formatBody(nil)
	c.Assert(err, IsNil)
	c.Assert(b.String(), Equals, "null")
}

func (s *S) TestNewRepository(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewRepository("proj1", []string{"someuser"}, false)
	c.Assert(err, IsNil)
	c.Assert(string(h.body), Equals, `{"name":"proj1","users":["someuser"],"ispublic":false}`)
	c.Assert(h.url, Equals, "/repository")
	c.Assert(h.method, Equals, "POST")
}

func (s *S) TestNewRepositoryWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewRepository("proj1", []string{"someuser"}, false)
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestNewUser(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewUser("someuser", map[string]string{"testkey": "ssh-rsa somekey"})
	c.Assert(err, IsNil)
	c.Assert(string(h.body), Equals, `{"name":"someuser","keys":{"testkey":"ssh-rsa somekey"}}`)
	c.Assert(h.url, Equals, "/user")
	c.Assert(h.method, Equals, "POST")
}

func (s *S) TestNewUserWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewUser("someuser", map[string]string{"testkey": "ssh-rsa somekey"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRemoveUser(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveUser("someuser")
	c.Assert(err, IsNil)
	c.Assert(string(h.body), Equals, "null")
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "DELETE")
}

func (s *S) TestRemoveUserWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveUser("someuser")
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRemoveRepository(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveRepository("project1")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository/project1")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestRemoveRepositoryWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveRepository("proj2")
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestAddKey(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	key := map[string]string{"pubkey": "ssh-rsa somekey me@myhost"}
	err := client.AddKey("username", key)
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/username/key")
	c.Assert(h.method, Equals, "POST")
	c.Assert(string(h.body), Equals, `{"pubkey":"ssh-rsa somekey me@myhost"}`)
	c.Assert(h.header.Get("Content-Type"), Equals, "application/json")
}

func (s *S) TestAddKeyWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.AddKey("proj2", map[string]string{"key": "ssh-rsa keycontent user@host"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRemoveKey(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveKey("username", "keyname")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/username/key/keyname")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestRemoveKeyWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveKey("proj2", "keyname")
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestGrantAccess(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	repositories := []string{"projectx", "projecty"}
	users := []string{"userx"}
	err := client.GrantAccess(repositories, users)
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository/grant")
	c.Assert(h.method, Equals, "POST")
	expected, err := json.Marshal(map[string][]string{"repositories": repositories, "users": users})
	c.Assert(err, IsNil)
	c.Assert(h.body, DeepEquals, expected)
	c.Assert(h.header.Get("Content-Type"), Equals, "application/json")
}

func (s *S) TestGrantAccessWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.GrantAccess([]string{"projectx", "projecty"}, []string{"userx"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRevokeAccess(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	repositories := []string{"projectx", "projecty"}
	users := []string{"userx"}
	err := client.RevokeAccess(repositories, users)
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository/revoke")
	c.Assert(h.method, Equals, "DELETE")
	expected, err := json.Marshal(map[string][]string{"repositories": repositories, "users": users})
	c.Assert(err, IsNil)
	c.Assert(h.body, DeepEquals, expected)
	c.Assert(h.header.Get("Content-Type"), Equals, "application/json")
}

func (s *S) TestRevokeAccessWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RevokeAccess([]string{"projectx", "projecty"}, []string{"usery"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}
<commit_msg>tests: added test for connection failure during POST<commit_after>package gandalf

import (
	"bytes"
	"encoding/json"
	. "launchpad.net/gocheck"
	"net/http/httptest"
)

func (s *S) TestDoRequest(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	body := bytes.NewBufferString(`{"foo":"bar"}`)
	response, err := client.doRequest("POST", "/test", body)
	c.Assert(err, IsNil)
	c.Assert(response.StatusCode, Equals, 200)
	c.Assert(string(h.body), Equals, `{"foo":"bar"}`)
	c.Assert(h.url, Equals, "/test")
}

func (s *S) TestDoRequestShouldNotSetContentTypeToJsonWhenBodyIsNil(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	response, err := client.doRequest("DELETE", "/test", nil)
	c.Assert(err, IsNil)
	c.Assert(response.StatusCode, Equals, 200)
	c.Assert(h.header.Get("Content-Type"), Not(Equals), "application/json")
}

func (s *S) TestDoRequestConnectionError(c *C) {
	client := Client{Endpoint: "http://127.0.0.1:747399"}
	response, err := client.doRequest("GET", "/", nil)
	c.Assert(response, IsNil)
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "Failed to connect to Gandalf server, it's probably down.")
}

func (s *S) TestPost(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	r := repository{Name: "test", Users: []string{"samwan"}}
	err := client.post(r, "/repository")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository")
	c.Assert(h.method, Equals, "POST")
	c.Assert(string(h.body), Equals, `{"name":"test","users":["samwan"],"ispublic":false}`)
}

func (s *S) TestPostWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	r := repository{Name: "test", Users: []string{"samwan"}}
	err := client.post(r, "/repository")
	c.Assert(err, ErrorMatches, "^Error performing requested operation\n$")
}

func (s *S) TestPostConnectionFailure(c *C) {
	client := Client{Endpoint: "http://127.0.0.1:747399"}
	err := client.post(nil, "/")
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "Failed to connect to Gandalf server, it's probably down.")
}

func (s *S) TestDelete(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.delete(nil, "/user/someuser")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestDeleteWithConnectionError(c *C) {
	client := Client{Endpoint: "http://127.0.0.1:747399"}
	err := client.delete(nil, "/users/something")
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "Failed to connect to Gandalf server, it's probably down.")
}

func (s *S) TestDeleteWithResponseError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.delete(nil, "/user/someuser")
	c.Assert(err, ErrorMatches, "^Error performing requested operation\n$")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestDeleteWithBody(c *C) {
	h := TestHandler{content: `some return message`}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.delete(map[string]string{"test": "foo"}, "/user/someuser")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, `{"test":"foo"}`)
}

func (s *S) TestGet(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.get("/user/someuser")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "GET")
}

func (s *S) TestGetWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.get("/user/someuser")
	c.Assert(err, ErrorMatches, "^Error performing requested operation\n$")
}

func (s *S) TestFormatBody(c *C) {
	b, err := (&Client{}).formatBody(map[string]string{"test": "foo"})
	c.Assert(err, IsNil)
	c.Assert(b.String(), Equals, `{"test":"foo"}`)
}

func (s *S) TestFormatBodyReturnJsonNullWithNilBody(c *C) {
	b, err := (&Client{}).formatBody(nil)
	c.Assert(err, IsNil)
	c.Assert(b.String(), Equals, "null")
}

func (s *S) TestNewRepository(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewRepository("proj1", []string{"someuser"}, false)
	c.Assert(err, IsNil)
	c.Assert(string(h.body), Equals, `{"name":"proj1","users":["someuser"],"ispublic":false}`)
	c.Assert(h.url, Equals, "/repository")
	c.Assert(h.method, Equals, "POST")
}

func (s *S) TestNewRepositoryWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewRepository("proj1", []string{"someuser"}, false)
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestNewUser(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewUser("someuser", map[string]string{"testkey": "ssh-rsa somekey"})
	c.Assert(err, IsNil)
	c.Assert(string(h.body), Equals, `{"name":"someuser","keys":{"testkey":"ssh-rsa somekey"}}`)
	c.Assert(h.url, Equals, "/user")
	c.Assert(h.method, Equals, "POST")
}

func (s *S) TestNewUserWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	_, err := client.NewUser("someuser", map[string]string{"testkey": "ssh-rsa somekey"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRemoveUser(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveUser("someuser")
	c.Assert(err, IsNil)
	c.Assert(string(h.body), Equals, "null")
	c.Assert(h.url, Equals, "/user/someuser")
	c.Assert(h.method, Equals, "DELETE")
}

func (s *S) TestRemoveUserWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveUser("someuser")
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRemoveRepository(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveRepository("project1")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository/project1")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestRemoveRepositoryWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveRepository("proj2")
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestAddKey(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	key := map[string]string{"pubkey": "ssh-rsa somekey me@myhost"}
	err := client.AddKey("username", key)
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/username/key")
	c.Assert(h.method, Equals, "POST")
	c.Assert(string(h.body), Equals, `{"pubkey":"ssh-rsa somekey me@myhost"}`)
	c.Assert(h.header.Get("Content-Type"), Equals, "application/json")
}

func (s *S) TestAddKeyWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.AddKey("proj2", map[string]string{"key": "ssh-rsa keycontent user@host"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRemoveKey(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveKey("username", "keyname")
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/user/username/key/keyname")
	c.Assert(h.method, Equals, "DELETE")
	c.Assert(string(h.body), Equals, "null")
}

func (s *S) TestRemoveKeyWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RemoveKey("proj2", "keyname")
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestGrantAccess(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	repositories := []string{"projectx", "projecty"}
	users := []string{"userx"}
	err := client.GrantAccess(repositories, users)
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository/grant")
	c.Assert(h.method, Equals, "POST")
	expected, err := json.Marshal(map[string][]string{"repositories": repositories, "users": users})
	c.Assert(err, IsNil)
	c.Assert(h.body, DeepEquals, expected)
	c.Assert(h.header.Get("Content-Type"), Equals, "application/json")
}

func (s *S) TestGrantAccessWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.GrantAccess([]string{"projectx", "projecty"}, []string{"userx"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}

func (s *S) TestRevokeAccess(c *C) {
	h := TestHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	repositories := []string{"projectx", "projecty"}
	users := []string{"userx"}
	err := client.RevokeAccess(repositories, users)
	c.Assert(err, IsNil)
	c.Assert(h.url, Equals, "/repository/revoke")
	c.Assert(h.method, Equals, "DELETE")
	expected, err := json.Marshal(map[string][]string{"repositories": repositories, "users": users})
	c.Assert(err, IsNil)
	c.Assert(h.body, DeepEquals, expected)
	c.Assert(h.header.Get("Content-Type"), Equals, "application/json")
}

func (s *S) TestRevokeAccessWithError(c *C) {
	h := ErrorHandler{}
	ts := httptest.NewServer(&h)
	defer ts.Close()
	client := Client{Endpoint: ts.URL}
	err := client.RevokeAccess([]string{"projectx", "projecty"}, []string{"usery"})
	expected := "^Error performing requested operation\n$"
	c.Assert(err, ErrorMatches, expected)
}
<|endoftext|>


<commit_before>package models

import (
	"errors"
	"time"

	"github.com/koding/bongo"
)

type ChannelMessageList struct {
	// unique identifier of the channel message list
	Id int64 `json:"id,string"`

	// Id of the channel
	ChannelId int64 `json:"channelId,string"     sql:"NOT NULL"`

	// Id of the message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

	// Addition date of the message to the channel
	AddedAt time.Time `json:"addedAt"            sql:"NOT NULL"`
}

func (c *ChannelMessageList) BeforeCreate() {
	c.AddedAt = time.Now()
}

func (c *ChannelMessageList) BeforeUpdate() {
	c.AddedAt = time.Now()
}

func (c *ChannelMessageList) AfterCreate() {
	bongo.B.AfterCreate(c)
}

func (c *ChannelMessageList) AfterUpdate() {
	bongo.B.AfterUpdate(c)
}

func (c ChannelMessageList) AfterDelete() {
	bongo.B.AfterDelete(c)
}

func (c ChannelMessageList) GetId() int64 {
	return c.Id
}

func (c ChannelMessageList) TableName() string {
	return "api.channel_message_list"
}

func NewChannelMessageList() *ChannelMessageList {
	return &ChannelMessageList{}
}

func (c *ChannelMessageList) ById(id int64) error {
	return bongo.B.ById(c, id)
}

func (c *ChannelMessageList) One(q *bongo.Query) error {
	return bongo.B.One(c, c, q)
}

func (c *ChannelMessageList) Update() error {
	return bongo.B.Update(c)
}

func (c *ChannelMessageList) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(c, data, q)
}

func (c *ChannelMessageList) UnreadCount(cp *ChannelParticipant) (int, error) {
	if cp.ChannelId == 0 {
		return 0, errors.New("ChannelId is not set")
	}

	if cp.AccountId == 0 {
		return 0, errors.New("AccountId is not set")
	}

	if cp.LastSeenAt.IsZero() {
		return 0, errors.New("Last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(c,
		"channel_id = ? and added_at > ?",
		cp.ChannelId,
		// todo change this format to get from a specific place
		cp.LastSeenAt.UTC().Format(time.RFC3339),
	)
}

func (c *ChannelMessageList) Create() error {
	return bongo.B.Create(c)
}

func (c *ChannelMessageList) Delete() error {
	return bongo.B.Delete(c)
}

func (c *ChannelMessageList) List(q *Query, populateUnreadCount bool) (*HistoryResponse, error) {
	messageList, err := c.getMessages(q)
	if err != nil {
		return nil, err
	}

	if populateUnreadCount {
		messageList = c.populateUnreadCount(messageList)
	}

	hr := NewHistoryResponse()
	hr.MessageList = messageList
	return hr, nil
}

// populateUnreadCount adds unread count into message containers
func (c *ChannelMessageList) populateUnreadCount(messageList []*ChannelMessageContainer) []*ChannelMessageContainer {
	channel := NewChannel()
	channel.Id = c.ChannelId

	for i, message := range messageList {
		cml, err := channel.FetchMessageList(message.Message.Id)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}

		count, err := NewMessageReply().UnreadCount(cml)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}
		messageList[i].UnreadRepliesCount = count
	}

	return messageList
}

func (c *ChannelMessageList) getMessages(q *Query) ([]*ChannelMessageContainer, error) {
	var messages []int64

	if c.ChannelId == 0 {
		return nil, errors.New("ChannelId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": c.ChannelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort:       map[string]string{"added_at": "DESC"},
	}

	bongoQuery := bongo.B.BuildQuery(c, query)
	if !q.From.IsZero() {
		bongoQuery = bongoQuery.Where("added_at < ?", q.From)
	}

	bongoQuery = bongoQuery.Pluck(query.Pluck, &messages)
	if err := bongoQuery.Error; err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessages, err := parent.FetchByIds(messages)
	if err != nil {
		return nil, err
	}

	populatedChannelMessages, err := c.populateChannelMessages(channelMessages, q)
	if err != nil {
		return nil, err
	}

	return populatedChannelMessages, nil
}

func (c *ChannelMessageList) populateChannelMessages(channelMessages []ChannelMessage, query *Query) ([]*ChannelMessageContainer, error) {
	channelMessageCount := len(channelMessages)

	populatedChannelMessages := make([]*ChannelMessageContainer, channelMessageCount)

	if channelMessageCount == 0 {
		return populatedChannelMessages, nil
	}

	for i := 0; i < channelMessageCount; i++ {
		cm := channelMessages[i]
		cmc, err := cm.BuildMessage(query)
		if err != nil {
			return nil, err
		}

		populatedChannelMessages[i] = cmc
	}
	return populatedChannelMessages, nil

}

func (c *ChannelMessageList) FetchMessageChannelIds(messageId int64) ([]int64, error) {
	var channelIds []int64

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": messageId,
		},
		Pluck: "channel_id",
	}

	err := bongo.B.Some(c, &channelIds, q)
	if err != nil {
		return nil, err
	}

	return channelIds, nil
}

func (c *ChannelMessageList) FetchMessageChannels(messageId int64) ([]Channel, error) {
	channelIds, err := c.FetchMessageChannelIds(messageId)
	if err != nil {
		return nil, err
	}

	return NewChannel().FetchByIds(channelIds)
}

func (c *ChannelMessageList) FetchMessageIdsByChannelId(channelId int64, q *Query) ([]int64, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"added_at": "DESC",
		},
	}

	var messageIds []int64
	if err := c.Some(&messageIds, query); err != nil {
		return nil, err
	}

	if messageIds == nil {
		return make([]int64, 0), nil
	}

	return messageIds, nil
}

// seperate this fucntion into modelhelper
// as setting it to a variadic function
func (c *ChannelMessageList) DeleteMessagesBySelector(selector map[string]interface{}) error {
	var cmls []ChannelMessageList

	err := bongo.B.Some(c, &cmls, &bongo.Query{Selector: selector})
	if err != nil {
		return err
	}

	for _, cml := range cmls {
		if err := cml.Delete(); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Social: fetch unread count by message id and addedAt of channel message list<commit_after>package models

import (
	"errors"
	"time"

	"github.com/koding/bongo"
)

type ChannelMessageList struct {
	// unique identifier of the channel message list
	Id int64 `json:"id,string"`

	// Id of the channel
	ChannelId int64 `json:"channelId,string"     sql:"NOT NULL"`

	// Id of the message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

	// Addition date of the message to the channel
	AddedAt time.Time `json:"addedAt"            sql:"NOT NULL"`
}

func (c *ChannelMessageList) BeforeCreate() {
	c.AddedAt = time.Now()
}

func (c *ChannelMessageList) BeforeUpdate() {
	c.AddedAt = time.Now()
}

func (c *ChannelMessageList) AfterCreate() {
	bongo.B.AfterCreate(c)
}

func (c *ChannelMessageList) AfterUpdate() {
	bongo.B.AfterUpdate(c)
}

func (c ChannelMessageList) AfterDelete() {
	bongo.B.AfterDelete(c)
}

func (c ChannelMessageList) GetId() int64 {
	return c.Id
}

func (c ChannelMessageList) TableName() string {
	return "api.channel_message_list"
}

func NewChannelMessageList() *ChannelMessageList {
	return &ChannelMessageList{}
}

func (c *ChannelMessageList) ById(id int64) error {
	return bongo.B.ById(c, id)
}

func (c *ChannelMessageList) One(q *bongo.Query) error {
	return bongo.B.One(c, c, q)
}

func (c *ChannelMessageList) Update() error {
	return bongo.B.Update(c)
}

func (c *ChannelMessageList) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(c, data, q)
}

func (c *ChannelMessageList) UnreadCount(cp *ChannelParticipant) (int, error) {
	if cp.ChannelId == 0 {
		return 0, errors.New("ChannelId is not set")
	}

	if cp.AccountId == 0 {
		return 0, errors.New("AccountId is not set")
	}

	if cp.LastSeenAt.IsZero() {
		return 0, errors.New("Last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(c,
		"channel_id = ? and added_at > ?",
		cp.ChannelId,
		// todo change this format to get from a specific place
		cp.LastSeenAt.UTC().Format(time.RFC3339),
	)
}

func (c *ChannelMessageList) Create() error {
	return bongo.B.Create(c)
}

func (c *ChannelMessageList) Delete() error {
	return bongo.B.Delete(c)
}

func (c *ChannelMessageList) List(q *Query, populateUnreadCount bool) (*HistoryResponse, error) {
	messageList, err := c.getMessages(q)
	if err != nil {
		return nil, err
	}

	if populateUnreadCount {
		messageList = c.populateUnreadCount(messageList)
	}

	hr := NewHistoryResponse()
	hr.MessageList = messageList
	return hr, nil
}

// populateUnreadCount adds unread count into message containers
func (c *ChannelMessageList) populateUnreadCount(messageList []*ChannelMessageContainer) []*ChannelMessageContainer {
	channel := NewChannel()
	channel.Id = c.ChannelId

	for i, message := range messageList {
		cml, err := channel.FetchMessageList(message.Message.Id)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}

		count, err := NewMessageReply().UnreadCount(cml.MessageId, cml.AddedAt)
		if err != nil {
			// helper.MustGetLogger().Error(err.Error())
			continue
		}
		messageList[i].UnreadRepliesCount = count
	}

	return messageList
}

func (c *ChannelMessageList) getMessages(q *Query) ([]*ChannelMessageContainer, error) {
	var messages []int64

	if c.ChannelId == 0 {
		return nil, errors.New("ChannelId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": c.ChannelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort:       map[string]string{"added_at": "DESC"},
	}

	bongoQuery := bongo.B.BuildQuery(c, query)
	if !q.From.IsZero() {
		bongoQuery = bongoQuery.Where("added_at < ?", q.From)
	}

	bongoQuery = bongoQuery.Pluck(query.Pluck, &messages)
	if err := bongoQuery.Error; err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessages, err := parent.FetchByIds(messages)
	if err != nil {
		return nil, err
	}

	populatedChannelMessages, err := c.populateChannelMessages(channelMessages, q)
	if err != nil {
		return nil, err
	}

	return populatedChannelMessages, nil
}

func (c *ChannelMessageList) populateChannelMessages(channelMessages []ChannelMessage, query *Query) ([]*ChannelMessageContainer, error) {
	channelMessageCount := len(channelMessages)

	populatedChannelMessages := make([]*ChannelMessageContainer, channelMessageCount)

	if channelMessageCount == 0 {
		return populatedChannelMessages, nil
	}

	for i := 0; i < channelMessageCount; i++ {
		cm := channelMessages[i]
		cmc, err := cm.BuildMessage(query)
		if err != nil {
			return nil, err
		}

		populatedChannelMessages[i] = cmc
	}
	return populatedChannelMessages, nil

}

func (c *ChannelMessageList) FetchMessageChannelIds(messageId int64) ([]int64, error) {
	var channelIds []int64

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": messageId,
		},
		Pluck: "channel_id",
	}

	err := bongo.B.Some(c, &channelIds, q)
	if err != nil {
		return nil, err
	}

	return channelIds, nil
}

func (c *ChannelMessageList) FetchMessageChannels(messageId int64) ([]Channel, error) {
	channelIds, err := c.FetchMessageChannelIds(messageId)
	if err != nil {
		return nil, err
	}

	return NewChannel().FetchByIds(channelIds)
}

func (c *ChannelMessageList) FetchMessageIdsByChannelId(channelId int64, q *Query) ([]int64, error) {
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id": channelId,
		},
		Pluck:      "message_id",
		Pagination: *bongo.NewPagination(q.Limit, q.Skip),
		Sort: map[string]string{
			"added_at": "DESC",
		},
	}

	var messageIds []int64
	if err := c.Some(&messageIds, query); err != nil {
		return nil, err
	}

	if messageIds == nil {
		return make([]int64, 0), nil
	}

	return messageIds, nil
}

// seperate this fucntion into modelhelper
// as setting it to a variadic function
func (c *ChannelMessageList) DeleteMessagesBySelector(selector map[string]interface{}) error {
	var cmls []ChannelMessageList

	err := bongo.B.Some(c, &cmls, &bongo.Query{Selector: selector})
	if err != nil {
		return err
	}

	for _, cml := range cmls {
		if err := cml.Delete(); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// client_test.go - Katzenpost client library tests.
// Copyright (C) 2019  David Stainton.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program 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 Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

// Package client provides a Katzenpost client library.
package client

import (
	"testing"
	"time"

	"github.com/katzenpost/kimchi"
	"github.com/stretchr/testify/require"
)

const basePort = 30000

// TestClientConnect tests that a client can connect and send a message to the loop service
func TestClientConnect(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	voting := false
	nVoting := 0
	nProvider := 2
	nMix := 6
	k := kimchi.NewKimchi(basePort+400, "", nil, voting, nVoting, nProvider, nMix)
	t.Logf("Running TestClientConnect.")
	k.Run()

	go func() {
		defer k.Shutdown()

		// create a client configuration
		cfg, username, linkKey, err := k.GetClientConfig()
		require.NoError(err)
		require.NotNil(cfg)

		<-time.After(90 * time.Second) // must wait for provider to fetch pki document
		t.Logf("Time is up!")

		// instantiate a client instance
		c, err := New(cfg)
		require.NotNil(cfg)
		require.NoError(err)

		// add client log output
		go k.LogTailer(username, cfg.Logging.File)

		// instantiate a session
		s, err := c.NewSession(linkKey)
		require.NoError(err)

		// get a PKI document? needs client method...
		desc, err := s.GetService("loop") // XXX: returns nil and no error?!
		require.NoError(err)

		// send a message
		t.Logf("desc.Provider: %s", desc.Provider)
		_, err = s.SendUnreliableMessage(desc.Name, desc.Provider, []byte("hello!"))
		t.Logf("Sent unreliable message to loop service")
		require.NoError(err)

		c.Shutdown()
		t.Logf("Shutdown requested")
		c.Wait()
	}()

	k.Wait()
	t.Logf("Terminated.")
}
<commit_msg>Add TestAutoRegisterRandomClient<commit_after>// client_test.go - Katzenpost client library tests.
// Copyright (C) 2019  David Stainton.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program 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 Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

// Package client provides a Katzenpost client library.
package client

import (
	"testing"
	"time"

	"github.com/katzenpost/kimchi"
	"github.com/stretchr/testify/require"
)

const basePort = 30000

// TestClientConnect tests that a client can connect and send a message to the loop service
func TestClientConnect(t *testing.T) {
	require := require.New(t)
	voting := false
	nVoting := 0
	nProvider := 2
	nMix := 6
	k := kimchi.NewKimchi(basePort+400, "", nil, voting, nVoting, nProvider, nMix)
	t.Logf("Running TestClientConnect.")
	k.Run()

	go func() {
		defer k.Shutdown()

		// create a client configuration
		cfg, username, linkKey, err := k.GetClientConfig()
		require.NoError(err)
		require.NotNil(cfg)

		<-time.After(90 * time.Second) // must wait for provider to fetch pki document
		t.Logf("Time is up!")

		// instantiate a client instance
		c, err := New(cfg)
		require.NotNil(cfg)
		require.NoError(err)

		// add client log output
		go k.LogTailer(username, cfg.Logging.File)

		// instantiate a session
		s, err := c.NewSession(linkKey)
		require.NoError(err)

		// look up a well known service
		desc, err := s.GetService("loop")
		require.NoError(err)

		// send a message
		t.Logf("desc.Provider: %s", desc.Provider)
		_, err = s.SendUnreliableMessage(desc.Name, desc.Provider, []byte("hello!"))
		t.Logf("Sent unreliable message to loop service")
		require.NoError(err)

		c.Shutdown()
		t.Logf("Shutdown requested")
		c.Wait()
	}()

	k.Wait()
	t.Logf("Terminated.")
}

// TestAutoRegisterRandomClient tests client registration
func TestAutoRegisterRandomClient(t *testing.T) {
	require := require.New(t)
	voting := false
	nVoting := 0
	nProvider := 2
	nMix := 6
	k := kimchi.NewKimchi(basePort+500, "", nil, voting, nVoting, nProvider, nMix)
	t.Logf("Running TestAutoRegisterRandomClient.")
	k.Run()

	go func() {
		defer k.Shutdown()
		<-time.After(70 * time.Second) // must wait for provider to fetch pki document

		cfg, err := k.GetClientNetconfig()
		require.NoError(err)

		_, linkKey := AutoRegisterRandomClient(cfg)
		require.NotNil(linkKey)

		// Verify that the client can connect
		c, err := New(cfg)
		require.NoError(err)

		// instantiate a session
		s, err := c.NewSession(linkKey)
		require.NoError(err)

		// look up a well known service
		desc, err := s.GetService("loop")
		require.NoError(err)
		t.Logf("Found %v kaetzchen on %v", desc.Name, desc.Provider)

		c.Shutdown()
		t.Logf("Shutdown requested")
		c.Wait()
	}()
	k.Wait()
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package crossdock

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestHandler(t *testing.T) {
	tests := []struct {
		failOnUnknown bool
		status        string
	}{
		{false, "skipped"},
		{true, "failed"},
	}

	for _, test := range tests {
		testHandler(t, test.failOnUnknown, test.status)
	}
}

func testHandler(t *testing.T, failOnUnknown bool, status string) {
	behaviors := Behaviors{
		"b1": func(t T) {
			t.Successf("ok")
		},
	}

	server := httptest.NewServer(Handler(behaviors, failOnUnknown))
	defer server.Close()

	verifyBehavior(t, server.URL, "b1", map[string]string{
		"status": "passed",
		"output": "ok",
	})

	verifyBehavior(t, server.URL, "b2", map[string]string{
		"status": status,
		"output": "unknown behavior \"b2\"",
	})
}

func verifyBehavior(t *testing.T, url string, behavior string, expectation map[string]string) {
	res, err := http.Get(fmt.Sprintf("%s/?behavior=%s", url, behavior))
	require.NoError(t, err)

	defer res.Body.Close()

	var answer []map[string]string
	err = json.NewDecoder(res.Body).Decode(&answer)
	require.NoError(t, err)

	assert.Equal(t, []map[string]string{expectation}, answer)
}
<commit_msg>nits from #7 (#8)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package crossdock

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestHandler(t *testing.T) {
	tests := []struct {
		failOnUnknown bool
		status        string
	}{
		{false, "skipped"},
		{true, "failed"},
	}

	for _, test := range tests {
		testHandler(t, test.failOnUnknown, test.status)
	}
}

func testHandler(t *testing.T, failOnUnknown bool, status string) {
	behaviors := Behaviors{
		"b1": func(t T) {
			t.Successf("ok")
		},
	}

	server := httptest.NewServer(Handler(behaviors, failOnUnknown))
	defer server.Close()

	runTestCase(t, server.URL, "b1", map[string]string{
		"status": "passed",
		"output": "ok",
	})

	runTestCase(t, server.URL, "b2", map[string]string{
		"status": status,
		"output": `unknown behavior "b2"`,
	})
}

func runTestCase(t *testing.T, url string, behavior string, expectation map[string]string) {
	res, err := http.Get(fmt.Sprintf("%s/?behavior=%s", url, behavior))
	require.NoError(t, err)

	defer res.Body.Close()

	var answer []map[string]string
	err = json.NewDecoder(res.Body).Decode(&answer)
	require.NoError(t, err)

	assert.Equal(t, []map[string]string{expectation}, answer)
}
<|endoftext|>


<commit_before>package garvis

import (
	"bytes"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
	glog "google.golang.org/appengine/log"

	"github.com/go-telegram-bot-api/telegram-bot-api"
)

type TextFilter struct {
	bot  *tgbotapi.BotAPI
	ctx  context.Context
	done chan bool
}

type TextFilterRule struct {
	ChatID    int64
	RxText    string
	TextReply string
	Count     int
	Limit     int
	UserID    int
}

type User struct {
	UserID   int
	UserName string
}

func (filter TextFilter) Run(update tgbotapi.Update) error {
	ctx := filter.ctx

	// It is not possible to match usernames to user ids from the API
	// for the case of mentions so we need to store username-ids in the database
	// for lookup
	updateUsers(ctx, update)

	query := datastore.NewQuery("Rule").Filter("ChatID = ", update.Message.Chat.ID)
	queryUserRules := query.Filter("UserID = ", update.Message.From.ID)
	queryStaticRules := query.Filter("UserID = ", 0)

	var userRules []TextFilterRule
	var staticRules []TextFilterRule
	userKeys, err := queryUserRules.GetAll(ctx, &userRules)
	staticKeys, err := queryStaticRules.GetAll(ctx, &staticRules)
	if err != nil {
		glog.Errorf(ctx, "client.GetAll: %v", err)
	}
	rules := append(userRules, staticRules...)
	keys := append(userKeys, staticKeys...)

	for i, rule := range rules {
		k := keys[i]

		rxRule := regexp.MustCompile(rule.RxText)

		if rxRule.MatchString(update.Message.Text) {
			rule.Count = rule.Count + 1
			if rule.Count >= rule.Limit {
				rule.Count = 0
				_, err = datastore.Put(ctx, k, &rule)
				if err != nil {
					glog.Errorf(ctx, "client.Put(reset): %v", err)
				}
				filter.Trigger(rule)
			}
			_, err = datastore.Put(ctx, k, &rule)
			if err != nil {
				glog.Errorf(ctx, "client.Put: %v", err)
			}
		}
	}

	filter.done <- true
	return nil
}

func (filter TextFilter) GetCommands(commands map[string]Filter) {
	commands["addrule"] = filter
	commands["deleterule"] = filter
	commands["listrules"] = filter
}

func (filter TextFilter) RunCommand(cmd string, cmdarg CommandArguments) {
	update := cmdarg.update
	var err error
	switch cmd {
	case "addrule":
		err = filter.addRule(update)
	case "deleterule":
		err = filter.deleteRule(update)
	case "listrules":
		err = filter.listRules(update)
	}
	if err != nil {
		glog.Errorf(filter.ctx, err.Error())
	}
}

func (filter TextFilter) Trigger(rule TextFilterRule) {
	msg := tgbotapi.NewMessage(rule.ChatID, rule.TextReply)
	filter.bot.Send(msg)
}

func (filter TextFilter) addRule(update tgbotapi.Update) (err error) {
	ctx := filter.ctx

	command := strings.SplitN(update.Message.Text, " ", 2)
	if len(command) < 2 {
		usage := "Usage: /addrule {regex matcher}~{reply}{#count (optional default: 1)}~{user (optional)}"
		msg := tgbotapi.NewMessage(update.Message.Chat.ID, usage)
		filter.bot.Send(msg)
		return nil
	}
	argstr := command[1]
	args := strings.SplitN(argstr, "~", 3)
	var userID int
	switch len(args) {
	case 2:
		userID = 0
	case 3:
		ents := update.Message.Entities
		for _, ent := range *ents {
			switch ent.Type {
			case "text_mention":
				userID = ent.User.ID
			case "mention":
				if userID, err = getUserID(ctx, args[2][1:]); err != nil {
					return err
				}
			}
		}
	default:
		return nil
	}
	arg1 := strings.SplitN(args[0], "#", 2)
	text := fmt.Sprintf("(?i)%v(?-i)", arg1[0])
	var limit int
	if len(arg1) < 2 {
		limit = 1
	} else {
		limit, _ = strconv.Atoi(arg1[1])
	}
	textreply := args[1]

	keyl, _, _ := datastore.AllocateIDs(ctx, "Rule", nil, 1)
	ruleKey := datastore.NewKey(ctx, "Rule", "", keyl, nil)
	rule := TextFilterRule{
		ChatID:    update.Message.Chat.ID,
		RxText:    text,
		TextReply: textreply,
		Count:     0,
		Limit:     int(limit),
		UserID:    userID,
	}

	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		var empty TextFilterRule
		if err := datastore.Get(ctx, ruleKey, &empty); err != datastore.ErrNoSuchEntity {
			return err
		}
		_, err := datastore.Put(ctx, ruleKey, &rule)
		return err
	}, nil)

	return err
}

func (filter TextFilter) deleteRule(update tgbotapi.Update) (err error) {
	ctx := filter.ctx

	command := strings.SplitN(update.Message.Text, " ", 2)
	if len(command) < 2 {
		usage := "Usage: /deleterule {ruleID}"
		msg := tgbotapi.NewMessage(update.Message.Chat.ID, usage)
		filter.bot.Send(msg)
		return nil
	}
	key, _ := strconv.ParseInt(command[1], 10, 64)
	ruleKey := datastore.NewKey(ctx, "Rule", "", key, nil)

	if err = datastore.Delete(filter.ctx, ruleKey); err != datastore.ErrNoSuchEntity {
		return err
	}

	return err
}

func (filter TextFilter) listRules(update tgbotapi.Update) error {
	ctx := filter.ctx

	query := datastore.NewQuery("Rule").Filter("ChatID = ", update.Message.Chat.ID)

	var buffer bytes.Buffer

	header := fmt.Sprintf("|%s|%s|%s|%s|%s|\n", "ID", "Regex", "Reply", "Count", "User(0 for all)")
	buffer.WriteString(header)
	buffer.WriteString(strings.Repeat("-", 32) + "\n")

	for t := query.Run(ctx); ; {
		var rule TextFilterRule
		k, err := t.Next(&rule)
		if err == datastore.Done {
			break
		}

		ruleText := fmt.Sprintf("|%v|%v|%v|%v|%v|\n", k.IntID(), rule.RxText, rule.TextReply, rule.Limit, rule.UserID)

		buffer.WriteString(ruleText)
	}

	msg := tgbotapi.NewMessage(update.Message.Chat.ID, buffer.String())
	filter.bot.Send(msg)

	return nil
}

func updateUsers(ctx context.Context, update tgbotapi.Update) (err error) {
	userKey := datastore.NewKey(ctx, "User", update.Message.From.UserName, 0, nil)
	user := User{
		UserID:   update.Message.From.ID,
		UserName: update.Message.From.UserName,
	}
	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		var empty User
		if err := datastore.Get(ctx, userKey, &empty); err != datastore.ErrNoSuchEntity {
			return err
		}
		_, err := datastore.Put(ctx, userKey, &user)
		return err
	}, nil)

	return err
}

func getUserID(ctx context.Context, mention string) (int, error) {
	var user User
	key := datastore.NewKey(ctx, "User", mention, 0, nil)
	err := datastore.Get(ctx, key, &user)
	if err != nil {
		glog.Errorf(ctx, "Error fetching user: %v", err)
		return -1, err
	}

	return user.UserID, nil
}
<commit_msg>Add addhiddenrule command in TextFilter<commit_after>package garvis

import (
	"bytes"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
	glog "google.golang.org/appengine/log"

	"github.com/go-telegram-bot-api/telegram-bot-api"
)

type TextFilter struct {
	bot  *tgbotapi.BotAPI
	ctx  context.Context
	done chan bool
}

type TextFilterRule struct {
	ChatID    int64
	RxText    string
	TextReply string
	Count     int
	Limit     int
	UserID    int
	Hidden    bool
	CreatorID int
}

type User struct {
	UserID   int
	UserName string
}

func (filter TextFilter) Run(update tgbotapi.Update) error {
	ctx := filter.ctx

	// It is not possible to match usernames to user ids from the API
	// for the case of mentions so we need to store username-ids in the database
	// for lookup
	updateUsers(ctx, update)

	query := datastore.NewQuery("Rule").Filter("ChatID = ", update.Message.Chat.ID)
	queryUserRules := query.Filter("UserID = ", update.Message.From.ID)
	queryStaticRules := query.Filter("UserID = ", 0)

	var userRules []TextFilterRule
	var staticRules []TextFilterRule
	userKeys, err := queryUserRules.GetAll(ctx, &userRules)
	staticKeys, err := queryStaticRules.GetAll(ctx, &staticRules)
	if err != nil {
		glog.Errorf(ctx, "client.GetAll: %v", err)
	}
	rules := append(userRules, staticRules...)
	keys := append(userKeys, staticKeys...)

	for i, rule := range rules {
		k := keys[i]

		rxRule := regexp.MustCompile(rule.RxText)

		if rxRule.MatchString(update.Message.Text) {
			rule.Count = rule.Count + 1
			if rule.Count >= rule.Limit {
				rule.Count = 0
				_, err = datastore.Put(ctx, k, &rule)
				if err != nil {
					glog.Errorf(ctx, "client.Put(reset): %v", err)
				}
				filter.Trigger(rule)
			}
			_, err = datastore.Put(ctx, k, &rule)
			if err != nil {
				glog.Errorf(ctx, "client.Put: %v", err)
			}
		}
	}

	filter.done <- true
	return nil
}

func (filter TextFilter) GetCommands(commands map[string]Filter) {
	commands["addrule"] = filter
	commands["deleterule"] = filter
	commands["listrules"] = filter
	commands["addhiddenrule"] = filter
}

func (filter TextFilter) RunCommand(cmd string, cmdarg CommandArguments) {
	update := cmdarg.update
	var err error
	switch cmd {
	case "addrule":
		err = filter.addRule(update, false)
	case "deleterule":
		err = filter.deleteRule(update)
	case "listrules":
		err = filter.listRules(update)
	case "addhiddenrule":
		err = filter.addRule(update, true)
	}
	if err != nil {
		glog.Errorf(filter.ctx, err.Error())
	}
}

func (filter TextFilter) Trigger(rule TextFilterRule) {
	msg := tgbotapi.NewMessage(rule.ChatID, rule.TextReply)
	filter.bot.Send(msg)
}

func (filter TextFilter) addRule(update tgbotapi.Update, hidden bool) (err error) {
	ctx := filter.ctx

	command := strings.SplitN(update.Message.Text, " ", 2)
	if len(command) < 2 {
		usage := "Usage: /addrule {regex matcher}~{reply}{#count (optional default: 1)}~{user (optional)}"
		msg := tgbotapi.NewMessage(update.Message.Chat.ID, usage)
		filter.bot.Send(msg)
		return nil
	}
	argstr := command[1]
	args := strings.SplitN(argstr, "~", 3)
	var userID int
	switch len(args) {
	case 2:
		userID = 0
	case 3:
		ents := update.Message.Entities
		for _, ent := range *ents {
			switch ent.Type {
			case "text_mention":
				userID = ent.User.ID
			case "mention":
				if userID, err = getUserID(ctx, args[2][1:]); err != nil {
					return err
				}
			}
		}
	default:
		return nil
	}
	arg1 := strings.SplitN(args[0], "#", 2)
	text := fmt.Sprintf("(?i)%v(?-i)", arg1[0])
	_, err = regexp.Compile(text)
	if err != nil {
		usage := "Invalid Regex"
		msg := tgbotapi.NewMessage(update.Message.Chat.ID, usage)
		filter.bot.Send(msg)
		return nil
	}
	var limit int
	if len(arg1) < 2 {
		limit = 1
	} else {
		limit, _ = strconv.Atoi(arg1[1])
	}
	textreply := args[1]

	keyl, _, _ := datastore.AllocateIDs(ctx, "Rule", nil, 1)
	ruleKey := datastore.NewKey(ctx, "Rule", "", keyl, nil)
	var creatorID int
	if creatorID, err = getUserID(ctx, update.Message.From.UserName); err != nil {
		return err
	}
	rule := TextFilterRule{
		ChatID:    update.Message.Chat.ID,
		RxText:    text,
		TextReply: textreply,
		Count:     0,
		Limit:     int(limit),
		UserID:    userID,
		CreatorID: creatorID,
		Hidden:    hidden,
	}

	_, err = datastore.Put(ctx, ruleKey, &rule)

	return err
}

func (filter TextFilter) deleteRule(update tgbotapi.Update) (err error) {
	ctx := filter.ctx

	command := strings.SplitN(update.Message.Text, " ", 2)
	if len(command) < 2 {
		usage := "Usage: /deleterule {ruleID}"
		msg := tgbotapi.NewMessage(update.Message.Chat.ID, usage)
		filter.bot.Send(msg)
		return nil
	}
	key, _ := strconv.ParseInt(command[1], 10, 64)
	ruleKey := datastore.NewKey(ctx, "Rule", "", key, nil)

	var userID int
	if userID, err = getUserID(ctx, update.Message.From.UserName); err != nil {
		return err
	}

	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		var rule TextFilterRule
		err := datastore.Get(ctx, ruleKey, &rule)
		if err != nil {
			return err
		}

		if rule.ChatID == update.Message.Chat.ID || rule.CreatorID == userID {
			err = datastore.Delete(ctx, ruleKey)
		}

		return err
	}, nil)

	return err
}

func (filter TextFilter) listRules(update tgbotapi.Update) (err error) {
	ctx := filter.ctx

	var buffer bytes.Buffer
	var query *datastore.Query
	var header string

	if update.Message.Chat.Type == "private" {
		var userID int
		if userID, err = getUserID(ctx, update.Message.From.UserName); err != nil {
			return err
		}
		query = datastore.NewQuery("Rule").Filter("CreatorID = ", userID)
		header = fmt.Sprintf("|%s|%s|%s|%s|%s|%s|\n", "Chat name", "ID", "Regex", "Reply", "Count", "User(0 for all)")
	} else {
		query = datastore.NewQuery("Rule").Filter("ChatID = ", update.Message.Chat.ID)
		header = fmt.Sprintf("|%s|%s|%s|%s|%s|\n", "ID", "Regex", "Reply", "Count", "User(0 for all)")
	}
	buffer.WriteString(header)

	for t := query.Run(ctx); ; {
		var rule TextFilterRule
		var ruleText string

		k, err := t.Next(&rule)
		if err == datastore.Done {
			break
		}

		if update.Message.Chat.Type == "private" {
			chat, _ := filter.bot.GetChat(tgbotapi.ChatConfig{ChatID: rule.ChatID})
			ruleText = fmt.Sprintf("|%v|%v|%v|%v|%v|%v|\n", chat.Title, k.IntID(), rule.RxText, rule.TextReply, rule.Limit, rule.UserID)
		} else if !rule.Hidden {
			ruleText = fmt.Sprintf("|%v|%v|%v|%v|%v|\n", k.IntID(), rule.RxText, rule.TextReply, rule.Limit, rule.UserID)
		}

		buffer.WriteString(ruleText)
	}

	msg := tgbotapi.NewMessage(update.Message.Chat.ID, buffer.String())
	filter.bot.Send(msg)

	return nil
}

func updateUsers(ctx context.Context, update tgbotapi.Update) (err error) {
	userKey := datastore.NewKey(ctx, "User", update.Message.From.UserName, 0, nil)
	user := User{
		UserID:   update.Message.From.ID,
		UserName: update.Message.From.UserName,
	}
	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		var empty User
		if err := datastore.Get(ctx, userKey, &empty); err != datastore.ErrNoSuchEntity {
			return err
		}
		_, err := datastore.Put(ctx, userKey, &user)
		return err
	}, nil)

	return err
}

func getUserID(ctx context.Context, mention string) (int, error) {
	var user User
	key := datastore.NewKey(ctx, "User", mention, 0, nil)
	err := datastore.Get(ctx, key, &user)
	if err != nil {
		glog.Errorf(ctx, "Error fetching user: %v", err)
		return -1, err
	}

	return user.UserID, nil
}
<|endoftext|>