<commit_before>package group

import (
	"container/list"
	"database/sql"
	"sort"

	"github.com/jmoiron/sqlx"
	"github.com/mijia/sweb/log"

	"github.com/laincloud/sso/ssolib/models"
	"github.com/laincloud/sso/ssolib/models/iuser"
)

type MemberRole int8

const (
	NORMAL MemberRole = iota
	ADMIN
)

type Member struct {
	iuser.User
	Role MemberRole
}

type GroupRole struct {
	Group
	Role MemberRole // some user's role in the group
}

const createUserGroupTableSQL = `
CREATE TABLE IF NOT EXISTS user_group (
	user_id INT NOT NULL,
	group_id INT NOT NULL,
	role TINYINT NOT NULL COMMENT '0: normal, 1: admin',
	PRIMARY KEY (user_id, group_id),
	KEY (group_id)
)`

func (g *Group) AddMember(ctx *models.Context, user iuser.User, role MemberRole) error {
	tx := ctx.DB.MustBegin()
	log.Debug("begin add:", user.GetId(), g.Id, role)
	_, err1 := tx.Exec(
		"INSERT INTO user_group (user_id, group_id, role) VALUES (?, ?, ?)",
		user.GetId(), g.Id, role)

	err2 := tx.Commit()

	if err1 != nil {
		return err1
	}

	if err2 != nil {
		return err2
	}

	return nil
}

func (g *Group) UpdateMember(ctx *models.Context, user iuser.User, role MemberRole) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec(
		"UPDATE user_group SET role=? WHERE user_id=? AND group_id=?",
		role, user.GetId(), g.Id)

	err2 := tx.Commit()

	if err1 != nil {
		return err1
	}

	if err2 != nil {
		return err2
	}

	return nil
}

func (g *Group) RemoveMember(ctx *models.Context, user iuser.User) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec("DELETE FROM user_group WHERE user_id=? AND group_id=?",
		user.GetId(), g.Id)

	if err2 := tx.Commit(); err2 != nil {
		return err2
	}

	if err1 != nil {
		return err1
	}

	return nil
}

// not recursive for nested groups
func (g *Group) ListMembers(ctx *models.Context) ([]Member, error) {
	back := ctx.Back
	members := []Member{}
	if g.GroupType == iuser.SSOLIBGROUP {
		rows, err := ctx.DB.Query(
			"SELECT user_id, role FROM user_group WHERE group_id=?",
			g.Id)
		if err != nil {
			if err == sql.ErrNoRows {
				return members, nil
			}
			return nil, err
		}

		for rows.Next() {
			var userId int
			var role MemberRole
			if err = rows.Scan(&userId, &role); err != nil {
				return nil, err
			}
			user, err := back.GetUser(userId)
			if err != nil {
				return nil, err
			}

			members = append(members, Member{user, role})
		}

		return members, nil
	} else if g.GroupType == iuser.BACKENDGROUP {
		var ubg iuser.BackendWithGroup
		var ok bool
		if ubg, ok = back.(iuser.BackendWithGroup); !ok {
			return nil, ErrBackendUnsupported
		}
		bg, err := ubg.GetBackendGroupByName(g.Name)
		if err != nil {
			panic(err)
		}
		users, err := bg.ListUsers()
		if err != nil {
			if err == iuser.ErrMethodNotSupported {
				return members, nil
			} else {
				panic(err)
			}
		}
		members = backendUsersToMembers(users)
		return members, nil
	} else {
		panic("here")
	}
	return nil, nil
}

// Return (true, role, nil) if u is member of g, otherwise return (false, 0, nil).
// error will be non-nil if anything unexpected happens.
// Must considering recursive if valid
func (g *Group) GetMember(ctx *models.Context, u iuser.User) (ok bool, role MemberRole, err error) {
	/*err = ctx.DB.QueryRow("SELECT role FROM user_group WHERE user_id=? AND group_id=?",
		u.GetId(), g.Id).Scan(&role)
	switch {
	case err == sql.ErrNoRows:
		err = nil
		return
	case err != nil:
		return
	default:
		ok = true
		return
	}*/
	groups, err := getGroupRolesRecursivelyOfUser(ctx, u, false)
	if err != nil {
		return false, 0, err
	}
	if role, ok = groups[g.Id]; ok {
		return true, role, nil
	} else {
		return false, 0, nil
	}
}

// must not recursive
func (g *Group) removeAllMembers(ctx *models.Context) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec("DELETE FROM user_group WHERE group_id=?", g.Id)

	if err2 := tx.Commit(); err2 != nil {
		return err2
	}

	if err1 != nil {
		return err1
	}
	return nil
}

// direct groups of the users, for now it is used only in "管理员特供"
func GetGroupsOfUserByIds(ctx *models.Context, userIds []int) (map[int][]Group, error) {
	groupMap := make(map[int][]Group)
	if len(userIds) == 0 {
		return groupMap, nil
	}

	userGroupIds := make(map[int][]int)
	groupIds := []int{}
	if query, args, err := sqlx.In("SELECT group_id, user_id FROM user_group WHERE user_id IN(?)", userIds); err != nil {
		if err == sql.ErrNoRows {
			err = nil
		}
		return groupMap, err
	} else {
		if rows, err := ctx.DB.Query(query, args...); err != nil {
			return groupMap, err
		} else {
			for rows.Next() {
				var (
					userId  int
					groupId int
				)
				if err := rows.Scan(&groupId, &userId); err != nil {
					return groupMap, err
				}
				groupIds = append(groupIds, groupId)
				userGroupIds[userId] = append(userGroupIds[userId], groupId)
			}
		}
	}

	groups, err := ListGroups(ctx, groupIds...)
	if err != nil {
		return groupMap, err
	}
	groupSet := make(map[int]Group)
	for _, g := range groups {
		groupSet[g.Id] = g
	}
	for userId, groups := range userGroupIds {
		for _, g := range groups {
			if group, ok := groupSet[g]; ok {
				groupMap[userId] = append(groupMap[userId], group)
			}
		}
	}
	return groupMap, nil
}

// should be recursive, but may be slow if we find recursive roles
func GetGroupRolesOfUser(ctx *models.Context, user iuser.User) ([]GroupRole, error) {
	roles := []GroupRole{}
	mapRoles, err := getGroupRolesRecursivelyOfUser(ctx, user, false)
	if err != nil {
		panic(err)
	}
	sortedGroupIds := []int{}
	for k, _ := range mapRoles {
		sortedGroupIds = append(sortedGroupIds, k)
	}
	sort.Ints(sortedGroupIds)
	for _, k := range sortedGroupIds {
		g, err := GetGroup(ctx, k)
		v := mapRoles[k]
		if err != nil {
			panic(err)
		}
		tmp := GroupRole{
			Group: *g,
			Role:  v,
		}
		roles = append(roles, tmp)
	}
	return roles, nil
}

func getGroupRolesDirectlyOfUser(ctx *models.Context, user iuser.User) ([]GroupRole, error) {
	groupRoles, err := getSSOLIBGroupRolesDirectlyOfUser(ctx, user)
	ub := ctx.Back
	var bGroups []iuser.BackendGroup
	if ubg, ok := ub.(iuser.BackendWithGroup); ok {
		bGroups, err = ubg.GetBackendGroupsOfUser(user)
		if err != nil {
			panic(err)
		}
	}
	ret := make([]GroupRole, (len(groupRoles) + len(bGroups)))
	copy(ret, groupRoles)
	for i, g := range bGroups {
		role := NORMAL
		ret[(i + len(groupRoles))] = GroupRole{
			Group: Group{
				Id:       g.GetId(),
				Name:     g.GetName(),
				FullName: g.GetRules().(string),
			},
			Role: role,
		}
	}
	return ret, nil
}

func getSSOLIBGroupRolesDirectlyOfUser(ctx *models.Context, user iuser.User) ([]GroupRole, error) {
	groupIds := []int{}
	grMap := make(map[int]MemberRole)
	if rows, err := ctx.DB.Query("SELECT group_id, role FROM user_group WHERE user_id=?", user.GetId()); err != nil {
		if err == sql.ErrNoRows {
			return []GroupRole{}, nil
		}
		return nil, err
	} else {
		for rows.Next() {
			var groupId, role int
			if err := rows.Scan(&groupId, &role); err == nil {
				groupIds = append(groupIds, groupId)
				grMap[groupId] = MemberRole(role)
			}
		}
	}

	groupRoles := make([]GroupRole, 0, len(groupIds))
	if len(groupIds) == 0 {
		return groupRoles, nil
	}
	if groups, err := ListGroups(ctx, groupIds...); err != nil {
		return nil, err
	} else {
		for _, g := range groups {
			gr := GroupRole{
				Group: g,
				Role:  grMap[g.Id],
			}
			groupRoles = append(groupRoles, gr)
		}
	}
	return groupRoles, nil
}

// must be recursive to find out all the groups
func GetGroupsOfUser(ctx *models.Context, user iuser.User) ([]Group, error) {
	groups := []Group{}
	mapGroups, err := getGroupsRecursivelyOfUser(ctx, user, false)
	if err != nil {
		panic(err)
	}
	for k, _ := range mapGroups {
		g, err := GetGroup(ctx, k)
		if err != nil {
			panic(err)
		}
		groups = append(groups, *g)
	}
	return groups, nil
}

func getGroupsDirectlyOfUser(ctx *models.Context, user iuser.User) ([]Group, error) {
	groupIds := []int{}
	err := ctx.DB.Select(&groupIds,
		"SELECT group_id FROM user_group WHERE user_id=?",
		user.GetId())
	groups := []Group{}
	if err != nil {
		if err == sql.ErrNoRows {
			return groups, nil
		}
		return nil, err
	}

	for _, gid := range groupIds {
		g, err := GetGroup(ctx, gid)
		if err != nil {
			if err != ErrGroupNotFound {
				return nil, err
			} else {
				log.Warnf("User %s in a non-exist group %d?", user.GetName(), gid)
			}
		} else {
			groups = append(groups, *g)
		}
	}

	ub := ctx.Back
	var bGroups []iuser.BackendGroup
	if ubg, ok := ub.(iuser.BackendWithGroup); ok {
		bGroups, err = ubg.GetBackendGroupsOfUser(user)
		if err != nil {
			panic(err)
		}
	}
	for _, g := range bGroups {
		groups = append(groups, Group{
			Id:       g.GetId(),
			Name:     g.GetName(),
			FullName: g.GetRules().(string),
		})
	}

	return groups, nil
}

func RemoveUserFromAllGroups(ctx *models.Context, user iuser.User) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec("DELETE FROM user_group WHERE user_id=?", user.GetId())

	if err2 := tx.Commit(); err2 != nil {
		return err2
	}

	if err1 != nil {
		return err1
	}

	return nil
}

func backendUsersToMembers(users []iuser.User) []Member {
	members := make([]Member, len(users))
	for i, u := range users {
		members[i] = Member{
			User: u,
			Role: NORMAL,
		}
	}
	return members
}

func getGroupRolesRecursivelyOfUser(ctx *models.Context, user iuser.User, adminOnly bool) (map[int]MemberRole, error) {
	if adminOnly {
		return getGroupsRecursivelyOfUser(ctx, user, true)
	} else {
		admins, err := getGroupsRecursivelyOfUser(ctx, user, true)
		if err != nil {
			panic(err)
		}
		ret, err := getGroupsRecursivelyOfUser(ctx, user, false)
		if err != nil {
			panic(err)
		}
		for k, v := range admins {
			ret[k] = v
		}
		return ret, nil
	}
}

func getGroupsRecursivelyOfUser(ctx *models.Context, user iuser.User, adminOnly bool) (map[int]MemberRole, error) {
	l := list.New()
	ret := make(map[int]MemberRole)
	if adminOnly {
		groupRoles, err := getGroupRolesDirectlyOfUser(ctx, user)
		if err != nil {
			return nil, err
		}
		for _, v := range groupRoles {
			if v.Role == ADMIN {
				ret[v.Id] = ADMIN
				l.PushBack(v.Id)
			}
		}
	} else {
		groups, err := getGroupsDirectlyOfUser(ctx, user)
		if err != nil {
			return nil, err
		}
		for _, v := range groups {
			ret[v.Id] = NORMAL
			l.PushBack(v.Id)
		}
	}
	for l.Len() > 0 {
		iter := l.Front()
		v := iter.Value.(int)
		l.Remove(iter)
		fathers, err := ListGroupFathersById(ctx, v)
		if err != nil {
			panic(err)
		}
		for _, g := range fathers {
			if adminOnly {
				if _, ok := ret[g.Id]; ok {
					continue
				}
				role, err := GetGroupMemberRole(ctx, g.Id, v)
				if err != nil {
					panic(err)
				}
				if role != ADMIN {
					continue
				}
				l.PushBack(g.Id)
				ret[g.Id] = ADMIN
			} else {
				if _, ok := ret[g.Id]; !ok {
					l.PushBack(g.Id)
					ret[g.Id] = NORMAL
				}
			}
		}
	}
	return ret, nil
}
<commit_msg>add log for a strange panic<commit_after>package group

import (
	"container/list"
	"database/sql"
	"sort"

	"github.com/jmoiron/sqlx"
	"github.com/mijia/sweb/log"

	"github.com/laincloud/sso/ssolib/models"
	"github.com/laincloud/sso/ssolib/models/iuser"
)

type MemberRole int8

const (
	NORMAL MemberRole = iota
	ADMIN
)

type Member struct {
	iuser.User
	Role MemberRole
}

type GroupRole struct {
	Group
	Role MemberRole // some user's role in the group
}

const createUserGroupTableSQL = `
CREATE TABLE IF NOT EXISTS user_group (
	user_id INT NOT NULL,
	group_id INT NOT NULL,
	role TINYINT NOT NULL COMMENT '0: normal, 1: admin',
	PRIMARY KEY (user_id, group_id),
	KEY (group_id)
)`

func (g *Group) AddMember(ctx *models.Context, user iuser.User, role MemberRole) error {
	tx := ctx.DB.MustBegin()
	log.Debug("begin add:", user.GetId(), g.Id, role)
	_, err1 := tx.Exec(
		"INSERT INTO user_group (user_id, group_id, role) VALUES (?, ?, ?)",
		user.GetId(), g.Id, role)

	err2 := tx.Commit()

	if err1 != nil {
		return err1
	}

	if err2 != nil {
		return err2
	}

	return nil
}

func (g *Group) UpdateMember(ctx *models.Context, user iuser.User, role MemberRole) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec(
		"UPDATE user_group SET role=? WHERE user_id=? AND group_id=?",
		role, user.GetId(), g.Id)

	err2 := tx.Commit()

	if err1 != nil {
		return err1
	}

	if err2 != nil {
		return err2
	}

	return nil
}

func (g *Group) RemoveMember(ctx *models.Context, user iuser.User) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec("DELETE FROM user_group WHERE user_id=? AND group_id=?",
		user.GetId(), g.Id)

	if err2 := tx.Commit(); err2 != nil {
		return err2
	}

	if err1 != nil {
		return err1
	}

	return nil
}

// not recursive for nested groups
func (g *Group) ListMembers(ctx *models.Context) ([]Member, error) {
	back := ctx.Back
	members := []Member{}
	if g.GroupType == iuser.SSOLIBGROUP {
		rows, err := ctx.DB.Query(
			"SELECT user_id, role FROM user_group WHERE group_id=?",
			g.Id)
		if err != nil {
			if err == sql.ErrNoRows {
				return members, nil
			}
			return nil, err
		}

		for rows.Next() {
			var userId int
			var role MemberRole
			if err = rows.Scan(&userId, &role); err != nil {
				return nil, err
			}
			user, err := back.GetUser(userId)
			if err != nil {
				return nil, err
			}

			members = append(members, Member{user, role})
		}

		return members, nil
	} else if g.GroupType == iuser.BACKENDGROUP {
		var ubg iuser.BackendWithGroup
		var ok bool
		if ubg, ok = back.(iuser.BackendWithGroup); !ok {
			return nil, ErrBackendUnsupported
		}
		bg, err := ubg.GetBackendGroupByName(g.Name)
		if err != nil {
			panic(err)
		}
		users, err := bg.ListUsers()
		if err != nil {
			if err == iuser.ErrMethodNotSupported {
				return members, nil
			} else {
				panic(err)
			}
		}
		members = backendUsersToMembers(users)
		return members, nil
	} else {
		panic("here")
	}
	return nil, nil
}

// Return (true, role, nil) if u is member of g, otherwise return (false, 0, nil).
// error will be non-nil if anything unexpected happens.
// Must considering recursive if valid
func (g *Group) GetMember(ctx *models.Context, u iuser.User) (ok bool, role MemberRole, err error) {
	/*err = ctx.DB.QueryRow("SELECT role FROM user_group WHERE user_id=? AND group_id=?",
		u.GetId(), g.Id).Scan(&role)
	switch {
	case err == sql.ErrNoRows:
		err = nil
		return
	case err != nil:
		return
	default:
		ok = true
		return
	}*/
	groups, err := getGroupRolesRecursivelyOfUser(ctx, u, false)
	if err != nil {
		return false, 0, err
	}
	if role, ok = groups[g.Id]; ok {
		return true, role, nil
	} else {
		return false, 0, nil
	}
}

// must not recursive
func (g *Group) removeAllMembers(ctx *models.Context) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec("DELETE FROM user_group WHERE group_id=?", g.Id)

	if err2 := tx.Commit(); err2 != nil {
		return err2
	}

	if err1 != nil {
		return err1
	}
	return nil
}

// direct groups of the users, for now it is used only in "管理员特供"
func GetGroupsOfUserByIds(ctx *models.Context, userIds []int) (map[int][]Group, error) {
	groupMap := make(map[int][]Group)
	if len(userIds) == 0 {
		return groupMap, nil
	}

	userGroupIds := make(map[int][]int)
	groupIds := []int{}
	if query, args, err := sqlx.In("SELECT group_id, user_id FROM user_group WHERE user_id IN(?)", userIds); err != nil {
		if err == sql.ErrNoRows {
			err = nil
		}
		return groupMap, err
	} else {
		if rows, err := ctx.DB.Query(query, args...); err != nil {
			return groupMap, err
		} else {
			for rows.Next() {
				var (
					userId  int
					groupId int
				)
				if err := rows.Scan(&groupId, &userId); err != nil {
					return groupMap, err
				}
				groupIds = append(groupIds, groupId)
				userGroupIds[userId] = append(userGroupIds[userId], groupId)
			}
		}
	}

	groups, err := ListGroups(ctx, groupIds...)
	if err != nil {
		return groupMap, err
	}
	groupSet := make(map[int]Group)
	for _, g := range groups {
		groupSet[g.Id] = g
	}
	for userId, groups := range userGroupIds {
		for _, g := range groups {
			if group, ok := groupSet[g]; ok {
				groupMap[userId] = append(groupMap[userId], group)
			}
		}
	}
	return groupMap, nil
}

// should be recursive, but may be slow if we find recursive roles
func GetGroupRolesOfUser(ctx *models.Context, user iuser.User) ([]GroupRole, error) {
	roles := []GroupRole{}
	mapRoles, err := getGroupRolesRecursivelyOfUser(ctx, user, false)
	if err != nil {
		log.Error("strange:", user)
		panic(err)
	}
	sortedGroupIds := []int{}
	for k, _ := range mapRoles {
		sortedGroupIds = append(sortedGroupIds, k)
	}
	sort.Ints(sortedGroupIds)
	for _, k := range sortedGroupIds {
		g, err := GetGroup(ctx, k)
		v := mapRoles[k]
		if err != nil {
			panic(err)
		}
		tmp := GroupRole{
			Group: *g,
			Role:  v,
		}
		roles = append(roles, tmp)
	}
	return roles, nil
}

func getGroupRolesDirectlyOfUser(ctx *models.Context, user iuser.User) ([]GroupRole, error) {
	groupRoles, err := getSSOLIBGroupRolesDirectlyOfUser(ctx, user)
	ub := ctx.Back
	var bGroups []iuser.BackendGroup
	if ubg, ok := ub.(iuser.BackendWithGroup); ok {
		bGroups, err = ubg.GetBackendGroupsOfUser(user)
		if err != nil {
			panic(err)
		}
	}
	ret := make([]GroupRole, (len(groupRoles) + len(bGroups)))
	copy(ret, groupRoles)
	for i, g := range bGroups {
		role := NORMAL
		ret[(i + len(groupRoles))] = GroupRole{
			Group: Group{
				Id:       g.GetId(),
				Name:     g.GetName(),
				FullName: g.GetRules().(string),
			},
			Role: role,
		}
	}
	return ret, nil
}

func getSSOLIBGroupRolesDirectlyOfUser(ctx *models.Context, user iuser.User) ([]GroupRole, error) {
	groupIds := []int{}
	grMap := make(map[int]MemberRole)
	if rows, err := ctx.DB.Query("SELECT group_id, role FROM user_group WHERE user_id=?", user.GetId()); err != nil {
		if err == sql.ErrNoRows {
			return []GroupRole{}, nil
		}
		return nil, err
	} else {
		for rows.Next() {
			var groupId, role int
			if err := rows.Scan(&groupId, &role); err == nil {
				groupIds = append(groupIds, groupId)
				grMap[groupId] = MemberRole(role)
			}
		}
	}

	groupRoles := make([]GroupRole, 0, len(groupIds))
	if len(groupIds) == 0 {
		return groupRoles, nil
	}
	if groups, err := ListGroups(ctx, groupIds...); err != nil {
		return nil, err
	} else {
		for _, g := range groups {
			gr := GroupRole{
				Group: g,
				Role:  grMap[g.Id],
			}
			groupRoles = append(groupRoles, gr)
		}
	}
	return groupRoles, nil
}

// must be recursive to find out all the groups
func GetGroupsOfUser(ctx *models.Context, user iuser.User) ([]Group, error) {
	groups := []Group{}
	mapGroups, err := getGroupsRecursivelyOfUser(ctx, user, false)
	if err != nil {
		panic(err)
	}
	for k, _ := range mapGroups {
		g, err := GetGroup(ctx, k)
		if err != nil {
			panic(err)
		}
		groups = append(groups, *g)
	}
	return groups, nil
}

func getGroupsDirectlyOfUser(ctx *models.Context, user iuser.User) ([]Group, error) {
	groupIds := []int{}
	err := ctx.DB.Select(&groupIds,
		"SELECT group_id FROM user_group WHERE user_id=?",
		user.GetId())
	groups := []Group{}
	if err != nil {
		if err == sql.ErrNoRows {
			return groups, nil
		}
		return nil, err
	}

	for _, gid := range groupIds {
		g, err := GetGroup(ctx, gid)
		if err != nil {
			if err != ErrGroupNotFound {
				return nil, err
			} else {
				log.Warnf("User %s in a non-exist group %d?", user.GetName(), gid)
			}
		} else {
			groups = append(groups, *g)
		}
	}

	ub := ctx.Back
	var bGroups []iuser.BackendGroup
	if ubg, ok := ub.(iuser.BackendWithGroup); ok {
		bGroups, err = ubg.GetBackendGroupsOfUser(user)
		if err != nil {
			panic(err)
		}
	}
	for _, g := range bGroups {
		groups = append(groups, Group{
			Id:       g.GetId(),
			Name:     g.GetName(),
			FullName: g.GetRules().(string),
		})
	}

	return groups, nil
}

func RemoveUserFromAllGroups(ctx *models.Context, user iuser.User) error {
	tx := ctx.DB.MustBegin()
	_, err1 := tx.Exec("DELETE FROM user_group WHERE user_id=?", user.GetId())

	if err2 := tx.Commit(); err2 != nil {
		return err2
	}

	if err1 != nil {
		return err1
	}

	return nil
}

func backendUsersToMembers(users []iuser.User) []Member {
	members := make([]Member, len(users))
	for i, u := range users {
		members[i] = Member{
			User: u,
			Role: NORMAL,
		}
	}
	return members
}

func getGroupRolesRecursivelyOfUser(ctx *models.Context, user iuser.User, adminOnly bool) (map[int]MemberRole, error) {
	if adminOnly {
		return getGroupsRecursivelyOfUser(ctx, user, true)
	} else {
		admins, err := getGroupsRecursivelyOfUser(ctx, user, true)
		if err != nil {
			panic(err)
		}
		ret, err := getGroupsRecursivelyOfUser(ctx, user, false)
		if err != nil {
			panic(err)
		}
		for k, v := range admins {
			ret[k] = v
		}
		return ret, nil
	}
}

func getGroupsRecursivelyOfUser(ctx *models.Context, user iuser.User, adminOnly bool) (map[int]MemberRole, error) {
	l := list.New()
	ret := make(map[int]MemberRole)
	if adminOnly {
		groupRoles, err := getGroupRolesDirectlyOfUser(ctx, user)
		if err != nil {
			return nil, err
		}
		for _, v := range groupRoles {
			if v.Role == ADMIN {
				ret[v.Id] = ADMIN
				l.PushBack(v.Id)
			}
		}
	} else {
		groups, err := getGroupsDirectlyOfUser(ctx, user)
		if err != nil {
			return nil, err
		}
		for _, v := range groups {
			ret[v.Id] = NORMAL
			l.PushBack(v.Id)
		}
	}
	for l.Len() > 0 {
		iter := l.Front()
		v := iter.Value.(int)
		l.Remove(iter)
		fathers, err := ListGroupFathersById(ctx, v)
		if err != nil {
			panic(err)
		}
		for _, g := range fathers {
			if adminOnly {
				if _, ok := ret[g.Id]; ok {
					continue
				}
				role, err := GetGroupMemberRole(ctx, g.Id, v)
				if err != nil {
					panic(err)
				}
				if role != ADMIN {
					continue
				}
				l.PushBack(g.Id)
				ret[g.Id] = ADMIN
			} else {
				if _, ok := ret[g.Id]; !ok {
					l.PushBack(g.Id)
					ret[g.Id] = NORMAL
				}
			}
		}
	}
	return ret, nil
}
<|endoftext|>


<commit_before>package pollers

/*

Simple External Poller Integration: If you can open a socket, you can write a poller.

Format: <RFC3339 date stamp> <what> <value>\n

The exact interpretation of these depends on the Outputter in use.

Example

In terminal A:
  SHH_POLLERS=listen ./shh

In a different terminal:
  (while true; do echo $(date "+%Y-%m-%dT%H:%M:%SZ") memfree $(grep MemFree /proc/meminfo | awk '{print $2}').0; sleep 5; done) | nc -U /tmp/shh

*/

import (
	"bufio"
	"github.com/freeformz/shh/config"
	"github.com/freeformz/shh/mm"
	"github.com/freeformz/shh/utils"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	listenNet   string
	listenLaddr string
)

func init() {
	tmp := strings.Split(config.Listen, ",")

	if len(tmp) != 2 {
		log.Fatal("SHH_LISTEN is not in the format: 'unix,#shh")
	}

	listenNet = tmp[0]
	listenLaddr = tmp[1]

	switch listenNet {
	case "tcp", "tcp4", "tcp6", "unix", "unixpacket":
		break
	default:
		log.Fatalf("SHH_LISTEN format (%s,%s) is not correct", listenNet, listenLaddr)
	}

	// If this is a path, remove it
	if listenNet == "unix" && utils.Exists(listenLaddr) {
		err := os.Remove(listenLaddr)
		if err != nil {
			log.Fatal(err)
		}
	}

}

// Used to track global listen stats
type ListenStats struct {
	sync.RWMutex
	counts map[string]interface{}
}

func (ls *ListenStats) New(what string, initialValue interface{}) {
	ls.Lock()
	defer ls.Unlock()
	switch initialValue.(type) {
	case float64, uint64:
		ls.counts[what] = initialValue
	case int:
		ls.counts[what] = uint64(initialValue.(int))
	}
}

func (ls *ListenStats) Increment(what string) {
	ls.Lock()
	defer ls.Unlock()
	v := ls.counts[what]
	switch v.(type) {
	case float64:
		tmp := v.(float64)
		tmp++
		ls.counts[what] = tmp
	case uint64, int:
		tmp := v.(uint64)
		tmp++
		ls.counts[what] = tmp
	}
}

func (ls *ListenStats) Decrement(what string) {
	ls.Lock()
	defer ls.Unlock()
	v := ls.counts[what]
	switch v.(type) {
	case float64:
		tmp := v.(float64)
		tmp--
		ls.counts[what] = tmp
	case uint64, int:
		tmp := v.(uint64)
		tmp--
		ls.counts[what] = tmp
	}
}

func (ls *ListenStats) CountOf(what string) interface{} {
	ls.RLock()
	defer ls.RUnlock()
	return ls.counts[what]
}

func (ls *ListenStats) Keys() <-chan string {
	ls.RLock()

	c := make(chan string)

	go func(c chan<- string) {
		defer ls.RUnlock()
		defer close(c)
		for k, _ := range ls.counts {
			c <- k
		}
	}(c)

	return c
}

type Listen struct {
	measurements chan<- *mm.Measurement
	listener     net.Listener
	stats        *ListenStats
}

func NewListenPoller(measurements chan<- *mm.Measurement) Listen {
	listener, err := net.Listen(listenNet, listenLaddr)

	if err != nil {
		log.Fatal(err)
	}

	ls := &ListenStats{counts: make(map[string]interface{})}
	ls.New("connection.count", 0.0)
	ls.New("time.parse.errors", 0)
	ls.New("value.parse.errors", 0)
	ls.New("metrics", 0)

	poller := Listen{measurements: measurements, listener: listener, stats: ls}

	go func(poller *Listen) {
		for {
			conn, err := poller.listener.Accept()
			if err != nil {
				log.Println(err)
				continue
			}

			go handleListenConnection(poller, conn)
		}
	}(&poller)

	return poller
}

func (poller Listen) Poll(tick time.Time) {
	for k := range poller.stats.Keys() {
		poller.measurements <- &mm.Measurement{tick, poller.Name(), strings.Split("stats."+k, "."), poller.stats.CountOf(k)}
	}
}

func handleListenConnection(poller *Listen, conn net.Conn) {
	defer conn.Close()

	var value interface{}

	poller.stats.Increment("connection.count")
	defer poller.stats.Decrement("connection.count")

	r := bufio.NewReader(conn)

	for {
		conn.SetDeadline(time.Now().Add(config.Interval).Add(config.Interval))
		line, err := r.ReadString('\n')
		if err != nil {
			break
		}

		fields := strings.Fields(line)
		if len(fields) == 3 {
			when, err := time.Parse(time.RFC3339, fields[0])
			if err != nil {
				poller.stats.Increment("time.parse.errors")
				break
			}
			value, err = strconv.ParseUint(fields[2], 10, 64)
			if err != nil {
				value, err = strconv.ParseFloat(fields[2], 64)
				if err != nil {
					poller.stats.Increment("value.parse.errors")
					break
				}
			}

			poller.stats.Increment("metrics")

			poller.measurements <- &mm.Measurement{when, poller.Name(), strings.Fields(fields[1]), value}
		}
	}
}

func (poller Listen) Name() string {
	return "listen"
}

func (poller Listen) Exit() {
	poller.listener.Close()
}
<commit_msg>Udate the listen example<commit_after>package pollers

/*

Simple External Poller Integration: If you can open a socket, you can write a poller.

Format: <RFC3339 date stamp> <what> <value>\n

The exact interpretation of these depends on the Outputter in use.

Example

In terminal A:
  SHH_POLLERS=listen ./shh

In a different terminal:
  (while true; do echo $(date "+%Y-%m-%dT%H:%M:%SZ") memfree $(grep MemFree /proc/meminfo | awk '{print $2}').0; sleep 5; done) | nc -U \#shh

*/

import (
	"bufio"
	"github.com/freeformz/shh/config"
	"github.com/freeformz/shh/mm"
	"github.com/freeformz/shh/utils"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	listenNet   string
	listenLaddr string
)

func init() {
	tmp := strings.Split(config.Listen, ",")

	if len(tmp) != 2 {
		log.Fatal("SHH_LISTEN is not in the format: 'unix,#shh")
	}

	listenNet = tmp[0]
	listenLaddr = tmp[1]

	switch listenNet {
	case "tcp", "tcp4", "tcp6", "unix", "unixpacket":
		break
	default:
		log.Fatalf("SHH_LISTEN format (%s,%s) is not correct", listenNet, listenLaddr)
	}

	// If this is a path, remove it
	if listenNet == "unix" && utils.Exists(listenLaddr) {
		err := os.Remove(listenLaddr)
		if err != nil {
			log.Fatal(err)
		}
	}

}

// Used to track global listen stats
type ListenStats struct {
	sync.RWMutex
	counts map[string]interface{}
}

func (ls *ListenStats) New(what string, initialValue interface{}) {
	ls.Lock()
	defer ls.Unlock()
	switch initialValue.(type) {
	case float64, uint64:
		ls.counts[what] = initialValue
	case int:
		ls.counts[what] = uint64(initialValue.(int))
	}
}

func (ls *ListenStats) Increment(what string) {
	ls.Lock()
	defer ls.Unlock()
	v := ls.counts[what]
	switch v.(type) {
	case float64:
		tmp := v.(float64)
		tmp++
		ls.counts[what] = tmp
	case uint64, int:
		tmp := v.(uint64)
		tmp++
		ls.counts[what] = tmp
	}
}

func (ls *ListenStats) Decrement(what string) {
	ls.Lock()
	defer ls.Unlock()
	v := ls.counts[what]
	switch v.(type) {
	case float64:
		tmp := v.(float64)
		tmp--
		ls.counts[what] = tmp
	case uint64, int:
		tmp := v.(uint64)
		tmp--
		ls.counts[what] = tmp
	}
}

func (ls *ListenStats) CountOf(what string) interface{} {
	ls.RLock()
	defer ls.RUnlock()
	return ls.counts[what]
}

func (ls *ListenStats) Keys() <-chan string {
	ls.RLock()

	c := make(chan string)

	go func(c chan<- string) {
		defer ls.RUnlock()
		defer close(c)
		for k, _ := range ls.counts {
			c <- k
		}
	}(c)

	return c
}

type Listen struct {
	measurements chan<- *mm.Measurement
	listener     net.Listener
	stats        *ListenStats
}

func NewListenPoller(measurements chan<- *mm.Measurement) Listen {
	listener, err := net.Listen(listenNet, listenLaddr)

	if err != nil {
		log.Fatal(err)
	}

	ls := &ListenStats{counts: make(map[string]interface{})}
	ls.New("connection.count", 0.0)
	ls.New("time.parse.errors", 0)
	ls.New("value.parse.errors", 0)
	ls.New("metrics", 0)

	poller := Listen{measurements: measurements, listener: listener, stats: ls}

	go func(poller *Listen) {
		for {
			conn, err := poller.listener.Accept()
			if err != nil {
				log.Println(err)
				continue
			}

			go handleListenConnection(poller, conn)
		}
	}(&poller)

	return poller
}

func (poller Listen) Poll(tick time.Time) {
	for k := range poller.stats.Keys() {
		poller.measurements <- &mm.Measurement{tick, poller.Name(), strings.Split("stats."+k, "."), poller.stats.CountOf(k)}
	}
}

func handleListenConnection(poller *Listen, conn net.Conn) {
	defer conn.Close()

	var value interface{}

	poller.stats.Increment("connection.count")
	defer poller.stats.Decrement("connection.count")

	r := bufio.NewReader(conn)

	for {
		conn.SetDeadline(time.Now().Add(config.Interval).Add(config.Interval))
		line, err := r.ReadString('\n')
		if err != nil {
			break
		}

		fields := strings.Fields(line)
		if len(fields) == 3 {
			when, err := time.Parse(time.RFC3339, fields[0])
			if err != nil {
				poller.stats.Increment("time.parse.errors")
				break
			}
			value, err = strconv.ParseUint(fields[2], 10, 64)
			if err != nil {
				value, err = strconv.ParseFloat(fields[2], 64)
				if err != nil {
					poller.stats.Increment("value.parse.errors")
					break
				}
			}

			poller.stats.Increment("metrics")

			poller.measurements <- &mm.Measurement{when, poller.Name(), strings.Fields(fields[1]), value}
		}
	}
}

func (poller Listen) Name() string {
	return "listen"
}

func (poller Listen) Exit() {
	poller.listener.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"expvar"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/NYTimes/gziphandler"
	"github.com/mat/besticon/besticon"
	"github.com/mat/besticon/besticon/iconserver/assets"
	"github.com/mat/besticon/lettericon"
)

func indexHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "" || r.URL.Path == "/" {
		renderHTMLTemplate(w, 200, indexHTML, nil)
	} else {
		renderHTMLTemplate(w, 404, notFoundHTML, nil)
	}
}

func iconsHandler(w http.ResponseWriter, r *http.Request) {
	url := r.FormValue(urlParam)
	if len(url) == 0 {
		http.Redirect(w, r, "/", 302)
		return
	}

	finder := besticon.IconFinder{}

	formats := r.FormValue("formats")
	if formats != "" {
		finder.FormatsAllowed = strings.Split(r.FormValue("formats"), ",")
	}

	e, icons := finder.FetchIcons(url)
	switch {
	case e != nil:
		renderHTMLTemplate(w, 404, iconsHTML, pageInfo{URL: url, Error: e})
	case len(icons) == 0:
		errNoIcons := errors.New("this poor site has no icons at all :-(")
		renderHTMLTemplate(w, 404, iconsHTML, pageInfo{URL: url, Error: errNoIcons})
	default:
		renderHTMLTemplate(w, 200, iconsHTML, pageInfo{Icons: icons, URL: url})
	}
}

func iconHandler(w http.ResponseWriter, r *http.Request) {
	url := r.FormValue("url")
	if len(url) == 0 {
		writeAPIError(w, 400, errors.New("need url parameter"), true)
		return
	}

	size := r.FormValue("size")
	if size == "" {
		writeAPIError(w, 400, errors.New("need size parameter"), true)
		return
	}
	minSize, err := strconv.Atoi(size)
	if err != nil || minSize < 0 || minSize > 500 {
		writeAPIError(w, 400, errors.New("bad size parameter"), true)
		return
	}

	finder := besticon.IconFinder{}
	formats := r.FormValue("formats")
	if formats != "" {
		finder.FormatsAllowed = strings.Split(r.FormValue("formats"), ",")
	}

	err, _ = finder.FetchIcons(url)
	if err != nil {
		writeAPIError(w, 404, err, true)
		return
	}

	icon := finder.IconWithMinSize(minSize)
	if icon != nil {
		http.Redirect(w, r, icon.URL, 302)
		return
	}

	fallbackIconURL := r.FormValue("fallback_icon_url")
	if fallbackIconURL != "" {
		http.Redirect(w, r, fallbackIconURL, 302)
		return
	}

	iconColor := finder.MainColorForIcons()
	letter := lettericon.MainLetterFromURL(url)
	redirectPath := lettericon.IconPath(letter, size, iconColor)
	http.Redirect(w, r, redirectPath, 302)
}

func popularHandler(w http.ResponseWriter, r *http.Request) {
	iconSize, err := strconv.Atoi(r.FormValue("iconsize"))
	if iconSize > 500 || iconSize < 10 || err != nil {
		iconSize = 120
	}

	pageInfo := struct {
		URLs        []string
		IconSize    int
		DisplaySize int
	}{
		besticon.PopularSites,
		iconSize,
		iconSize / 2,
	}
	renderHTMLTemplate(w, 200, popularHTML, pageInfo)
}

const (
	urlParam    = "url"
	prettyParam = "pretty"
	maxAge      = "max_age"
)

const defaultMaxAge = time.Duration(604800) * time.Second // 7 days

func alliconsHandler(w http.ResponseWriter, r *http.Request) {
	url := r.FormValue(urlParam)
	if len(url) == 0 {
		errMissingURL := errors.New("need url query parameter")
		writeAPIError(w, 400, errMissingURL, true)
		return
	}

	finder := besticon.IconFinder{}
	formats := r.FormValue("formats")
	if formats != "" {
		finder.FormatsAllowed = strings.Split(r.FormValue("formats"), ",")
	}

	e, icons := finder.FetchIcons(url)
	if e != nil {
		writeAPIError(w, 404, e, true)
		return
	}

	pretty, err := strconv.ParseBool(r.FormValue(prettyParam))
	prettyPrint := (err == nil) && pretty

	writeAPIIcons(w, url, icons, prettyPrint)
}

func lettericonHandler(w http.ResponseWriter, r *http.Request) {
	charParam, col, size := lettericon.ParseIconPath(r.URL.Path)
	if charParam != "" {
		w.Header().Add("Content-Type", "image/png")
		lettericon.Render(charParam, col, size, w)
	} else {
		writeAPIError(w, 400, errors.New("wrong format for lettericons/ path, must look like lettericons/M-144-EFC25D.png"), true)
	}
}

func obsoleteAPIHandler(w http.ResponseWriter, r *http.Request) {
	if r.FormValue("i_am_feeling_lucky") == "yes" {
		http.Redirect(w, r, fmt.Sprintf("/icon?size=120&%s", r.URL.RawQuery), 302)
	} else {
		http.Redirect(w, r, fmt.Sprintf("/allicons.json?%s", r.URL.RawQuery), 302)
	}
}

func writeAPIError(w http.ResponseWriter, httpStatus int, e error, pretty bool) {
	data := struct {
		Error string `json:"error"`
	}{
		e.Error(),
	}

	if pretty {
		renderJSONResponsePretty(w, httpStatus, data)
	} else {
		renderJSONResponse(w, httpStatus, data)
	}
}

func writeAPIIcons(w http.ResponseWriter, url string, icons []besticon.Icon, pretty bool) {
	// Don't return whole image data
	newIcons := []besticon.Icon{}
	for _, ico := range icons {
		newIcon := ico
		newIcon.ImageData = nil
		newIcons = append(newIcons, newIcon)
	}

	data := &struct {
		URL   string          `json:"url"`
		Icons []besticon.Icon `json:"icons"`
	}{
		url,
		newIcons,
	}

	if pretty {
		renderJSONResponsePretty(w, 200, data)
	} else {
		renderJSONResponse(w, 200, data)
	}
}

const (
	contentType     = "Content-Type"
	applicationJSON = "application/json"
	cacheControl    = "Cache-Control"
)

func renderJSONResponse(w http.ResponseWriter, httpStatus int, data interface{}) {
	w.Header().Add(contentType, applicationJSON)
	w.WriteHeader(httpStatus)
	enc := json.NewEncoder(w)
	enc.Encode(data)
}

func renderJSONResponsePretty(w http.ResponseWriter, httpStatus int, data interface{}) {
	w.Header().Add(contentType, applicationJSON)
	w.WriteHeader(httpStatus)
	b, _ := json.MarshalIndent(data, "", "  ")
	w.Write(b)
}

type pageInfo struct {
	URL   string
	Icons []besticon.Icon
	Error error
}

func (pi pageInfo) Host() string {
	u := pi.URL
	url, _ := url.Parse(u)
	if url != nil && url.Host != "" {
		return url.Host
	}
	return pi.URL
}

func (pi pageInfo) Best() string {
	if len(pi.Icons) > 0 {
		best := pi.Icons[0]
		return best.URL
	}
	return ""
}

func renderHTMLTemplate(w http.ResponseWriter, httpStatus int, templ *template.Template, data interface{}) {
	w.Header().Add(contentType, "text/html; charset=utf-8")
	w.WriteHeader(httpStatus)

	err := templ.Execute(w, data)
	if err != nil {
		err = fmt.Errorf("server: could not generate output: %s", err)
		logger.Print(err)
		w.Write([]byte(err.Error()))
	}
}

func startServer(port string) {
	registerGzipHandler("/", indexHandler)
	registerGzipHandler("/icons", iconsHandler)
	registerHandler("/icon", iconHandler)
	registerGzipHandler("/popular", popularHandler)
	registerGzipHandler("/allicons.json", alliconsHandler)
	registerHandler("/lettericons/", lettericonHandler)
	registerHandler("/api/icons", obsoleteAPIHandler)

	serveAsset("/pure-0.5.0-min.css", "besticon/iconserver/assets/pure-0.5.0-min.css", oneYear)
	serveAsset("/grids-responsive-0.5.0-min.css", "besticon/iconserver/assets/grids-responsive-0.5.0-min.css", oneYear)
	serveAsset("/main-min.css", "besticon/iconserver/assets/main-min.css", oneYear)

	serveAsset("/icon.svg", "besticon/iconserver/assets/icon.svg", oneYear)
	serveAsset("/favicon.ico", "besticon/iconserver/assets/favicon.ico", oneYear)
	serveAsset("/apple-touch-icon.png", "besticon/iconserver/assets/apple-touch-icon.png", oneYear)
	serveAsset("/robots.txt", "besticon/iconserver/assets/robots.txt", nocache)
	serveAsset("/test-lettericons", "besticon/iconserver/assets/test-lettericons.html", nocache)

	addr := "0.0.0.0:" + port
	logger.Print("Starting server on ", addr, "...")
	e := http.ListenAndServe(addr, newLoggingMux())
	if e != nil {
		logger.Fatalf("cannot start server: %s\n", e)
	}
}

const (
	oneYear = 365 * 24 * 3600
	nocache = -1
)

func serveAsset(path string, assetPath string, maxAgeSeconds int) {
	registerGzipHandler(path, func(w http.ResponseWriter, r *http.Request) {
		assetInfo, err := assets.AssetInfo(assetPath)
		if err != nil {
			panic(err)
		}

		if maxAgeSeconds != nocache {
			w.Header().Add(cacheControl, fmt.Sprintf("max-age=%d", maxAgeSeconds))
		}

		http.ServeContent(w, r, assetInfo.Name(), assetInfo.ModTime(),
			bytes.NewReader(assets.MustAsset(assetPath)))
	})
}

func registerHandler(path string, f http.HandlerFunc) {
	http.Handle(path, newExpvarHandler(path, f))
}

func registerGzipHandler(path string, f http.HandlerFunc) {
	http.Handle(path, gziphandler.GzipHandler(newExpvarHandler(path, f)))
}

func main() {
	fmt.Printf("iconserver %s (%s) - http://icons.better-idea.org\n", besticon.VersionString, runtime.Version())
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}
	startServer(port)
}

func init() {
	indexHTML = templateFromAsset("besticon/iconserver/assets/index.html", "index.html")
	iconsHTML = templateFromAsset("besticon/iconserver/assets/icons.html", "icons.html")
	popularHTML = templateFromAsset("besticon/iconserver/assets/popular.html", "popular.html")
	notFoundHTML = templateFromAsset("besticon/iconserver/assets/not_found.html", "not_found.html")
}

func templateFromAsset(assetPath, templateName string) *template.Template {
	bytes := assets.MustAsset(assetPath)
	return template.Must(template.New(templateName).Funcs(funcMap).Parse(string(bytes)))
}

var indexHTML *template.Template
var iconsHTML *template.Template
var popularHTML *template.Template
var notFoundHTML *template.Template

var funcMap = template.FuncMap{
	"ImgWidth": imgWidth,
}

func imgWidth(i *besticon.Icon) int {
	return i.Width / 2.0
}

func init() {
	besticon.SetCacheMaxSize(128)

	expvar.Publish("cacheBytes", expvar.Func(func() interface{} { return besticon.GetCacheStats().Bytes }))
	expvar.Publish("cacheItems", expvar.Func(func() interface{} { return besticon.GetCacheStats().Items }))
	expvar.Publish("cacheGets", expvar.Func(func() interface{} { return besticon.GetCacheStats().Gets }))
	expvar.Publish("cacheHits", expvar.Func(func() interface{} { return besticon.GetCacheStats().Hits }))
	expvar.Publish("cacheEvictions", expvar.Func(func() interface{} { return besticon.GetCacheStats().Evictions }))
}
<commit_msg>Use lower in memory cache size<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"expvar"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/NYTimes/gziphandler"
	"github.com/mat/besticon/besticon"
	"github.com/mat/besticon/besticon/iconserver/assets"
	"github.com/mat/besticon/lettericon"
)

func indexHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "" || r.URL.Path == "/" {
		renderHTMLTemplate(w, 200, indexHTML, nil)
	} else {
		renderHTMLTemplate(w, 404, notFoundHTML, nil)
	}
}

func iconsHandler(w http.ResponseWriter, r *http.Request) {
	url := r.FormValue(urlParam)
	if len(url) == 0 {
		http.Redirect(w, r, "/", 302)
		return
	}

	finder := besticon.IconFinder{}

	formats := r.FormValue("formats")
	if formats != "" {
		finder.FormatsAllowed = strings.Split(r.FormValue("formats"), ",")
	}

	e, icons := finder.FetchIcons(url)
	switch {
	case e != nil:
		renderHTMLTemplate(w, 404, iconsHTML, pageInfo{URL: url, Error: e})
	case len(icons) == 0:
		errNoIcons := errors.New("this poor site has no icons at all :-(")
		renderHTMLTemplate(w, 404, iconsHTML, pageInfo{URL: url, Error: errNoIcons})
	default:
		renderHTMLTemplate(w, 200, iconsHTML, pageInfo{Icons: icons, URL: url})
	}
}

func iconHandler(w http.ResponseWriter, r *http.Request) {
	url := r.FormValue("url")
	if len(url) == 0 {
		writeAPIError(w, 400, errors.New("need url parameter"), true)
		return
	}

	size := r.FormValue("size")
	if size == "" {
		writeAPIError(w, 400, errors.New("need size parameter"), true)
		return
	}
	minSize, err := strconv.Atoi(size)
	if err != nil || minSize < 0 || minSize > 500 {
		writeAPIError(w, 400, errors.New("bad size parameter"), true)
		return
	}

	finder := besticon.IconFinder{}
	formats := r.FormValue("formats")
	if formats != "" {
		finder.FormatsAllowed = strings.Split(r.FormValue("formats"), ",")
	}

	err, _ = finder.FetchIcons(url)
	if err != nil {
		writeAPIError(w, 404, err, true)
		return
	}

	icon := finder.IconWithMinSize(minSize)
	if icon != nil {
		http.Redirect(w, r, icon.URL, 302)
		return
	}

	fallbackIconURL := r.FormValue("fallback_icon_url")
	if fallbackIconURL != "" {
		http.Redirect(w, r, fallbackIconURL, 302)
		return
	}

	iconColor := finder.MainColorForIcons()
	letter := lettericon.MainLetterFromURL(url)
	redirectPath := lettericon.IconPath(letter, size, iconColor)
	http.Redirect(w, r, redirectPath, 302)
}

func popularHandler(w http.ResponseWriter, r *http.Request) {
	iconSize, err := strconv.Atoi(r.FormValue("iconsize"))
	if iconSize > 500 || iconSize < 10 || err != nil {
		iconSize = 120
	}

	pageInfo := struct {
		URLs        []string
		IconSize    int
		DisplaySize int
	}{
		besticon.PopularSites,
		iconSize,
		iconSize / 2,
	}
	renderHTMLTemplate(w, 200, popularHTML, pageInfo)
}

const (
	urlParam    = "url"
	prettyParam = "pretty"
	maxAge      = "max_age"
)

const defaultMaxAge = time.Duration(604800) * time.Second // 7 days

func alliconsHandler(w http.ResponseWriter, r *http.Request) {
	url := r.FormValue(urlParam)
	if len(url) == 0 {
		errMissingURL := errors.New("need url query parameter")
		writeAPIError(w, 400, errMissingURL, true)
		return
	}

	finder := besticon.IconFinder{}
	formats := r.FormValue("formats")
	if formats != "" {
		finder.FormatsAllowed = strings.Split(r.FormValue("formats"), ",")
	}

	e, icons := finder.FetchIcons(url)
	if e != nil {
		writeAPIError(w, 404, e, true)
		return
	}

	pretty, err := strconv.ParseBool(r.FormValue(prettyParam))
	prettyPrint := (err == nil) && pretty

	writeAPIIcons(w, url, icons, prettyPrint)
}

func lettericonHandler(w http.ResponseWriter, r *http.Request) {
	charParam, col, size := lettericon.ParseIconPath(r.URL.Path)
	if charParam != "" {
		w.Header().Add("Content-Type", "image/png")
		lettericon.Render(charParam, col, size, w)
	} else {
		writeAPIError(w, 400, errors.New("wrong format for lettericons/ path, must look like lettericons/M-144-EFC25D.png"), true)
	}
}

func obsoleteAPIHandler(w http.ResponseWriter, r *http.Request) {
	if r.FormValue("i_am_feeling_lucky") == "yes" {
		http.Redirect(w, r, fmt.Sprintf("/icon?size=120&%s", r.URL.RawQuery), 302)
	} else {
		http.Redirect(w, r, fmt.Sprintf("/allicons.json?%s", r.URL.RawQuery), 302)
	}
}

func writeAPIError(w http.ResponseWriter, httpStatus int, e error, pretty bool) {
	data := struct {
		Error string `json:"error"`
	}{
		e.Error(),
	}

	if pretty {
		renderJSONResponsePretty(w, httpStatus, data)
	} else {
		renderJSONResponse(w, httpStatus, data)
	}
}

func writeAPIIcons(w http.ResponseWriter, url string, icons []besticon.Icon, pretty bool) {
	// Don't return whole image data
	newIcons := []besticon.Icon{}
	for _, ico := range icons {
		newIcon := ico
		newIcon.ImageData = nil
		newIcons = append(newIcons, newIcon)
	}

	data := &struct {
		URL   string          `json:"url"`
		Icons []besticon.Icon `json:"icons"`
	}{
		url,
		newIcons,
	}

	if pretty {
		renderJSONResponsePretty(w, 200, data)
	} else {
		renderJSONResponse(w, 200, data)
	}
}

const (
	contentType     = "Content-Type"
	applicationJSON = "application/json"
	cacheControl    = "Cache-Control"
)

func renderJSONResponse(w http.ResponseWriter, httpStatus int, data interface{}) {
	w.Header().Add(contentType, applicationJSON)
	w.WriteHeader(httpStatus)
	enc := json.NewEncoder(w)
	enc.Encode(data)
}

func renderJSONResponsePretty(w http.ResponseWriter, httpStatus int, data interface{}) {
	w.Header().Add(contentType, applicationJSON)
	w.WriteHeader(httpStatus)
	b, _ := json.MarshalIndent(data, "", "  ")
	w.Write(b)
}

type pageInfo struct {
	URL   string
	Icons []besticon.Icon
	Error error
}

func (pi pageInfo) Host() string {
	u := pi.URL
	url, _ := url.Parse(u)
	if url != nil && url.Host != "" {
		return url.Host
	}
	return pi.URL
}

func (pi pageInfo) Best() string {
	if len(pi.Icons) > 0 {
		best := pi.Icons[0]
		return best.URL
	}
	return ""
}

func renderHTMLTemplate(w http.ResponseWriter, httpStatus int, templ *template.Template, data interface{}) {
	w.Header().Add(contentType, "text/html; charset=utf-8")
	w.WriteHeader(httpStatus)

	err := templ.Execute(w, data)
	if err != nil {
		err = fmt.Errorf("server: could not generate output: %s", err)
		logger.Print(err)
		w.Write([]byte(err.Error()))
	}
}

func startServer(port string) {
	registerGzipHandler("/", indexHandler)
	registerGzipHandler("/icons", iconsHandler)
	registerHandler("/icon", iconHandler)
	registerGzipHandler("/popular", popularHandler)
	registerGzipHandler("/allicons.json", alliconsHandler)
	registerHandler("/lettericons/", lettericonHandler)
	registerHandler("/api/icons", obsoleteAPIHandler)

	serveAsset("/pure-0.5.0-min.css", "besticon/iconserver/assets/pure-0.5.0-min.css", oneYear)
	serveAsset("/grids-responsive-0.5.0-min.css", "besticon/iconserver/assets/grids-responsive-0.5.0-min.css", oneYear)
	serveAsset("/main-min.css", "besticon/iconserver/assets/main-min.css", oneYear)

	serveAsset("/icon.svg", "besticon/iconserver/assets/icon.svg", oneYear)
	serveAsset("/favicon.ico", "besticon/iconserver/assets/favicon.ico", oneYear)
	serveAsset("/apple-touch-icon.png", "besticon/iconserver/assets/apple-touch-icon.png", oneYear)
	serveAsset("/robots.txt", "besticon/iconserver/assets/robots.txt", nocache)
	serveAsset("/test-lettericons", "besticon/iconserver/assets/test-lettericons.html", nocache)

	addr := "0.0.0.0:" + port
	logger.Print("Starting server on ", addr, "...")
	e := http.ListenAndServe(addr, newLoggingMux())
	if e != nil {
		logger.Fatalf("cannot start server: %s\n", e)
	}
}

const (
	oneYear = 365 * 24 * 3600
	nocache = -1
)

func serveAsset(path string, assetPath string, maxAgeSeconds int) {
	registerGzipHandler(path, func(w http.ResponseWriter, r *http.Request) {
		assetInfo, err := assets.AssetInfo(assetPath)
		if err != nil {
			panic(err)
		}

		if maxAgeSeconds != nocache {
			w.Header().Add(cacheControl, fmt.Sprintf("max-age=%d", maxAgeSeconds))
		}

		http.ServeContent(w, r, assetInfo.Name(), assetInfo.ModTime(),
			bytes.NewReader(assets.MustAsset(assetPath)))
	})
}

func registerHandler(path string, f http.HandlerFunc) {
	http.Handle(path, newExpvarHandler(path, f))
}

func registerGzipHandler(path string, f http.HandlerFunc) {
	http.Handle(path, gziphandler.GzipHandler(newExpvarHandler(path, f)))
}

func main() {
	fmt.Printf("iconserver %s (%s) - http://icons.better-idea.org\n", besticon.VersionString, runtime.Version())
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}
	startServer(port)
}

func init() {
	indexHTML = templateFromAsset("besticon/iconserver/assets/index.html", "index.html")
	iconsHTML = templateFromAsset("besticon/iconserver/assets/icons.html", "icons.html")
	popularHTML = templateFromAsset("besticon/iconserver/assets/popular.html", "popular.html")
	notFoundHTML = templateFromAsset("besticon/iconserver/assets/not_found.html", "not_found.html")
}

func templateFromAsset(assetPath, templateName string) *template.Template {
	bytes := assets.MustAsset(assetPath)
	return template.Must(template.New(templateName).Funcs(funcMap).Parse(string(bytes)))
}

var indexHTML *template.Template
var iconsHTML *template.Template
var popularHTML *template.Template
var notFoundHTML *template.Template

var funcMap = template.FuncMap{
	"ImgWidth": imgWidth,
}

func imgWidth(i *besticon.Icon) int {
	return i.Width / 2.0
}

func init() {
	besticon.SetCacheMaxSize(32)

	expvar.Publish("cacheBytes", expvar.Func(func() interface{} { return besticon.GetCacheStats().Bytes }))
	expvar.Publish("cacheItems", expvar.Func(func() interface{} { return besticon.GetCacheStats().Items }))
	expvar.Publish("cacheGets", expvar.Func(func() interface{} { return besticon.GetCacheStats().Gets }))
	expvar.Publish("cacheHits", expvar.Func(func() interface{} { return besticon.GetCacheStats().Hits }))
	expvar.Publish("cacheEvictions", expvar.Func(func() interface{} { return besticon.GetCacheStats().Evictions }))
}
<|endoftext|>


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

package pretty

import (
	"encoding"
	"fmt"
	"reflect"
	"sort"
)

func isZeroVal(val reflect.Value) bool {
	if !val.CanInterface() {
		return false
	}
	z := reflect.Zero(val.Type()).Interface()
	return reflect.DeepEqual(val.Interface(), z)
}

func (c *Config) val2node(val reflect.Value) node {
	// TODO(kevlar): pointer tracking?

	if val.CanInterface() {
		v := val.Interface()
		if s, ok := v.(fmt.Stringer); ok && c.PrintStringers {
			return stringVal(s.String())
		}
		if t, ok := v.(encoding.TextMarshaler); ok && !c.NoTextMarshalers {
			if raw, err := t.MarshalText(); err == nil { // if NOT an error
				return stringVal(string(raw))
			}
		}
	}

	switch kind := val.Kind(); kind {
	case reflect.Ptr, reflect.Interface:
		if val.IsNil() {
			return rawVal("nil")
		}
		return c.val2node(val.Elem())
	case reflect.String:
		return stringVal(val.String())
	case reflect.Slice, reflect.Array:
		n := list{}
		length := val.Len()
		for i := 0; i < length; i++ {
			n = append(n, c.val2node(val.Index(i)))
		}
		return n
	case reflect.Map:
		n := keyvals{}
		keys := val.MapKeys()
		for _, key := range keys {
			// TODO(kevlar): Support arbitrary type keys?
			n = append(n, keyval{compactString(c.val2node(key)), c.val2node(val.MapIndex(key))})
		}
		sort.Sort(n)
		return n
	case reflect.Struct:
		n := keyvals{}
		typ := val.Type()
		fields := typ.NumField()
		for i := 0; i < fields; i++ {
			sf := typ.Field(i)
			if !c.IncludeUnexported && sf.PkgPath != "" {
				continue
			}
			field := val.Field(i)
			if c.SkipZeroFields && isZeroVal(field) {
				continue
			}
			n = append(n, keyval{sf.Name, c.val2node(field)})
		}
		return n
	case reflect.Bool:
		if val.Bool() {
			return rawVal("true")
		}
		return rawVal("false")
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return rawVal(fmt.Sprintf("%d", val.Int()))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return rawVal(fmt.Sprintf("%d", val.Uint()))
	case reflect.Uintptr:
		return rawVal(fmt.Sprintf("0x%X", val.Uint()))
	case reflect.Float32, reflect.Float64:
		return rawVal(fmt.Sprintf("%v", val.Float()))
	case reflect.Complex64, reflect.Complex128:
		return rawVal(fmt.Sprintf("%v", val.Complex()))
	}

	// Fall back to the default %#v if we can
	if val.CanInterface() {
		return rawVal(fmt.Sprintf("%#v", val.Interface()))
	}

	return rawVal(val.String())
}
<commit_msg>Correctly handle bare nil values<commit_after>// Copyright 2013 Google Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pretty

import (
	"encoding"
	"fmt"
	"reflect"
	"sort"
)

func isZeroVal(val reflect.Value) bool {
	if !val.CanInterface() {
		return false
	}
	z := reflect.Zero(val.Type()).Interface()
	return reflect.DeepEqual(val.Interface(), z)
}

func (c *Config) val2node(val reflect.Value) node {
	// TODO(kevlar): pointer tracking?

	if !val.IsValid() {
		return rawVal("nil")
	}

	if val.CanInterface() {
		v := val.Interface()
		if s, ok := v.(fmt.Stringer); ok && c.PrintStringers {
			return stringVal(s.String())
		}
		if t, ok := v.(encoding.TextMarshaler); ok && !c.NoTextMarshalers {
			if raw, err := t.MarshalText(); err == nil { // if NOT an error
				return stringVal(string(raw))
			}
		}
	}

	switch kind := val.Kind(); kind {
	case reflect.Ptr, reflect.Interface:
		if val.IsNil() {
			return rawVal("nil")
		}
		return c.val2node(val.Elem())
	case reflect.String:
		return stringVal(val.String())
	case reflect.Slice, reflect.Array:
		n := list{}
		length := val.Len()
		for i := 0; i < length; i++ {
			n = append(n, c.val2node(val.Index(i)))
		}
		return n
	case reflect.Map:
		n := keyvals{}
		keys := val.MapKeys()
		for _, key := range keys {
			// TODO(kevlar): Support arbitrary type keys?
			n = append(n, keyval{compactString(c.val2node(key)), c.val2node(val.MapIndex(key))})
		}
		sort.Sort(n)
		return n
	case reflect.Struct:
		n := keyvals{}
		typ := val.Type()
		fields := typ.NumField()
		for i := 0; i < fields; i++ {
			sf := typ.Field(i)
			if !c.IncludeUnexported && sf.PkgPath != "" {
				continue
			}
			field := val.Field(i)
			if c.SkipZeroFields && isZeroVal(field) {
				continue
			}
			n = append(n, keyval{sf.Name, c.val2node(field)})
		}
		return n
	case reflect.Bool:
		if val.Bool() {
			return rawVal("true")
		}
		return rawVal("false")
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return rawVal(fmt.Sprintf("%d", val.Int()))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return rawVal(fmt.Sprintf("%d", val.Uint()))
	case reflect.Uintptr:
		return rawVal(fmt.Sprintf("0x%X", val.Uint()))
	case reflect.Float32, reflect.Float64:
		return rawVal(fmt.Sprintf("%v", val.Float()))
	case reflect.Complex64, reflect.Complex128:
		return rawVal(fmt.Sprintf("%v", val.Complex()))
	}

	// Fall back to the default %#v if we can
	if val.CanInterface() {
		return rawVal(fmt.Sprintf("%#v", val.Interface()))
	}

	return rawVal(val.String())
}
<|endoftext|>


<commit_before>package proj

import (
	"fmt"
	"testing"
)

func TestMercToLatlong(t *testing.T) {
	merc, err := NewProj("+proj=merc +ellps=clrk66 +lat_ts=33")
	defer merc.Close()
	if err != nil {
		t.Error(err)
	}

	ll, err := NewProj("+proj=latlong +datum=WGS84")
	defer ll.Close()
	if err != nil {
		t.Error(err)
	}

	x, y, err := Transform2(ll, merc, DegToRad(-16), DegToRad(20.25))
	if err != nil {
		t.Error(err)
	} else {
		s := fmt.Sprintf("%.2f %.2f", x, y)
		s1 := "-1495284.21 1920596.79"
		if s != s1 {
			t.Errorf("MercToLatlong = %v, want %v", s, s1)
		}
	}
}

func TestInvalidErrorProblem(t *testing.T) {
	merc, err := NewProj("+proj=merc +ellps=clrk66 +lat_ts=33")
	defer merc.Close()
	if err != nil {
		t.Error(err)
	}

	ll, err := NewProj("+proj=latlong +datum=WGS84")
	defer ll.Close()
	if err != nil {
		t.Error(err)
	}

	_, _, err = Transform2(ll, merc, DegToRad(3000), DegToRad(500))
	if err == nil {
		t.Error("err should not be nil")
	}

	// Try create a new projection after an error
	merc2, err := NewProj("+proj=merc +ellps=clrk66 +lat_ts=33")
	defer merc2.Close()
	if err != nil {
		t.Error(err)
	}
}
<commit_msg>Revert "Fix to make tests succeed with proj-4.9.1"<commit_after>package proj

import (
	"fmt"
	"testing"
)

func TestMercToLatlong(t *testing.T) {
	merc, err := NewProj("+proj=merc +ellps=clrk66 +lat_ts=33")
	defer merc.Close()
	if err != nil {
		t.Error(err)
	}

	ll, err := NewProj("+proj=latlong")
	defer ll.Close()
	if err != nil {
		t.Error(err)
	}

	x, y, err := Transform2(ll, merc, DegToRad(-16), DegToRad(20.25))
	if err != nil {
		t.Error(err)
	} else {
		s := fmt.Sprintf("%.2f %.2f", x, y)
		s1 := "-1495284.21 1920596.79"
		if s != s1 {
			t.Errorf("MercToLatlong = %v, want %v", s, s1)
		}
	}
}

func TestInvalidErrorProblem(t *testing.T) {
	merc, err := NewProj("+proj=merc +ellps=clrk66 +lat_ts=33")
	defer merc.Close()
	if err != nil {
		t.Error(err)
	}

	ll, err := NewProj("+proj=latlong")
	defer ll.Close()
	if err != nil {
		t.Error(err)
	}

	_, _, err = Transform2(ll, merc, DegToRad(3000), DegToRad(500))
	if err == nil {
		t.Error("err should not be nil")
	}

	// Try create a new projection after an error
	merc2, err := NewProj("+proj=merc +ellps=clrk66 +lat_ts=33")
	defer merc2.Close()
	if err != nil {
		t.Error(err)
	}
}
<|endoftext|>


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

package monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get container information"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory containing configs for each container"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 110,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) > 1 && m[1] == s.Service {
		return msg
	}
	return s.errMsg(nil)
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) verifyContactGroup(group string) bool {
	output, err := exec.Command("/usr/bin/cmk_admin", "-l").Output()
	if err != nil {
		fmt.Printf("%d %s - Error listing existing contact_groups for validation, please try again later! Error: %s\n", Critical, config.CheckName, err.Error())
		return false
	}
	for _, l := range strings.Split(string(output), "\n") {
		cg := strings.TrimSpace(strings.TrimPrefix(l, "*"))
		if cg == c.ContactGroup {
			return true
		}
	}
	return false
}

func (c *ContainerCheck) getContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, config.CheckName, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, config.CheckName, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, config.CheckName)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, config.CheckName)
			return
		}
		group, ok := val.(string)
		if ok {
			group = strings.ToLower(group)
			if c.verifyContactGroup(group) {
				c.ContactGroup = group
			} else {
				fmt.Printf("%d %s - Specified contact_group does not exist in cmk! Falling back to default group %s.\n", Critical, config.CheckName, config.DefaultGroup)
			}
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, config.CheckName)
		}
	}
}

func (c *ContainerCheck) setContactGroup(name string) {
	if len(c.ContactGroup) == 0 {
		c.getContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, config.CheckName, name, err.Error())
		} else {
			os.Create(inventoryPath)
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	c.setContactGroup(c.Name)
	defer func() { done <- true }()
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		c.setContactGroup(serviceName)
		go c.serviceCheck(s).checkWithTimeout(results, t)
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<commit_msg>fix contact group function names<commit_after>/* Copyright 2014 Ooyala, Inc. All rights reserved.
 *
 * This file is licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
 * except in compliance with the License. You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is
 * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and limitations under the License.
 */

package monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get container information"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory containing configs for each container"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 110,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) > 1 && m[1] == s.Service {
		return msg
	}
	return s.errMsg(nil)
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) verifyContactGroup(group string) bool {
	output, err := exec.Command("/usr/bin/cmk_admin", "-l").Output()
	if err != nil {
		fmt.Printf("%d %s - Error listing existing contact_groups for validation, please try again later! Error: %s\n", Critical, config.CheckName, err.Error())
		return false
	}
	for _, l := range strings.Split(string(output), "\n") {
		cg := strings.TrimSpace(strings.TrimPrefix(l, "*"))
		if cg == c.ContactGroup {
			return true
		}
	}
	return false
}

func (c *ContainerCheck) parseContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, config.CheckName, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, config.CheckName, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, config.CheckName)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, config.CheckName)
			return
		}
		group, ok := val.(string)
		if ok {
			group = strings.ToLower(group)
			if c.verifyContactGroup(group) {
				c.ContactGroup = group
			} else {
				fmt.Printf("%d %s - Specified contact_group does not exist in cmk! Falling back to default group %s.\n", Critical, config.CheckName, config.DefaultGroup)
			}
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, config.CheckName)
		}
	}
}

func (c *ContainerCheck) updateContactGroup(name string) {
	if len(c.ContactGroup) == 0 {
		c.parseContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, config.CheckName, name, err.Error())
		} else {
			os.Create(inventoryPath)
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	c.updateContactGroup(c.Name)
	defer func() { done <- true }()
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		c.updateContactGroup(serviceName)
		go c.serviceCheck(s).checkWithTimeout(results, t)
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<|endoftext|>