<commit_before>package library

import (
	"github.com/anaminus/rbxmk"
	lua "github.com/yuin/gopher-lua"
)

func init() { register(Base, -1) }

var Base = rbxmk.Library{
	Name: "",
	Open: func(s rbxmk.State) *lua.LTable {
		openFilteredLibs(s.L, filteredStdLib)
		return nil
	},
}

type libFilter struct {
	Name     string
	OpenFunc lua.LGFunction
	Filter   map[lua.LValue]bool
}

var filteredStdLib = []libFilter{
	{lua.BaseLibName, lua.OpenBase, map[lua.LValue]bool{
		lua.LString("_G"):       true,
		lua.LString("_VERSION"): true,
		lua.LString("assert"):   true,
		lua.LString("error"):    true,
		lua.LString("ipairs"):   true,
		lua.LString("next"):     true,
		lua.LString("pairs"):    true,
		lua.LString("pcall"):    true,
		lua.LString("print"):    true,
		lua.LString("select"):   true,
		lua.LString("tonumber"): true,
		lua.LString("tostring"): true,
		lua.LString("type"):     true,
		lua.LString("unpack"):   true,
		lua.LString("xpcall"):   true,
		// lua.LString("collectgarbage"): true,
		// lua.LString("dofile"):         true,
		// lua.LString("getfenv"):        true,
		// lua.LString("getmetatable"):   true,
		// lua.LString("load"):           true,
		// lua.LString("loadfile"):       true,
		// lua.LString("loadstring"):     true,
		// lua.LString("module"):         true,
		// lua.LString("rawequal"):       true,
		// lua.LString("rawget"):         true,
		// lua.LString("rawset"):         true,
		// lua.LString("require"):        true,
		// lua.LString("setfenv"):        true,
		// lua.LString("setmetatable"):   true,
	}},
	// {lua.CoroutineLibName, lua.OpenCoroutine, map[lua.LValue]bool{
	// 	lua.LString("create"):  true,
	// 	lua.LString("resume"):  true,
	// 	lua.LString("running"): true,
	// 	lua.LString("status"):  true,
	// 	lua.LString("wrap"):    true,
	// 	lua.LString("yield"):   true,
	// }},
	// {lua.DebugLibName, lua.OpenDebug, map[lua.LValue]bool{
	// 	lua.LString("debug"):        true,
	// 	lua.LString("getfenv"):      true,
	// 	lua.LString("gethook"):      true,
	// 	lua.LString("getinfo"):      true,
	// 	lua.LString("getlocal"):     true,
	// 	lua.LString("getmetatable"): true,
	// 	lua.LString("getregistry"):  true,
	// 	lua.LString("getupvalue"):   true,
	// 	lua.LString("setfenv"):      true,
	// 	lua.LString("sethook"):      true,
	// 	lua.LString("setlocal"):     true,
	// 	lua.LString("setmetatable"): true,
	// 	lua.LString("setupvalue"):   true,
	// 	lua.LString("traceback"):    true,
	// }},
	// {lua.IoLibName, lua.OpenIo, map[lua.LValue]bool{
	// 	lua.LString("close"):   true,
	// 	lua.LString("flush"):   true,
	// 	lua.LString("input"):   true,
	// 	lua.LString("lines"):   true,
	// 	lua.LString("open"):    true,
	// 	lua.LString("output"):  true,
	// 	lua.LString("popen"):   true,
	// 	lua.LString("read"):    true,
	// 	lua.LString("stderr"):  true,
	// 	lua.LString("stdin"):   true,
	// 	lua.LString("stdout"):  true,
	// 	lua.LString("tmpfile"): true,
	// 	lua.LString("type"):    true,
	// 	lua.LString("write"):   true,
	// }},
	{lua.MathLibName, lua.OpenMath, map[lua.LValue]bool{
		lua.LString("abs"):        true,
		lua.LString("acos"):       true,
		lua.LString("asin"):       true,
		lua.LString("atan"):       true,
		lua.LString("atan2"):      true,
		lua.LString("ceil"):       true,
		lua.LString("cos"):        true,
		lua.LString("cosh"):       true,
		lua.LString("deg"):        true,
		lua.LString("exp"):        true,
		lua.LString("floor"):      true,
		lua.LString("fmod"):       true,
		lua.LString("frexp"):      true,
		lua.LString("huge"):       true,
		lua.LString("ldexp"):      true,
		lua.LString("log"):        true,
		lua.LString("log10"):      true,
		lua.LString("max"):        true,
		lua.LString("min"):        true,
		lua.LString("modf"):       true,
		lua.LString("pi"):         true,
		lua.LString("pow"):        true,
		lua.LString("rad"):        true,
		lua.LString("random"):     true,
		lua.LString("randomseed"): true,
		lua.LString("sin"):        true,
		lua.LString("sinh"):       true,
		lua.LString("sqrt"):       true,
		lua.LString("tan"):        true,
		lua.LString("tanh"):       true,
	}},
	{lua.OsLibName, lua.OpenOs, map[lua.LValue]bool{
		lua.LString("clock"):    true,
		lua.LString("date"):     true,
		lua.LString("difftime"): true,
		lua.LString("time"):     true,
		// lua.LString("execute"):   true,
		// lua.LString("exit"):      true,
		// lua.LString("getenv"):    true,
		// lua.LString("remove"):    true,
		// lua.LString("rename"):    true,
		// lua.LString("setlocale"): true,
		// lua.LString("tmpname"):   true,
	}},
	// {lua.LoadLibName, lua.OpenPackage, map[lua.LValue]bool{
	// 	lua.LString("cpath"):   true,
	// 	lua.LString("loaded"):  true,
	// 	lua.LString("loaders"): true,
	// 	lua.LString("loadlib"): true,
	// 	lua.LString("path"):    true,
	// 	lua.LString("preload"): true,
	// 	lua.LString("seeall"):  true,
	// }},
	{lua.StringLibName, lua.OpenString, map[lua.LValue]bool{
		lua.LString("byte"):    true,
		lua.LString("char"):    true,
		lua.LString("find"):    true,
		lua.LString("format"):  true,
		lua.LString("gmatch"):  true,
		lua.LString("gsub"):    true,
		lua.LString("len"):     true,
		lua.LString("lower"):   true,
		lua.LString("match"):   true,
		lua.LString("rep"):     true,
		lua.LString("reverse"): true,
		lua.LString("sub"):     true,
		lua.LString("upper"):   true,
		// lua.LString("dump"): true,
	}},
	{lua.TabLibName, lua.OpenTable, map[lua.LValue]bool{
		lua.LString("concat"): true,
		lua.LString("insert"): true,
		lua.LString("maxn"):   true,
		lua.LString("remove"): true,
		lua.LString("sort"):   true,
	}},
	// {lua.ChannelLibName, lua.OpenChannel, map[lua.LValue]bool{
	// 	lua.LString("make"):   true,
	// 	lua.LString("select"): true,
	// }},
}

func openFilteredLibs(l *lua.LState, libs []libFilter, upvalues ...lua.LValue) {
	for _, lib := range libs {
		l.Push(l.NewClosure(lib.OpenFunc, upvalues...))
		// LState.OpenLibs passes the library name as an argument for whatever
		// reason.
		l.Push(lua.LString(lib.Name))

		if lib.Filter == nil {
			l.Call(1, 0)
			continue
		}
		l.Call(1, 1)
		table := l.CheckTable(1)
		l.Pop(1)
		for k, _ := table.Next(lua.LNil); k != lua.LNil; k, _ = table.Next(k) {
			if !lib.Filter[k] {
				table.RawSet(k, lua.LNil)
			}
		}
	}
}
<commit_msg>Include getmetatable and setmetatable.<commit_after>package library

import (
	"github.com/anaminus/rbxmk"
	lua "github.com/yuin/gopher-lua"
)

func init() { register(Base, -1) }

var Base = rbxmk.Library{
	Name: "",
	Open: func(s rbxmk.State) *lua.LTable {
		openFilteredLibs(s.L, filteredStdLib)
		return nil
	},
}

type libFilter struct {
	Name     string
	OpenFunc lua.LGFunction
	Filter   map[lua.LValue]bool
}

var filteredStdLib = []libFilter{
	{lua.BaseLibName, lua.OpenBase, map[lua.LValue]bool{
		lua.LString("_G"):       true,
		lua.LString("_VERSION"): true,
		lua.LString("assert"):   true,
		lua.LString("error"):    true,
		lua.LString("ipairs"):   true,
		lua.LString("next"):     true,
		lua.LString("pairs"):    true,
		lua.LString("pcall"):    true,
		lua.LString("print"):    true,
		lua.LString("select"):   true,
		lua.LString("tonumber"): true,
		lua.LString("tostring"): true,
		lua.LString("type"):     true,
		lua.LString("unpack"):   true,
		lua.LString("xpcall"):   true,
		// lua.LString("collectgarbage"): true,
		// lua.LString("dofile"):         true,
		// lua.LString("getfenv"):        true,
		lua.LString("getmetatable"): true,
		// lua.LString("load"):           true,
		// lua.LString("loadfile"):       true,
		// lua.LString("loadstring"):     true,
		// lua.LString("module"):         true,
		// lua.LString("rawequal"):       true,
		// lua.LString("rawget"):         true,
		// lua.LString("rawset"):         true,
		// lua.LString("require"):        true,
		// lua.LString("setfenv"):        true,
		lua.LString("setmetatable"): true,
	}},
	// {lua.CoroutineLibName, lua.OpenCoroutine, map[lua.LValue]bool{
	// 	lua.LString("create"):  true,
	// 	lua.LString("resume"):  true,
	// 	lua.LString("running"): true,
	// 	lua.LString("status"):  true,
	// 	lua.LString("wrap"):    true,
	// 	lua.LString("yield"):   true,
	// }},
	// {lua.DebugLibName, lua.OpenDebug, map[lua.LValue]bool{
	// 	lua.LString("debug"):        true,
	// 	lua.LString("getfenv"):      true,
	// 	lua.LString("gethook"):      true,
	// 	lua.LString("getinfo"):      true,
	// 	lua.LString("getlocal"):     true,
	// 	lua.LString("getmetatable"): true,
	// 	lua.LString("getregistry"):  true,
	// 	lua.LString("getupvalue"):   true,
	// 	lua.LString("setfenv"):      true,
	// 	lua.LString("sethook"):      true,
	// 	lua.LString("setlocal"):     true,
	// 	lua.LString("setmetatable"): true,
	// 	lua.LString("setupvalue"):   true,
	// 	lua.LString("traceback"):    true,
	// }},
	// {lua.IoLibName, lua.OpenIo, map[lua.LValue]bool{
	// 	lua.LString("close"):   true,
	// 	lua.LString("flush"):   true,
	// 	lua.LString("input"):   true,
	// 	lua.LString("lines"):   true,
	// 	lua.LString("open"):    true,
	// 	lua.LString("output"):  true,
	// 	lua.LString("popen"):   true,
	// 	lua.LString("read"):    true,
	// 	lua.LString("stderr"):  true,
	// 	lua.LString("stdin"):   true,
	// 	lua.LString("stdout"):  true,
	// 	lua.LString("tmpfile"): true,
	// 	lua.LString("type"):    true,
	// 	lua.LString("write"):   true,
	// }},
	{lua.MathLibName, lua.OpenMath, map[lua.LValue]bool{
		lua.LString("abs"):        true,
		lua.LString("acos"):       true,
		lua.LString("asin"):       true,
		lua.LString("atan"):       true,
		lua.LString("atan2"):      true,
		lua.LString("ceil"):       true,
		lua.LString("cos"):        true,
		lua.LString("cosh"):       true,
		lua.LString("deg"):        true,
		lua.LString("exp"):        true,
		lua.LString("floor"):      true,
		lua.LString("fmod"):       true,
		lua.LString("frexp"):      true,
		lua.LString("huge"):       true,
		lua.LString("ldexp"):      true,
		lua.LString("log"):        true,
		lua.LString("log10"):      true,
		lua.LString("max"):        true,
		lua.LString("min"):        true,
		lua.LString("modf"):       true,
		lua.LString("pi"):         true,
		lua.LString("pow"):        true,
		lua.LString("rad"):        true,
		lua.LString("random"):     true,
		lua.LString("randomseed"): true,
		lua.LString("sin"):        true,
		lua.LString("sinh"):       true,
		lua.LString("sqrt"):       true,
		lua.LString("tan"):        true,
		lua.LString("tanh"):       true,
	}},
	{lua.OsLibName, lua.OpenOs, map[lua.LValue]bool{
		lua.LString("clock"):    true,
		lua.LString("date"):     true,
		lua.LString("difftime"): true,
		lua.LString("time"):     true,
		// lua.LString("execute"):   true,
		// lua.LString("exit"):      true,
		// lua.LString("getenv"):    true,
		// lua.LString("remove"):    true,
		// lua.LString("rename"):    true,
		// lua.LString("setlocale"): true,
		// lua.LString("tmpname"):   true,
	}},
	// {lua.LoadLibName, lua.OpenPackage, map[lua.LValue]bool{
	// 	lua.LString("cpath"):   true,
	// 	lua.LString("loaded"):  true,
	// 	lua.LString("loaders"): true,
	// 	lua.LString("loadlib"): true,
	// 	lua.LString("path"):    true,
	// 	lua.LString("preload"): true,
	// 	lua.LString("seeall"):  true,
	// }},
	{lua.StringLibName, lua.OpenString, map[lua.LValue]bool{
		lua.LString("byte"):    true,
		lua.LString("char"):    true,
		lua.LString("find"):    true,
		lua.LString("format"):  true,
		lua.LString("gmatch"):  true,
		lua.LString("gsub"):    true,
		lua.LString("len"):     true,
		lua.LString("lower"):   true,
		lua.LString("match"):   true,
		lua.LString("rep"):     true,
		lua.LString("reverse"): true,
		lua.LString("sub"):     true,
		lua.LString("upper"):   true,
		// lua.LString("dump"): true,
	}},
	{lua.TabLibName, lua.OpenTable, map[lua.LValue]bool{
		lua.LString("concat"): true,
		lua.LString("insert"): true,
		lua.LString("maxn"):   true,
		lua.LString("remove"): true,
		lua.LString("sort"):   true,
	}},
	// {lua.ChannelLibName, lua.OpenChannel, map[lua.LValue]bool{
	// 	lua.LString("make"):   true,
	// 	lua.LString("select"): true,
	// }},
}

func openFilteredLibs(l *lua.LState, libs []libFilter, upvalues ...lua.LValue) {
	for _, lib := range libs {
		l.Push(l.NewClosure(lib.OpenFunc, upvalues...))
		// LState.OpenLibs passes the library name as an argument for whatever
		// reason.
		l.Push(lua.LString(lib.Name))

		if lib.Filter == nil {
			l.Call(1, 0)
			continue
		}
		l.Call(1, 1)
		table := l.CheckTable(1)
		l.Pop(1)
		for k, _ := table.Next(lua.LNil); k != lua.LNil; k, _ = table.Next(k) {
			if !lib.Filter[k] {
				table.RawSet(k, lua.LNil)
			}
		}
	}
}
<|endoftext|>


<commit_before>package bot

import (
	"github.com/bwmarrin/discordgo"
	"go.uber.org/zap"
	"fmt"
	"github.com/FederationOfFathers/dashboard/messaging"
)

type DiscordAPI struct {
	Config DiscordCfg
	discord *discordgo.Session
	assignmentMsgs map[string]map[string]string
}

type EmojiRole struct {
	EmojiId string `yaml:"emojiId"`
	RoleId string `yaml:"roleId"`
}

type DiscordEmojiRoleGroup struct {
	MessageTitle string `yaml:"messageTitle"`
	MessageBody string `yaml:"messageBody"`
	Roles []EmojiRole `yaml:"roles"`

}

type DiscordRoleCfg struct {
	ChannelId string `yaml:"channelId"`
	EmojiRoleGroups []DiscordEmojiRoleGroup `yaml:"emojiRoles"`

}

type DiscordCfg struct {
	ClientId string `yaml:"appClientId"`
	Token string `yaml:"botToken"`
	StreamChannelId string `yaml:"streamChannelId"`
	GuildId string `yaml:"guildId"`
	RoleCfg DiscordRoleCfg `yaml:"roleConfig"`
}

func NewDiscordAPI(cfg DiscordCfg) DiscordAPI {
	return DiscordAPI{
		Config: cfg,
	}
}

func (d DiscordAPI) roleAssignmentHandler(s *discordgo.Session, event *discordgo.MessageReactionAdd) {

	// skip if the event was from the bot/app
	if event.UserID == d.Config.ClientId {
		return
	}

	// only handle if the message is one we have configured
	if roles, ok := d.assignmentMsgs[event.MessageID]; ok {

		// Unicode emojis use the unicode character (name) as the id. Others use the name and integer as the id.
		emojiId := event.Emoji.Name
		if event.Emoji.ID != "" {
			emojiId = fmt.Sprintf(":%s:%s", event.Emoji.Name, event.Emoji.ID)
		}

		if roleId, ok := roles[emojiId]; ok{
			// get the user from the server/guild
			member, err := d.discord.GuildMember(d.Config.GuildId, event.UserID)
			if err != nil {
				Logger.Error("Unable to get member", zap.Error(err))
				return
			}

			userHadRole := false
			// if the member has the mapped role, remove it
			for _, id := range member.Roles {
				if id == roleId {
					d.removeRoleFromUser(event.UserID, roleId)
					userHadRole = true // true, even if an attempt was made and failed
				}
			}

			// if the member does not have the mapped role, add it
			if !userHadRole {
				d.addRoleToUser(event.UserID, roleId)
				d.discord.GuildMemberRoleAdd(d.Config.GuildId, event.UserID, roleId)
				Logger.Info("added role to user", zap.String("userId", event.UserID), zap.String("roleId", roleId))

			}
		}

		// remove the users reaction. If the add/remove failed, they can click it again to re-trigger
		d.discord.MessageReactionRemove(
			event.ChannelID,
			event.MessageID,
			emojiId,
			event.UserID,
		)
	}

}

// Needs to be called before any other API function work
func (d *DiscordAPI) Connect() {
	dg, err := discordgo.New("Bot " + d.Config.Token)
	if err != nil {
		Logger.Error("Unable to create discord connection", zap.Error(err))
		return
	}

	d.discord = dg
	dg.Open()
}

func (d *DiscordAPI) StartRoleHandlers(){
	d.clearRoleChannel()
	d.createRoleMessages()
	d.discord.AddHandler(d.roleAssignmentHandler)
	Logger.Info("Role assignment handler started")
}

// Needs to be called to disconnect from discord
func (d *DiscordAPI) Shutdown() {
	d.discord.Close()
}

func (d DiscordAPI) PostStreamMessage(sm messaging.StreamMessage) error {
	if d.discord == nil {
		return fmt.Errorf("discord API not connected")
	}
	if d.Config.StreamChannelId == "" {
		return fmt.Errorf("stream channel id not configured")
	}
	author := discordgo.MessageEmbedAuthor{
		Name: fmt.Sprintf("%s is live!", sm.Username),
	}
	thumbnail := discordgo.MessageEmbedThumbnail{
		URL: sm.UserLogo,
	}
	footer := discordgo.MessageEmbedFooter{
		Text: fmt.Sprintf("%s | %s", sm.Platform, sm.Timestamp),
		IconURL: sm.PlatformLogo,
	}
	messageEmbed := discordgo.MessageEmbed{
		Description: sm.URL,
		Color: sm.PlatformColorInt,
		URL: sm.URL,
		Author: &author,
		Thumbnail: &thumbnail,
		Footer: &footer,
		Fields: []*discordgo.MessageEmbedField{
			{
				Name: "Game",
				Value: fmt.Sprintf("%s - %s", sm.Game, sm.Description),
				Inline: false,
			},
		},
	}
	_, err := d.discord.ChannelMessageSendEmbed(d.Config.StreamChannelId, &messageEmbed)
	return err
}

// removes all messages entered by this bot in the channel. Uses the ClientID of the bot/app
func (d *DiscordAPI) clearRoleChannel() {
	channelId := d.Config.RoleCfg.ChannelId

	messages, err := d.discord.ChannelMessages(channelId,50,"","","")
	if err != nil {
		Logger.Error("Unable to access channel messages", zap.String("channelId", channelId), zap.Error(err))
		return
	}
	for _, message := range messages {
		if message.Author.ID == d.Config.ClientId {
			err := d.discord.ChannelMessageDelete(channelId, message.ID)
			if err == nil {
				Logger.Info("Removed bot message", zap.String("messageId", message.ID))
			} else {
				Logger.Error("Unable to remove message", zap.String("messageId", message.ID), zap.Error(err))
			}

		}
	}
}

// creates messages in the channel and adds the emojis
func (d *DiscordAPI) createRoleMessages() {
	d.assignmentMsgs = make(map[string]map[string]string)
	for _, group := range d.Config.RoleCfg.EmojiRoleGroups {
		messageEmbed := discordgo.MessageEmbed{
			Title: group.MessageTitle,
			Description: group.MessageBody,
			Color: 15581239,
		}
		// create the message
		message, err := d.discord.ChannelMessageSendEmbed(d.Config.RoleCfg.ChannelId, &messageEmbed)
		if err != nil {
			Logger.Error("Unable to create message", zap.Error(err))
			continue
		}
		Logger.Info("Added role message", zap.String("message", group.MessageTitle), zap.String("messageId", message.ID))

		// add the emojis
		emojiRoles := make(map[string]string)
		for _, role := range group.Roles {
			emojiRoles[role.EmojiId] = role.RoleId
			err := d.discord.MessageReactionAdd(d.Config.RoleCfg.ChannelId, message.ID, role.EmojiId)
			if err != nil {
				Logger.Error("Unable to add emoji to message",
					zap.String("emojiId", role.EmojiId),
					zap.String("messageId", message.ID),
					zap.Error(err),
				)
			}
		}

		d.assignmentMsgs[message.ID] = emojiRoles
	}
}

func (d DiscordAPI) addRoleToUser(userId string, roleId string) {
	err := d.discord.GuildMemberRoleAdd(d.Config.GuildId, userId, roleId)
	if err != nil {
		Logger.Error("could not add role",
			zap.String("userId", userId),
			zap.String("roleId", roleId),
			zap.Error(err))
	} else {
		Logger.Info("added role to user",
			zap.String("userId", userId),
			zap.String("roleId", roleId),
		)
	}
}

func (d DiscordAPI) removeRoleFromUser(userId string, roleId string) {
	err := d.discord.GuildMemberRoleRemove(d.Config.GuildId, userId, roleId)
	if err != nil {
		Logger.Error("could not remove role", zap.Error(err))
	} else {
		Logger.Info("removed role from user",
			zap.String("userId", userId),
			zap.String("roleId", roleId),
		)
	}
}<commit_msg>Logging all available roles<commit_after>package bot

import (
	"github.com/bwmarrin/discordgo"
	"go.uber.org/zap"
	"fmt"
	"github.com/FederationOfFathers/dashboard/messaging"
)

type DiscordAPI struct {
	Config DiscordCfg
	discord *discordgo.Session
	assignmentMsgs map[string]map[string]string
}

type EmojiRole struct {
	EmojiId string `yaml:"emojiId"`
	RoleId string `yaml:"roleId"`
}

type DiscordEmojiRoleGroup struct {
	MessageTitle string `yaml:"messageTitle"`
	MessageBody string `yaml:"messageBody"`
	Roles []EmojiRole `yaml:"roles"`

}

type DiscordRoleCfg struct {
	ChannelId string `yaml:"channelId"`
	EmojiRoleGroups []DiscordEmojiRoleGroup `yaml:"emojiRoles"`

}

type DiscordCfg struct {
	ClientId string `yaml:"appClientId"`
	Token string `yaml:"botToken"`
	StreamChannelId string `yaml:"streamChannelId"`
	GuildId string `yaml:"guildId"`
	RoleCfg DiscordRoleCfg `yaml:"roleConfig"`
}

func NewDiscordAPI(cfg DiscordCfg) DiscordAPI {
	return DiscordAPI{
		Config: cfg,
	}
}

func (d DiscordAPI) roleAssignmentHandler(s *discordgo.Session, event *discordgo.MessageReactionAdd) {

	// skip if the event was from the bot/app
	if event.UserID == d.Config.ClientId {
		return
	}

	// only handle if the message is one we have configured
	if roles, ok := d.assignmentMsgs[event.MessageID]; ok {

		// Unicode emojis use the unicode character (name) as the id. Others use the name and integer as the id.
		emojiId := event.Emoji.Name
		if event.Emoji.ID != "" {
			emojiId = fmt.Sprintf(":%s:%s", event.Emoji.Name, event.Emoji.ID)
		}

		if roleId, ok := roles[emojiId]; ok{
			// get the user from the server/guild
			member, err := d.discord.GuildMember(d.Config.GuildId, event.UserID)
			if err != nil {
				Logger.Error("Unable to get member", zap.Error(err))
				return
			}

			userHadRole := false
			// if the member has the mapped role, remove it
			for _, id := range member.Roles {
				if id == roleId {
					d.removeRoleFromUser(event.UserID, roleId)
					userHadRole = true // true, even if an attempt was made and failed
				}
			}

			// if the member does not have the mapped role, add it
			if !userHadRole {
				d.addRoleToUser(event.UserID, roleId)
				d.discord.GuildMemberRoleAdd(d.Config.GuildId, event.UserID, roleId)
				Logger.Info("added role to user", zap.String("userId", event.UserID), zap.String("roleId", roleId))

			}
		}

		// remove the users reaction. If the add/remove failed, they can click it again to re-trigger
		d.discord.MessageReactionRemove(
			event.ChannelID,
			event.MessageID,
			emojiId,
			event.UserID,
		)
	}

}

// Needs to be called before any other API function work
func (d *DiscordAPI) Connect() {
	dg, err := discordgo.New("Bot " + d.Config.Token)
	if err != nil {
		Logger.Error("Unable to create discord connection", zap.Error(err))
		return
	}

	d.discord = dg
	dg.Open()
}

func (d *DiscordAPI) StartRoleHandlers(){
	d.listRoles()
	d.clearRoleChannel()
	d.createRoleMessages()
	d.discord.AddHandler(d.roleAssignmentHandler)
	Logger.Info("Role assignment handler started")
}

// Needs to be called to disconnect from discord
func (d *DiscordAPI) Shutdown() {
	d.discord.Close()
}

func (d DiscordAPI) PostStreamMessage(sm messaging.StreamMessage) error {
	if d.discord == nil {
		return fmt.Errorf("discord API not connected")
	}
	if d.Config.StreamChannelId == "" {
		return fmt.Errorf("stream channel id not configured")
	}
	author := discordgo.MessageEmbedAuthor{
		Name: fmt.Sprintf("%s is live!", sm.Username),
	}
	thumbnail := discordgo.MessageEmbedThumbnail{
		URL: sm.UserLogo,
	}
	footer := discordgo.MessageEmbedFooter{
		Text: fmt.Sprintf("%s | %s", sm.Platform, sm.Timestamp),
		IconURL: sm.PlatformLogo,
	}
	messageEmbed := discordgo.MessageEmbed{
		Description: sm.URL,
		Color: sm.PlatformColorInt,
		URL: sm.URL,
		Author: &author,
		Thumbnail: &thumbnail,
		Footer: &footer,
		Fields: []*discordgo.MessageEmbedField{
			{
				Name: "Game",
				Value: fmt.Sprintf("%s - %s", sm.Game, sm.Description),
				Inline: false,
			},
		},
	}
	_, err := d.discord.ChannelMessageSendEmbed(d.Config.StreamChannelId, &messageEmbed)
	return err
}

// removes all messages entered by this bot in the channel. Uses the ClientID of the bot/app
func (d *DiscordAPI) clearRoleChannel() {
	channelId := d.Config.RoleCfg.ChannelId

	messages, err := d.discord.ChannelMessages(channelId,50,"","","")
	if err != nil {
		Logger.Error("Unable to access channel messages", zap.String("channelId", channelId), zap.Error(err))
		return
	}
	for _, message := range messages {
		if message.Author.ID == d.Config.ClientId {
			err := d.discord.ChannelMessageDelete(channelId, message.ID)
			if err == nil {
				Logger.Info("Removed bot message", zap.String("messageId", message.ID))
			} else {
				Logger.Error("Unable to remove message", zap.String("messageId", message.ID), zap.Error(err))
			}

		}
	}
}

// creates messages in the channel and adds the emojis
func (d *DiscordAPI) createRoleMessages() {
	d.assignmentMsgs = make(map[string]map[string]string)
	for _, group := range d.Config.RoleCfg.EmojiRoleGroups {
		messageEmbed := discordgo.MessageEmbed{
			Title: group.MessageTitle,
			Description: group.MessageBody,
			Color: 15581239,
		}
		// create the message
		message, err := d.discord.ChannelMessageSendEmbed(d.Config.RoleCfg.ChannelId, &messageEmbed)
		if err != nil {
			Logger.Error("Unable to create message", zap.Error(err))
			continue
		}
		Logger.Info("Added role message", zap.String("message", group.MessageTitle), zap.String("messageId", message.ID))

		// add the emojis
		emojiRoles := make(map[string]string)
		for _, role := range group.Roles {
			emojiRoles[role.EmojiId] = role.RoleId
			err := d.discord.MessageReactionAdd(d.Config.RoleCfg.ChannelId, message.ID, role.EmojiId)
			if err != nil {
				Logger.Error("Unable to add emoji to message",
					zap.String("emojiId", role.EmojiId),
					zap.String("messageId", message.ID),
					zap.Error(err),
				)
			}
		}

		d.assignmentMsgs[message.ID] = emojiRoles
	}
}

func (d DiscordAPI) addRoleToUser(userId string, roleId string) {
	err := d.discord.GuildMemberRoleAdd(d.Config.GuildId, userId, roleId)
	if err != nil {
		Logger.Error("could not add role",
			zap.String("userId", userId),
			zap.String("roleId", roleId),
			zap.Error(err))
	} else {
		Logger.Info("added role to user",
			zap.String("userId", userId),
			zap.String("roleId", roleId),
		)
	}
}

func (d DiscordAPI) removeRoleFromUser(userId string, roleId string) {
	err := d.discord.GuildMemberRoleRemove(d.Config.GuildId, userId, roleId)
	if err != nil {
		Logger.Error("could not remove role", zap.Error(err))
	} else {
		Logger.Info("removed role from user",
			zap.String("userId", userId),
			zap.String("roleId", roleId),
		)
	}
}
func (d *DiscordAPI) listRoles() {
	roles, err := d.discord.GuildRoles(d.Config.GuildId)

	if err != nil {
		Logger.Error("Could not get roles", zap.Error(err))
		return
	}

	for i, role := range roles {
		Logger.Info(fmt.Sprintf("Role %d", i), zap.Any("role",role))
	}
}<|endoftext|>


<commit_before>package store_test

import (
	"bytes"
	"fmt"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/store"
	"launchpad.net/juju-core/testing"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

func (s *StoreSuite) dummyBranch(c *C, suffix string) bzrDir {
	tmpDir := c.MkDir()
	if suffix != "" {
		tmpDir = filepath.Join(tmpDir, suffix)
		err := os.MkdirAll(tmpDir, 0755)
		c.Assert(err, IsNil)
	}
	branch := bzrDir(tmpDir)
	branch.init()

	copyCharmDir(branch.path(), testing.Charms.Dir("dummy"))
	branch.add()
	branch.commit("Imported charm.")
	return branch
}

var urls = []*charm.URL{
	charm.MustParseURL("cs:~joe/oneiric/dummy"),
	charm.MustParseURL("cs:oneiric/dummy"),
}

type fakePlugin struct {
	oldEnv string
}

func (p *fakePlugin) install(dir string, content string) {
	p.oldEnv = os.Getenv("BZR_PLUGINS_AT")
	err := ioutil.WriteFile(filepath.Join(dir, "__init__.py"), []byte(content), 0644)
	if err != nil {
		panic(err)
	}
	os.Setenv("BZR_PLUGINS_AT", "fakePlugin@" + dir)
}

func (p *fakePlugin) uninstall() {
	os.Setenv("BZR_PLUGINS_AT", p.oldEnv)
}

func (s *StoreSuite) TestPublish(c *C) {
	// Ensure that the streams are parsed separately by inserting
	// garbage on stderr. The wanted information is still there.
	plugin := fakePlugin{}
	plugin.install(c.MkDir(), `import sys; sys.stderr.write("STDERR STUFF FROM TEST\n")`)
	defer plugin.uninstall()

	branch := s.dummyBranch(c, "")

	err := store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, IsNil)

	for _, url := range urls {
		info, rc, err := s.store.OpenCharm(url)
		c.Assert(err, IsNil)
		defer rc.Close()
		c.Assert(info.Revision(), Equals, 0)
		c.Assert(info.Meta().Name, Equals, "dummy")

		data, err := ioutil.ReadAll(rc)
		c.Assert(err, IsNil)

		bundle, err := charm.ReadBundleBytes(data)
		c.Assert(err, IsNil)
		c.Assert(bundle.Revision(), Equals, 0)
		c.Assert(bundle.Meta().Name, Equals, "dummy")
	}

	// Attempt to publish the same content again while providing the wrong
	// tip revision. It must pick the real revision from the branch and
	// note this was previously published.
	err = store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, Equals, store.ErrRedundantUpdate)

	// Bump the content revision and lie again about the known tip revision.
	// This time, though, pretend it's the same as the real branch revision
	// previously published. It must error and not publish the new revision
	// because it will use the revision provided as a parameter to check if
	// publishing was attempted before. This is the mechanism that enables
	// stopping fast without having to download every single branch. Real
	// revision is picked in the next scan.
	digest1 := branch.digest()
	branch.change()
	err = store.PublishBazaarBranch(s.store, urls, branch.path(), digest1)
	c.Assert(err, Equals, store.ErrRedundantUpdate)

	// Now allow it to publish the new content by providing an unseen revision.
	err = store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, IsNil)
	digest2 := branch.digest()

	info, err := s.store.CharmInfo(urls[0])
	c.Assert(err, IsNil)
	c.Assert(info.Revision(), Equals, 1)
	c.Assert(info.Meta().Name, Equals, "dummy")

	// There are two events published, for each of the successful attempts.
	// The failures are ignored given that they are artifacts of the
	// publishing mechanism rather than actual problems.
	_, err = s.store.CharmEvent(urls[0], "wrong-rev")
	c.Assert(err, Equals, store.ErrNotFound)
	for i, digest := range []string{digest1, digest2} {
		event, err := s.store.CharmEvent(urls[0], digest)
		c.Assert(err, IsNil)
		c.Assert(event.Kind, Equals, store.EventPublished)
		c.Assert(event.Revision, Equals, i)
		c.Assert(event.Errors, IsNil)
		c.Assert(event.Warnings, IsNil)
	}
}

func (s *StoreSuite) TestPublishErrorFromBzr(c *C) {
	// In TestPublish we ensure that the streams are parsed
	// separately by inserting garbage on stderr. Now make
	// sure that stderr isn't simply trashed, as we want to
	// know about what a real error tells us.
	plugin := fakePlugin{}
	plugin.install(c.MkDir(), `import sys; sys.stderr.write("STDERR STUFF FROM TEST\n"); sys.exit(1)`)
	defer plugin.uninstall()

	branch := s.dummyBranch(c, "")

	err := store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, ErrorMatches, "(?s).*STDERR STUFF.*")
}

func (s *StoreSuite) TestPublishErrorInCharm(c *C) {
	branch := s.dummyBranch(c, "")

	// Corrupt the charm.
	branch.remove("metadata.yaml")
	branch.commit("Removed metadata.yaml.")

	// Attempt to publish the erroneous content.
	err := store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, ErrorMatches, ".*/metadata.yaml: no such file or directory")

	// The event should be logged as well, since this was an error in the charm
	// that won't go away and must be communicated to the author.
	event, err := s.store.CharmEvent(urls[0], branch.digest())
	c.Assert(err, IsNil)
	c.Assert(event.Kind, Equals, store.EventPublishError)
	c.Assert(event.Revision, Equals, 0)
	c.Assert(event.Errors, NotNil)
	c.Assert(event.Errors[0], Matches, ".*/metadata.yaml: no such file or directory")
	c.Assert(event.Warnings, IsNil)
}

type bzrDir string

func (dir bzrDir) path(args ...string) string {
	return filepath.Join(append([]string{string(dir)}, args...)...)
}

func (dir bzrDir) run(args ...string) []byte {
	cmd := exec.Command("bzr", args...)
	cmd.Dir = string(dir)
	output, err := cmd.Output()
	if err != nil {
		panic(fmt.Sprintf("command failed: bzr %s\n%s", strings.Join(args, " "), output))
	}
	return output
}

func (dir bzrDir) init() {
	dir.run("init")
}

func (dir bzrDir) add(paths ...string) {
	dir.run(append([]string{"add"}, paths...)...)
}

func (dir bzrDir) remove(paths ...string) {
	dir.run(append([]string{"rm"}, paths...)...)
}

func (dir bzrDir) commit(msg string) {
	dir.run("commit", "-m", msg)
}

func (dir bzrDir) write(path string, data string) {
	err := ioutil.WriteFile(dir.path(path), []byte(data), 0644)
	if err != nil {
		panic(err)
	}
}

func (dir bzrDir) change() {
	t := time.Now().String()
	dir.write("timestamp", t)
	dir.add("timestamp")
	dir.commit("Revision bumped at " + t)
}

func (dir bzrDir) digest() string {
	output := dir.run("revision-info")
	f := bytes.Fields(output)
	if len(f) != 2 {
		panic("revision-info returned bad output: " + string(output))
	}
	return string(f[1])
}

func copyCharmDir(dst string, dir *charm.Dir) {
	var b bytes.Buffer
	err := dir.BundleTo(&b)
	if err != nil {
		panic(err)
	}
	bundle, err := charm.ReadBundleBytes(b.Bytes())
	if err != nil {
		panic(err)
	}
	err = bundle.ExpandTo(dst)
	if err != nil {
		panic(err)
	}
}
<commit_msg>store: fix tests<commit_after>package store_test

import (
	"bytes"
	"fmt"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/store"
	"launchpad.net/juju-core/testing"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

func (s *StoreSuite) dummyBranch(c *C, suffix string) bzrDir {
	tmpDir := c.MkDir()
	if suffix != "" {
		tmpDir = filepath.Join(tmpDir, suffix)
		err := os.MkdirAll(tmpDir, 0755)
		c.Assert(err, IsNil)
	}
	branch := bzrDir(tmpDir)
	branch.init()

	copyCharmDir(branch.path(), testing.Charms.Dir("dummy"))
	branch.add()
	branch.commit("Imported charm.")
	return branch
}

var urls = []*charm.URL{
	charm.MustParseURL("cs:~joe/oneiric/dummy"),
	charm.MustParseURL("cs:oneiric/dummy"),
}

type fakePlugin struct {
	oldEnv string
}

func (p *fakePlugin) install(dir string, content string) {
	p.oldEnv = os.Getenv("BZR_PLUGINS_AT")
	err := ioutil.WriteFile(filepath.Join(dir, "__init__.py"), []byte(content), 0644)
	if err != nil {
		panic(err)
	}
	os.Setenv("BZR_PLUGINS_AT", "fakePlugin@" + dir)
}

func (p *fakePlugin) uninstall() {
	os.Setenv("BZR_PLUGINS_AT", p.oldEnv)
}

func (s *StoreSuite) TestPublish(c *C) {
	branch := s.dummyBranch(c, "")

	// Ensure that the streams are parsed separately by inserting
	// garbage on stderr. The wanted information is still there.
	plugin := fakePlugin{}
	plugin.install(c.MkDir(), `import sys; sys.stderr.write("STDERR STUFF FROM TEST\n")`)
	defer plugin.uninstall()

	err := store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, IsNil)

	for _, url := range urls {
		info, rc, err := s.store.OpenCharm(url)
		c.Assert(err, IsNil)
		defer rc.Close()
		c.Assert(info.Revision(), Equals, 0)
		c.Assert(info.Meta().Name, Equals, "dummy")

		data, err := ioutil.ReadAll(rc)
		c.Assert(err, IsNil)

		bundle, err := charm.ReadBundleBytes(data)
		c.Assert(err, IsNil)
		c.Assert(bundle.Revision(), Equals, 0)
		c.Assert(bundle.Meta().Name, Equals, "dummy")
	}

	// Attempt to publish the same content again while providing the wrong
	// tip revision. It must pick the real revision from the branch and
	// note this was previously published.
	err = store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, Equals, store.ErrRedundantUpdate)

	// Bump the content revision and lie again about the known tip revision.
	// This time, though, pretend it's the same as the real branch revision
	// previously published. It must error and not publish the new revision
	// because it will use the revision provided as a parameter to check if
	// publishing was attempted before. This is the mechanism that enables
	// stopping fast without having to download every single branch. Real
	// revision is picked in the next scan.
	digest1 := branch.digest()
	branch.change()
	err = store.PublishBazaarBranch(s.store, urls, branch.path(), digest1)
	c.Assert(err, Equals, store.ErrRedundantUpdate)

	// Now allow it to publish the new content by providing an unseen revision.
	err = store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, IsNil)
	digest2 := branch.digest()

	info, err := s.store.CharmInfo(urls[0])
	c.Assert(err, IsNil)
	c.Assert(info.Revision(), Equals, 1)
	c.Assert(info.Meta().Name, Equals, "dummy")

	// There are two events published, for each of the successful attempts.
	// The failures are ignored given that they are artifacts of the
	// publishing mechanism rather than actual problems.
	_, err = s.store.CharmEvent(urls[0], "wrong-rev")
	c.Assert(err, Equals, store.ErrNotFound)
	for i, digest := range []string{digest1, digest2} {
		event, err := s.store.CharmEvent(urls[0], digest)
		c.Assert(err, IsNil)
		c.Assert(event.Kind, Equals, store.EventPublished)
		c.Assert(event.Revision, Equals, i)
		c.Assert(event.Errors, IsNil)
		c.Assert(event.Warnings, IsNil)
	}
}

func (s *StoreSuite) TestPublishErrorFromBzr(c *C) {
	branch := s.dummyBranch(c, "")

	// In TestPublish we ensure that the streams are parsed
	// separately by inserting garbage on stderr. Now make
	// sure that stderr isn't simply trashed, as we want to
	// know about what a real error tells us.
	plugin := fakePlugin{}
	plugin.install(c.MkDir(), `import sys; sys.stderr.write("STDERR STUFF FROM TEST\n"); sys.exit(1)`)
	defer plugin.uninstall()

	err := store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, ErrorMatches, "(?s).*STDERR STUFF.*")
}

func (s *StoreSuite) TestPublishErrorInCharm(c *C) {
	branch := s.dummyBranch(c, "")

	// Corrupt the charm.
	branch.remove("metadata.yaml")
	branch.commit("Removed metadata.yaml.")

	// Attempt to publish the erroneous content.
	err := store.PublishBazaarBranch(s.store, urls, branch.path(), "wrong-rev")
	c.Assert(err, ErrorMatches, ".*/metadata.yaml: no such file or directory")

	// The event should be logged as well, since this was an error in the charm
	// that won't go away and must be communicated to the author.
	event, err := s.store.CharmEvent(urls[0], branch.digest())
	c.Assert(err, IsNil)
	c.Assert(event.Kind, Equals, store.EventPublishError)
	c.Assert(event.Revision, Equals, 0)
	c.Assert(event.Errors, NotNil)
	c.Assert(event.Errors[0], Matches, ".*/metadata.yaml: no such file or directory")
	c.Assert(event.Warnings, IsNil)
}

type bzrDir string

func (dir bzrDir) path(args ...string) string {
	return filepath.Join(append([]string{string(dir)}, args...)...)
}

func (dir bzrDir) run(args ...string) []byte {
	cmd := exec.Command("bzr", args...)
	cmd.Dir = string(dir)
	output, err := cmd.Output()
	if err != nil {
		panic(fmt.Sprintf("command failed: bzr %s\n%s", strings.Join(args, " "), output))
	}
	return output
}

func (dir bzrDir) init() {
	dir.run("init")
}

func (dir bzrDir) add(paths ...string) {
	dir.run(append([]string{"add"}, paths...)...)
}

func (dir bzrDir) remove(paths ...string) {
	dir.run(append([]string{"rm"}, paths...)...)
}

func (dir bzrDir) commit(msg string) {
	dir.run("commit", "-m", msg)
}

func (dir bzrDir) write(path string, data string) {
	err := ioutil.WriteFile(dir.path(path), []byte(data), 0644)
	if err != nil {
		panic(err)
	}
}

func (dir bzrDir) change() {
	t := time.Now().String()
	dir.write("timestamp", t)
	dir.add("timestamp")
	dir.commit("Revision bumped at " + t)
}

func (dir bzrDir) digest() string {
	output := dir.run("revision-info")
	f := bytes.Fields(output)
	if len(f) != 2 {
		panic("revision-info returned bad output: " + string(output))
	}
	return string(f[1])
}

func copyCharmDir(dst string, dir *charm.Dir) {
	var b bytes.Buffer
	err := dir.BundleTo(&b)
	if err != nil {
		panic(err)
	}
	bundle, err := charm.ReadBundleBytes(b.Bytes())
	if err != nil {
		panic(err)
	}
	err = bundle.ExpandTo(dst)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libhttpserver

import (
	"net/http"
	"strings"
)

type contentTypeOverridingResponseWriter struct {
	original http.ResponseWriter
}

var _ http.ResponseWriter = (*contentTypeOverridingResponseWriter)(nil)

func newContentTypeOverridingResponseWriter(
	original http.ResponseWriter) *contentTypeOverridingResponseWriter {
	return &contentTypeOverridingResponseWriter{
		original: original,
	}
}

func (w *contentTypeOverridingResponseWriter) overrideMimeType(
	mimeType string) (newMimeType string) {
	// Send text/plain for all HTML and JS files to avoid them being executed
	// by the frontend WebView.
	ty := strings.ToLower(mimeType)
	switch {
	// First anything textual as text/plain.
	// Includes javascript, html, xml. (note that the type may be e.g. application/xhtml+xml)
	case strings.Contains(ty, "javascript") ||
		strings.Contains(ty, "xml") ||
		strings.Contains(ty, "html") ||
		strings.HasPrefix(ty, "text/"):
		return "text/plain"
	// Pass multimedia types through, and pdf too.
	case strings.HasPrefix(ty, "audio/") ||
		strings.HasPrefix(ty, "image/") ||
		strings.HasPrefix(ty, "video/") ||
		ty == "application/pdf":
		return ty
	// Otherwise default to binary data.
	default:
		return "application/octet-stream"
	}
}

func (w *contentTypeOverridingResponseWriter) override() {
	t := w.original.Header().Get("Content-Type")
	if len(t) > 0 {
		w.original.Header().Set("Content-Type", w.overrideMimeType(t))
	}
	w.original.Header().Set("X-Content-Type-Options", "nosniff")
}

func (w *contentTypeOverridingResponseWriter) Header() http.Header {
	return w.original.Header()
}

func (w *contentTypeOverridingResponseWriter) WriteHeader(statusCode int) {
	w.override()
	w.original.WriteHeader(statusCode)
}

func (w *contentTypeOverridingResponseWriter) Write(data []byte) (int, error) {
	w.override()
	return w.original.Write(data)
}

var additionalMimeTypes = map[string]string{
	".go":    "text/plain",
	".py":    "text/plain",
	".zsh":   "text/plain",
	".fish":  "text/plain",
	".cs":    "text/plain",
	".rb":    "text/plain",
	".m":     "text/plain",
	".mm":    "text/plain",
	".swift": "text/plain",
	".flow":  "text/plain",
	".php":   "text/plain",
	".pl":    "text/plain",
	".sh":    "text/plain",
	".js":    "text/plain",
	".json":  "text/plain",
	".sql":   "text/plain",
	".rs":    "text/plain",
	".xml":   "text/plain",
	".tex":   "text/plain",
	".pub":   "text/plain",
}
<commit_msg>libhttpserver: Use Content-Disposition as discussed on chat<commit_after>// Copyright 2018 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libhttpserver

import (
	"net/http"
	"strings"
)

type contentTypeOverridingResponseWriter struct {
	original http.ResponseWriter
}

var _ http.ResponseWriter = (*contentTypeOverridingResponseWriter)(nil)

func newContentTypeOverridingResponseWriter(
	original http.ResponseWriter) *contentTypeOverridingResponseWriter {
	return &contentTypeOverridingResponseWriter{
		original: original,
	}
}

func (w *contentTypeOverridingResponseWriter) calculateOverride(
	mimeType string) (newMimeType, disposition string) {
	// Send text/plain for all HTML and JS files to avoid them being executed
	// by the frontend WebView.
	ty := strings.ToLower(mimeType)
	switch {
	// First anything textual as text/plain.
	// Includes javascript, html, xml. (note that the type may be e.g. application/xhtml+xml)
	case strings.Contains(ty, "javascript") ||
		strings.Contains(ty, "xml") ||
		strings.Contains(ty, "html") ||
		strings.HasPrefix(ty, "text/"):
		return "text/plain", "inline"
	// Pass multimedia types through, and pdf too.
	case strings.HasPrefix(ty, "audio/") ||
		strings.HasPrefix(ty, "image/") ||
		strings.HasPrefix(ty, "video/") ||
		ty == "application/pdf":
		return ty, "inline"
	// Otherwise default to text + attachment.
	default:
		return "text/plain", "attachment"
	}
}

func (w *contentTypeOverridingResponseWriter) override() {
	t := w.original.Header().Get("Content-Type")
	if len(t) > 0 {
		ct, disp := w.calculateOverride(t)
		w.original.Header().Set("Content-Type", ct)
		w.original.Header().Set("Content-Disposition", disp)
	}
	w.original.Header().Set("X-Content-Type-Options", "nosniff")
}

func (w *contentTypeOverridingResponseWriter) Header() http.Header {
	return w.original.Header()
}

func (w *contentTypeOverridingResponseWriter) WriteHeader(statusCode int) {
	w.override()
	w.original.WriteHeader(statusCode)
}

func (w *contentTypeOverridingResponseWriter) Write(data []byte) (int, error) {
	w.override()
	return w.original.Write(data)
}

var additionalMimeTypes = map[string]string{
	".go":    "text/plain",
	".py":    "text/plain",
	".zsh":   "text/plain",
	".fish":  "text/plain",
	".cs":    "text/plain",
	".rb":    "text/plain",
	".m":     "text/plain",
	".mm":    "text/plain",
	".swift": "text/plain",
	".flow":  "text/plain",
	".php":   "text/plain",
	".pl":    "text/plain",
	".sh":    "text/plain",
	".js":    "text/plain",
	".json":  "text/plain",
	".sql":   "text/plain",
	".rs":    "text/plain",
	".xml":   "text/plain",
	".tex":   "text/plain",
	".pub":   "text/plain",
}
<|endoftext|>


<commit_before>package logger

import (
	"fmt"
	"os"

	"github.com/bugsnag/bugsnag-go"
	"github.com/juju/loggo"
)

func init() {
	if os.Getenv("DEBUG") != "" {
		// set the default logger to info
		loggo.GetLogger("").SetLogLevel(loggo.DEBUG)
	} else {
		// set the default logger to info
		loggo.GetLogger("").SetLogLevel(loggo.INFO)
	}
}

// Logger wrapper for the internal logger with some extra helpers
type Logger struct {
	loggo.Logger
}

// GetLogger builds a ninja logger with the given name
func GetLogger(name string) *Logger {
	l := loggo.GetLogger(name)
	return &Logger{l}
}

// HandleError This notifies bugsnag and logs the error.
func (l *Logger) HandleError(err error, msg string) {
	l.Errorf("%s : %v", msg, err)
	bugsnag.Notify(err)
}

// FatalError This notifies bugsnag and logs the error then quits.
func (l *Logger) FatalError(err error, msg string) {
	l.Errorf("%s : %v", msg, err)
	bugsnag.Notify(err)
	os.Exit(1)
}

// HandleErrorf This notifies bugsnag and logs the error based on the args.
func (l *Logger) HandleErrorf(err error, msg string, args ...interface{}) {
	l.Errorf(msg, args)
	bugsnag.Notify(err)
}

// FatalErrorf This notifies bugsnag and logs the error based on the args then quits
func (l *Logger) FatalErrorf(err error, msg string, args ...interface{}) {
	l.Errorf(msg, args)
	bugsnag.Notify(err)
	os.Exit(1)
}

// FatalErrorf This notifies bugsnag and logs the error based on the args then quits
func (l *Logger) Fatalf(msg string, args ...interface{}) {
	l.Errorf(msg, args)
	bugsnag.Notify(fmt.Errorf(msg, args))
	os.Exit(1)
}
<commit_msg>Move magic debug ENV var into get logger to ensure it overrides everything.<commit_after>package logger

import (
	"fmt"
	"os"

	"github.com/bugsnag/bugsnag-go"
	"github.com/juju/loggo"
)

// Logger wrapper for the internal logger with some extra helpers
type Logger struct {
	loggo.Logger
}

// GetLogger builds a ninja logger with the given name
func GetLogger(name string) *Logger {

	// if we are setting or configuring the ROOT logger
	if name == "" {
		// check for the magic DEBUG env variable.
		if os.Getenv("DEBUG") != "" {
			// set the default logger to debug
			loggo.GetLogger("").SetLogLevel(loggo.DEBUG)
			loggo.GetLogger("").Infof("enabling DEBUG")
		}
	}

	l := loggo.GetLogger(name)
	return &Logger{l}
}

// HandleError This notifies bugsnag and logs the error.
func (l *Logger) HandleError(err error, msg string) {
	l.Errorf("%s : %v", msg, err)
	bugsnag.Notify(err)
}

// FatalError This notifies bugsnag and logs the error then quits.
func (l *Logger) FatalError(err error, msg string) {
	l.Errorf("%s : %v", msg, err)
	bugsnag.Notify(err)
	os.Exit(1)
}

// HandleErrorf This notifies bugsnag and logs the error based on the args.
func (l *Logger) HandleErrorf(err error, msg string, args ...interface{}) {
	l.Errorf(msg, args)
	bugsnag.Notify(err)
}

// FatalErrorf This notifies bugsnag and logs the error based on the args then quits
func (l *Logger) FatalErrorf(err error, msg string, args ...interface{}) {
	l.Errorf(msg, args)
	bugsnag.Notify(err)
	os.Exit(1)
}

// FatalErrorf This notifies bugsnag and logs the error based on the args then quits
func (l *Logger) Fatalf(msg string, args ...interface{}) {
	l.Errorf(msg, args)
	bugsnag.Notify(fmt.Errorf(msg, args))
	os.Exit(1)
}
<|endoftext|>


<commit_before>package engine

//
// engine.GPGImportKeyEngine is a class that selects key from the GPG keyring via
// shell-out to the gpg command line client. It's useful in `client mykey select`
// and other places in which the user picks existing PGP keys on the existing
// system for use in Keybase tasks.
//

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

type GPGImportKeyArg struct {
	Query      string
	Signer     libkb.GenericKey
	AllowMulti bool
	SkipImport bool
	OnlyImport bool
	Me         *libkb.User
	Lks        *libkb.LKSec
}

type GPGImportKeyEngine struct {
	last                   *libkb.PGPKeyBundle
	arg                    *GPGImportKeyArg
	duplicatedFingerprints []libkb.PGPFingerprint
	libkb.Contextified
}

func NewGPGImportKeyEngine(arg *GPGImportKeyArg, g *libkb.GlobalContext) *GPGImportKeyEngine {
	return &GPGImportKeyEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (e *GPGImportKeyEngine) Prereqs() Prereqs {
	return Prereqs{
		Session: true,
	}
}

func (e *GPGImportKeyEngine) Name() string {
	return "GPGImportKeyEngine"
}

func (e *GPGImportKeyEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.GPGUIKind,
		libkb.SecretUIKind,
	}
}

func (e *GPGImportKeyEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		NewPGPKeyImportEngine(PGPKeyImportEngineArg{}),
		NewPGPUpdateEngine(PGPUpdateEngineArg{}),
	}
}

func (e *GPGImportKeyEngine) WantsGPG(ctx *Context) (bool, error) {
	gpg := e.G().GetGpgClient()
	canExec, err := gpg.CanExec()
	if err != nil {
		return false, err
	}
	if !canExec {
		return false, nil
	}

	// they have gpg

	res, err := ctx.GPGUI.WantToAddGPGKey(0)
	if err != nil {
		return false, err
	}
	return res, nil
}

func (e *GPGImportKeyEngine) Run(ctx *Context) (err error) {
	gpg := e.G().GetGpgClient()

	me := e.arg.Me
	if me == nil {
		if me, err = libkb.LoadMe(libkb.NewLoadUserPubOptionalArg(e.G())); err != nil {
			return err
		}
	}

	if !e.arg.OnlyImport {
		if err = PGPCheckMulti(me, e.arg.AllowMulti); err != nil {
			return err
		}
	}

	if err = gpg.Configure(); err != nil {
		return err
	}
	index, warns, err := gpg.Index(true, e.arg.Query)
	if err != nil {
		return err
	}
	warns.Warn()

	var gks []keybase1.GPGKey
	for _, key := range index.Keys {
		gk := keybase1.GPGKey{
			Algorithm:  fmt.Sprintf("%d%s", key.Bits, key.AlgoString()),
			KeyID:      key.GetFingerprint().ToKeyID(),
			Expiration: key.ExpirationString(),
			Identities: key.GetPGPIdentities(),
		}
		gks = append(gks, gk)
	}

	if len(gks) == 0 {
		return fmt.Errorf("No PGP keys available to choose from.")
	}

	res, err := ctx.GPGUI.SelectKeyAndPushOption(keybase1.SelectKeyAndPushOptionArg{Keys: gks})
	if err != nil {
		return err
	}
	e.G().Log.Debug("SelectKey result: %+v", res)

	var selected *libkb.GpgPrimaryKey
	for _, key := range index.Keys {
		if key.GetFingerprint().ToKeyID() == res.KeyID {
			selected = key
			break
		}
	}

	if selected == nil {
		return nil
	}

	publicKeys := me.GetActivePGPKeys(false)
	duplicate := false
	for _, key := range publicKeys {
		if key.GetFingerprint().Eq(*(selected.GetFingerprint())) {
			duplicate = true
			break
		}
	}
	if duplicate {
		// This key's already been posted to the server.
		res, err := ctx.GPGUI.ConfirmDuplicateKeyChosen(0)
		if err != nil {
			return err
		}
		if !res {
			return libkb.SibkeyAlreadyExistsError{}
		}
		// We're sending a key update, then.
		fp := fmt.Sprintf("%s", *(selected.GetFingerprint()))
		eng := NewPGPUpdateEngine([]string{fp}, false, e.G())
		err = RunEngine(eng, ctx)
		e.duplicatedFingerprints = eng.duplicatedFingerprints

		return err
	}

	bundle, err := gpg.ImportKey(true, *(selected.GetFingerprint()))
	if err != nil {
		return fmt.Errorf("ImportKey error: %s", err)
	}

	if err := bundle.Unlock("Import of key into keybase keyring", ctx.SecretUI); err != nil {
		return err
	}

	e.G().Log.Info("Bundle unlocked: %s", selected.GetFingerprint().ToKeyID())

	eng := NewPGPKeyImportEngine(PGPKeyImportEngineArg{
		Pregen:     bundle,
		SigningKey: e.arg.Signer,
		Me:         me,
		AllowMulti: e.arg.AllowMulti,
		NoSave:     e.arg.SkipImport,
		OnlySave:   e.arg.OnlyImport,
		Lks:        e.arg.Lks,
	})

	if err = RunEngine(eng, ctx); err != nil {

		// It's important to propagate a CanceledError unmolested,
		// since the UI needs to know that. See:
		//  https://github.com/keybase/client/issues/226
		if _, ok := err.(libkb.CanceledError); !ok {
			err = libkb.KeyGenError{Msg: err.Error()}
		}
		return
	}

	e.G().Log.Info("Key %s imported", selected.GetFingerprint().ToKeyID())

	e.last = bundle

	return nil
}

func (e *GPGImportKeyEngine) LastKey() *libkb.PGPKeyBundle {
	return e.last
}
<commit_msg>Fix up subconsumers<commit_after>package engine

//
// engine.GPGImportKeyEngine is a class that selects key from the GPG keyring via
// shell-out to the gpg command line client. It's useful in `client mykey select`
// and other places in which the user picks existing PGP keys on the existing
// system for use in Keybase tasks.
//

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

type GPGImportKeyArg struct {
	Query      string
	Signer     libkb.GenericKey
	AllowMulti bool
	SkipImport bool
	OnlyImport bool
	Me         *libkb.User
	Lks        *libkb.LKSec
}

type GPGImportKeyEngine struct {
	last                   *libkb.PGPKeyBundle
	arg                    *GPGImportKeyArg
	duplicatedFingerprints []libkb.PGPFingerprint
	libkb.Contextified
}

func NewGPGImportKeyEngine(arg *GPGImportKeyArg, g *libkb.GlobalContext) *GPGImportKeyEngine {
	return &GPGImportKeyEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (e *GPGImportKeyEngine) Prereqs() Prereqs {
	return Prereqs{
		Session: true,
	}
}

func (e *GPGImportKeyEngine) Name() string {
	return "GPGImportKeyEngine"
}

func (e *GPGImportKeyEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.GPGUIKind,
		libkb.SecretUIKind,
	}
}

func (e *GPGImportKeyEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&PGPKeyImportEngine{},
		&PGPUpdateEngine{},
	}
}

func (e *GPGImportKeyEngine) WantsGPG(ctx *Context) (bool, error) {
	gpg := e.G().GetGpgClient()
	canExec, err := gpg.CanExec()
	if err != nil {
		return false, err
	}
	if !canExec {
		return false, nil
	}

	// they have gpg

	res, err := ctx.GPGUI.WantToAddGPGKey(0)
	if err != nil {
		return false, err
	}
	return res, nil
}

func (e *GPGImportKeyEngine) Run(ctx *Context) (err error) {
	gpg := e.G().GetGpgClient()

	me := e.arg.Me
	if me == nil {
		if me, err = libkb.LoadMe(libkb.NewLoadUserPubOptionalArg(e.G())); err != nil {
			return err
		}
	}

	if !e.arg.OnlyImport {
		if err = PGPCheckMulti(me, e.arg.AllowMulti); err != nil {
			return err
		}
	}

	if err = gpg.Configure(); err != nil {
		return err
	}
	index, warns, err := gpg.Index(true, e.arg.Query)
	if err != nil {
		return err
	}
	warns.Warn()

	var gks []keybase1.GPGKey
	for _, key := range index.Keys {
		gk := keybase1.GPGKey{
			Algorithm:  fmt.Sprintf("%d%s", key.Bits, key.AlgoString()),
			KeyID:      key.GetFingerprint().ToKeyID(),
			Expiration: key.ExpirationString(),
			Identities: key.GetPGPIdentities(),
		}
		gks = append(gks, gk)
	}

	if len(gks) == 0 {
		return fmt.Errorf("No PGP keys available to choose from.")
	}

	res, err := ctx.GPGUI.SelectKeyAndPushOption(keybase1.SelectKeyAndPushOptionArg{Keys: gks})
	if err != nil {
		return err
	}
	e.G().Log.Debug("SelectKey result: %+v", res)

	var selected *libkb.GpgPrimaryKey
	for _, key := range index.Keys {
		if key.GetFingerprint().ToKeyID() == res.KeyID {
			selected = key
			break
		}
	}

	if selected == nil {
		return nil
	}

	publicKeys := me.GetActivePGPKeys(false)
	duplicate := false
	for _, key := range publicKeys {
		if key.GetFingerprint().Eq(*(selected.GetFingerprint())) {
			duplicate = true
			break
		}
	}
	if duplicate {
		// This key's already been posted to the server.
		res, err := ctx.GPGUI.ConfirmDuplicateKeyChosen(0)
		if err != nil {
			return err
		}
		if !res {
			return libkb.SibkeyAlreadyExistsError{}
		}
		// We're sending a key update, then.
		fp := fmt.Sprintf("%s", *(selected.GetFingerprint()))
		eng := NewPGPUpdateEngine([]string{fp}, false, e.G())
		err = RunEngine(eng, ctx)
		e.duplicatedFingerprints = eng.duplicatedFingerprints

		return err
	}

	bundle, err := gpg.ImportKey(true, *(selected.GetFingerprint()))
	if err != nil {
		return fmt.Errorf("ImportKey error: %s", err)
	}

	if err := bundle.Unlock("Import of key into keybase keyring", ctx.SecretUI); err != nil {
		return err
	}

	e.G().Log.Info("Bundle unlocked: %s", selected.GetFingerprint().ToKeyID())

	eng := NewPGPKeyImportEngine(PGPKeyImportEngineArg{
		Pregen:     bundle,
		SigningKey: e.arg.Signer,
		Me:         me,
		AllowMulti: e.arg.AllowMulti,
		NoSave:     e.arg.SkipImport,
		OnlySave:   e.arg.OnlyImport,
		Lks:        e.arg.Lks,
	})

	if err = RunEngine(eng, ctx); err != nil {

		// It's important to propagate a CanceledError unmolested,
		// since the UI needs to know that. See:
		//  https://github.com/keybase/client/issues/226
		if _, ok := err.(libkb.CanceledError); !ok {
			err = libkb.KeyGenError{Msg: err.Error()}
		}
		return
	}

	e.G().Log.Info("Key %s imported", selected.GetFingerprint().ToKeyID())

	e.last = bundle

	return nil
}

func (e *GPGImportKeyEngine) LastKey() *libkb.PGPKeyBundle {
	return e.last
}
<|endoftext|>


<commit_before>/*
   Copyright 2017 Shlomi Noach, GitHub Inc.

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

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

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

package logic

import (
	"encoding/json"

	"github.com/github/orchestrator/go/inst"

	"github.com/openark/golib/log"
)

// AsyncRequest represents an entry in the async_request table
type CommandApplier struct {
}

func NewCommandApplier() *CommandApplier {
	applier := &CommandApplier{}
	return applier
}

func (applier *CommandApplier) ApplyCommand(op string, value []byte) interface{} {
	log.Debugf("orchestrator/raft: applying command: %s", op)
	switch op {
	case "register-node":
		return applier.registerNode(value)
	case "discover":
		return applier.discover(value)
	case "forget":
		return applier.forget(value)
	case "begin-downtime":
		return applier.beginDowntime(value)
	case "end-downtime":
		return applier.endDowntime(value)
	case "register-candidate":
		return applier.registerCandidate(value)
	case "ack-recovery":
		return applier.ackRecovery(value)
	case "register-hostname-unresolve":
		return applier.registerHostnameUnresolve(value)
	}
	return log.Errorf("Unknown command op: %s", op)
}

func (applier *CommandApplier) registerNode(value []byte) interface{} {
	return nil
}

func (applier *CommandApplier) discover(value []byte) interface{} {
	instanceKey := inst.InstanceKey{}
	if err := json.Unmarshal(value, &instanceKey); err != nil {
		return log.Errore(err)
	}
	discoverInstance(instanceKey)
	return nil
}

func (applier *CommandApplier) forget(value []byte) interface{} {
	instanceKey := inst.InstanceKey{}
	if err := json.Unmarshal(value, &instanceKey); err != nil {
		return log.Errore(err)
	}
	err := inst.ForgetInstance(&instanceKey)
	return err
}

func (applier *CommandApplier) beginDowntime(value []byte) interface{} {
	downtime := inst.Downtime{}
	if err := json.Unmarshal(value, &downtime); err != nil {
		return log.Errore(err)
	}
	err := inst.BeginDowntime(&downtime)
	return err
}

func (applier *CommandApplier) endDowntime(value []byte) interface{} {
	instanceKey := inst.InstanceKey{}
	if err := json.Unmarshal(value, &instanceKey); err != nil {
		return log.Errore(err)
	}
	_, err := inst.EndDowntime(&instanceKey)
	return err
}

func (applier *CommandApplier) registerCandidate(value []byte) interface{} {
	candidate := inst.CandidateDatabaseInstance{}
	if err := json.Unmarshal(value, &candidate); err != nil {
		return log.Errore(err)
	}
	err := inst.RegisterCandidateInstance(&candidate)
	return err
}

func (applier *CommandApplier) ackRecovery(value []byte) interface{} {
	ack := RecoveryAcknowledgement{}
	err := json.Unmarshal(value, &ack)
	if err != nil {
		return log.Errore(err)
	}
	if ack.ClusterName != "" {
		_, err = AcknowledgeClusterRecoveries(ack.ClusterName, ack.Owner, ack.Comment)
	}
	if ack.Key.IsValid() {
		_, err = AcknowledgeInstanceRecoveries(&ack.Key, ack.Owner, ack.Comment)
	}
	return err
}

func (applier *CommandApplier) registerHostnameUnresolve(value []byte) interface{} {
	ack := RecoveryAcknowledgement{}
	err := json.Unmarshal(value, &ack)
	if err != nil {
		return log.Errore(err)
	}
	if ack.ClusterName != "" {
		_, err = AcknowledgeClusterRecoveries(ack.ClusterName, ack.Owner, ack.Comment)
	}
	if ack.Key.IsValid() {
		_, err = AcknowledgeInstanceRecoveries(&ack.Key, ack.Owner, ack.Comment)
	}
	return err
}
<commit_msg>fixed register-hostname-unresolve<commit_after>/*
   Copyright 2017 Shlomi Noach, GitHub Inc.

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

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

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

package logic

import (
	"encoding/json"

	"github.com/github/orchestrator/go/inst"

	"github.com/openark/golib/log"
)

// AsyncRequest represents an entry in the async_request table
type CommandApplier struct {
}

func NewCommandApplier() *CommandApplier {
	applier := &CommandApplier{}
	return applier
}

func (applier *CommandApplier) ApplyCommand(op string, value []byte) interface{} {
	log.Debugf("orchestrator/raft: applying command: %s", op)
	switch op {
	case "register-node":
		return applier.registerNode(value)
	case "discover":
		return applier.discover(value)
	case "forget":
		return applier.forget(value)
	case "begin-downtime":
		return applier.beginDowntime(value)
	case "end-downtime":
		return applier.endDowntime(value)
	case "register-candidate":
		return applier.registerCandidate(value)
	case "ack-recovery":
		return applier.ackRecovery(value)
	case "register-hostname-unresolve":
		return applier.registerHostnameUnresolve(value)
	}
	return log.Errorf("Unknown command op: %s", op)
}

func (applier *CommandApplier) registerNode(value []byte) interface{} {
	return nil
}

func (applier *CommandApplier) discover(value []byte) interface{} {
	instanceKey := inst.InstanceKey{}
	if err := json.Unmarshal(value, &instanceKey); err != nil {
		return log.Errore(err)
	}
	discoverInstance(instanceKey)
	return nil
}

func (applier *CommandApplier) forget(value []byte) interface{} {
	instanceKey := inst.InstanceKey{}
	if err := json.Unmarshal(value, &instanceKey); err != nil {
		return log.Errore(err)
	}
	err := inst.ForgetInstance(&instanceKey)
	return err
}

func (applier *CommandApplier) beginDowntime(value []byte) interface{} {
	downtime := inst.Downtime{}
	if err := json.Unmarshal(value, &downtime); err != nil {
		return log.Errore(err)
	}
	err := inst.BeginDowntime(&downtime)
	return err
}

func (applier *CommandApplier) endDowntime(value []byte) interface{} {
	instanceKey := inst.InstanceKey{}
	if err := json.Unmarshal(value, &instanceKey); err != nil {
		return log.Errore(err)
	}
	_, err := inst.EndDowntime(&instanceKey)
	return err
}

func (applier *CommandApplier) registerCandidate(value []byte) interface{} {
	candidate := inst.CandidateDatabaseInstance{}
	if err := json.Unmarshal(value, &candidate); err != nil {
		return log.Errore(err)
	}
	err := inst.RegisterCandidateInstance(&candidate)
	return err
}

func (applier *CommandApplier) ackRecovery(value []byte) interface{} {
	ack := RecoveryAcknowledgement{}
	err := json.Unmarshal(value, &ack)
	if err != nil {
		return log.Errore(err)
	}
	if ack.ClusterName != "" {
		_, err = AcknowledgeClusterRecoveries(ack.ClusterName, ack.Owner, ack.Comment)
	}
	if ack.Key.IsValid() {
		_, err = AcknowledgeInstanceRecoveries(&ack.Key, ack.Owner, ack.Comment)
	}
	return err
}

func (applier *CommandApplier) registerHostnameUnresolve(value []byte) interface{} {
	registration := inst.HostnameRegistration{}
	if err := json.Unmarshal(value, &registration); err != nil {
		return log.Errore(err)
	}
	err := inst.RegisterHostnameUnresolve(&registration)
	return err
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package types

import (
	"go/ast"
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
)

var gcPath string // Go compiler path

func init() {
	// determine compiler
	var gc string
	switch runtime.GOARCH {
	case "386":
		gc = "8g"
	case "amd64":
		gc = "6g"
	case "arm":
		gc = "5g"
	default:
		gcPath = "unknown-GOARCH-compiler"
		return
	}
	gcPath = filepath.Join(build.ToolDir, gc)
}

func compile(t *testing.T, dirname, filename string) string {
	cmd := exec.Command(gcPath, filename)
	cmd.Dir = dirname
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Logf("%s", out)
		t.Fatalf("%s %s failed: %s", gcPath, filename, err)
	}
	archCh, _ := build.ArchChar(runtime.GOARCH)
	// filename should end with ".go"
	return filepath.Join(dirname, filename[:len(filename)-2]+archCh)
}

// Use the same global imports map for all tests. The effect is
// as if all tested packages were imported into a single package.
var imports = make(map[string]*Package)

func testPath(t *testing.T, path string) bool {
	t0 := time.Now()
	_, err := GcImport(imports, path)
	if err != nil {
		t.Errorf("testPath(%s): %s", path, err)
		return false
	}
	t.Logf("testPath(%s): %v", path, time.Since(t0))
	return true
}

const maxTime = 30 * time.Second

func testDir(t *testing.T, dir string, endTime time.Time) (nimports int) {
	dirname := filepath.Join(runtime.GOROOT(), "pkg", runtime.GOOS+"_"+runtime.GOARCH, dir)
	list, err := ioutil.ReadDir(dirname)
	if err != nil {
		t.Fatalf("testDir(%s): %s", dirname, err)
	}
	for _, f := range list {
		if time.Now().After(endTime) {
			t.Log("testing time used up")
			return
		}
		switch {
		case !f.IsDir():
			// try extensions
			for _, ext := range pkgExts {
				if strings.HasSuffix(f.Name(), ext) {
					name := f.Name()[0 : len(f.Name())-len(ext)] // remove extension
					if testPath(t, filepath.Join(dir, name)) {
						nimports++
					}
				}
			}
		case f.IsDir():
			nimports += testDir(t, filepath.Join(dir, f.Name()), endTime)
		}
	}
	return
}

func TestGcImport(t *testing.T) {
	// On cross-compile builds, the path will not exist.
	// Need to use GOHOSTOS, which is not available.
	if _, err := os.Stat(gcPath); err != nil {
		t.Skipf("skipping test: %v", err)
	}

	if outFn := compile(t, "testdata", "exports.go"); outFn != "" {
		defer os.Remove(outFn)
	}

	nimports := 0
	if testPath(t, "./testdata/exports") {
		nimports++
	}
	nimports += testDir(t, "", time.Now().Add(maxTime)) // installed packages
	t.Logf("tested %d imports", nimports)
}

var importedObjectTests = []struct {
	name string
	kind ast.ObjKind
	typ  string
}{
	{"unsafe.Pointer", ast.Typ, "unsafe.Pointer"},
	{"math.Pi", ast.Con, "untyped float"},
	{"io.Reader", ast.Typ, "interface{Read(p []byte) (n int, err error)}"},
	{"io.ReadWriter", ast.Typ, "interface{Read(p []byte) (n int, err error); Write(p []byte) (n int, err error)}"},
	{"math.Sin", ast.Fun, "func(x·2 float64) (_ float64)"},
	// TODO(gri) add more tests
}

func TestGcImportedTypes(t *testing.T) {
	// This package does not yet know how to read gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}
	for _, test := range importedObjectTests {
		s := strings.Split(test.name, ".")
		if len(s) != 2 {
			t.Fatal("inconsistent test data")
		}
		importPath := s[0]
		objName := s[1]

		pkg, err := GcImport(imports, importPath)
		if err != nil {
			t.Error(err)
			continue
		}

		obj := pkg.scope.Lookup(objName)

		// TODO(gri) should define an accessor on Object
		var kind ast.ObjKind
		var typ Type
		switch obj := obj.(type) {
		case *Const:
			kind = ast.Con
			typ = obj.typ
		case *TypeName:
			kind = ast.Typ
			typ = obj.typ
		case *Var:
			kind = ast.Var
			typ = obj.typ
		case *Func:
			kind = ast.Fun
			typ = obj.typ
		default:
			unreachable()
		}

		if kind != test.kind {
			t.Errorf("%s: got kind = %q; want %q", test.name, kind, test.kind)
		}

		str := typeString(typ.Underlying())
		if str != test.typ {
			t.Errorf("%s: got type = %q; want %q", test.name, typ, test.typ)
		}
	}
}
<commit_msg>go.tools/go/types: simplified GcImportedTypes test<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package types

import (
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
)

var gcPath string // Go compiler path

func init() {
	// determine compiler
	var gc string
	switch runtime.GOARCH {
	case "386":
		gc = "8g"
	case "amd64":
		gc = "6g"
	case "arm":
		gc = "5g"
	default:
		gcPath = "unknown-GOARCH-compiler"
		return
	}
	gcPath = filepath.Join(build.ToolDir, gc)
}

func compile(t *testing.T, dirname, filename string) string {
	cmd := exec.Command(gcPath, filename)
	cmd.Dir = dirname
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Logf("%s", out)
		t.Fatalf("%s %s failed: %s", gcPath, filename, err)
	}
	archCh, _ := build.ArchChar(runtime.GOARCH)
	// filename should end with ".go"
	return filepath.Join(dirname, filename[:len(filename)-2]+archCh)
}

// Use the same global imports map for all tests. The effect is
// as if all tested packages were imported into a single package.
var imports = make(map[string]*Package)

func testPath(t *testing.T, path string) bool {
	t0 := time.Now()
	_, err := GcImport(imports, path)
	if err != nil {
		t.Errorf("testPath(%s): %s", path, err)
		return false
	}
	t.Logf("testPath(%s): %v", path, time.Since(t0))
	return true
}

const maxTime = 30 * time.Second

func testDir(t *testing.T, dir string, endTime time.Time) (nimports int) {
	dirname := filepath.Join(runtime.GOROOT(), "pkg", runtime.GOOS+"_"+runtime.GOARCH, dir)
	list, err := ioutil.ReadDir(dirname)
	if err != nil {
		t.Fatalf("testDir(%s): %s", dirname, err)
	}
	for _, f := range list {
		if time.Now().After(endTime) {
			t.Log("testing time used up")
			return
		}
		switch {
		case !f.IsDir():
			// try extensions
			for _, ext := range pkgExts {
				if strings.HasSuffix(f.Name(), ext) {
					name := f.Name()[0 : len(f.Name())-len(ext)] // remove extension
					if testPath(t, filepath.Join(dir, name)) {
						nimports++
					}
				}
			}
		case f.IsDir():
			nimports += testDir(t, filepath.Join(dir, f.Name()), endTime)
		}
	}
	return
}

func TestGcImport(t *testing.T) {
	// On cross-compile builds, the path will not exist.
	// Need to use GOHOSTOS, which is not available.
	if _, err := os.Stat(gcPath); err != nil {
		t.Skipf("skipping test: %v", err)
	}

	if outFn := compile(t, "testdata", "exports.go"); outFn != "" {
		defer os.Remove(outFn)
	}

	nimports := 0
	if testPath(t, "./testdata/exports") {
		nimports++
	}
	nimports += testDir(t, "", time.Now().Add(maxTime)) // installed packages
	t.Logf("tested %d imports", nimports)
}

var importedObjectTests = []struct {
	name string
	want string
}{
	{"unsafe.Pointer", "type Pointer unsafe.Pointer"},
	{"math.Pi", "const Pi untyped float"},
	{"io.Reader", "type Reader interface{Read(p []byte) (n int, err error)}"},
	{"io.ReadWriter", "type ReadWriter interface{Read(p []byte) (n int, err error); Write(p []byte) (n int, err error)}"},
	{"math.Sin", "func math.Sin(x·2 float64) (_ float64)"},
	// TODO(gri) add more tests
}

func TestGcImportedTypes(t *testing.T) {
	// This package does not yet know how to read gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}
	for _, test := range importedObjectTests {
		s := strings.Split(test.name, ".")
		if len(s) != 2 {
			t.Fatal("inconsistent test data")
		}
		importPath := s[0]
		objName := s[1]

		pkg, err := GcImport(imports, importPath)
		if err != nil {
			t.Error(err)
			continue
		}

		obj := pkg.scope.Lookup(objName)
		if obj == nil {
			t.Errorf("%s: object not found", test.name)
			continue
		}

		got := obj.String()
		if got != test.want {
			t.Errorf("%s: got %q; want %q", test.name, got, test.want)
		}
	}
}
<|endoftext|>


<commit_before>package gorillamux

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"reflect"

	"github.com/domonda/errors"
	"github.com/h2non/filetype"
	"github.com/h2non/filetype/types"

	command "github.com/ungerik/go-command"
)

type ResultsWriter interface {
	WriteResults(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error
}

type ResultsWriterFunc func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error

func (f ResultsWriterFunc) WriteResults(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	return f(args, vars, resultVals, resultErr, writer, request)
}

func encodeJSON(response interface{}) ([]byte, error) {
	if PrettyPrint {
		return json.MarshalIndent(response, "", PrettyPrintIndent)
	}
	return json.Marshal(response)
}

var RespondJSON ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf []byte
	for _, resultVal := range resultVals {
		b, err := encodeJSON(resultVal.Interface())
		if err != nil {
			return err
		}
		buf = append(buf, b...)
	}
	writer.Header().Set("Content-Type", "application/json; charset=utf-8")
	writer.Write(buf)
	return nil
}

// RespondBinary responds with contentType using the binary data from results of type []byte, string, or io.Reader.
func RespondBinary(contentType string) ResultsWriterFunc {
	return func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) (err error) {
		if resultErr != nil {
			return resultErr
		}
		var buf bytes.Buffer
		for _, resultVal := range resultVals {
			switch data := resultVal.Interface().(type) {
			case []byte:
				_, err = buf.Write(data)
			case string:
				_, err = buf.WriteString(data)
			case io.Reader:
				_, err = io.Copy(&buf, data)
			default:
				return errors.Errorf("RespondBinary does not support result type %s", resultVal.Type())
			}
			if err != nil {
				return err
			}
		}
		writer.Header().Set("Content-Type", contentType)
		writer.Write(buf.Bytes())
		return nil
	}
}

func RespondJSONField(fieldName string) ResultsWriterFunc {
	return func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) (err error) {
		if resultErr != nil {
			return resultErr
		}
		var buf []byte
		m := make(map[string]interface{})
		if len(resultVals) > 0 {
			m[fieldName] = resultVals[0].Interface()
		}
		buf, err = encodeJSON(m)
		if err != nil {
			return err
		}
		writer.Header().Set("Content-Type", "application/json; charset=utf-8")
		writer.Write(buf)
		return nil
	}
}

func encodeXML(response interface{}) ([]byte, error) {
	if PrettyPrint {
		return xml.MarshalIndent(response, "", PrettyPrintIndent)
	}
	return xml.Marshal(response)
}

var RespondXML ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf []byte
	for _, resultVal := range resultVals {
		b, err := encodeXML(resultVal.Interface())
		if err != nil {
			return err
		}
		buf = append(buf, b...)
	}
	writer.Header().Set("Content-Type", "application/xml; charset=utf-8")
	writer.Write(buf)
	return nil
}

var RespondPlaintext ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf bytes.Buffer
	for _, resultVal := range resultVals {
		result := resultVal.Interface()
		if b, ok := result.([]byte); ok {
			buf.Write(b)
		} else {
			fmt.Fprintf(&buf, "%s", result)
		}
	}
	writer.Header().Add("Content-Type", "text/plain; charset=utf-8")
	writer.Write(buf.Bytes())
	return nil
}

var RespondHTML ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf bytes.Buffer
	for _, resultVal := range resultVals {
		result := resultVal.Interface()
		if b, ok := result.([]byte); ok {
			buf.Write(b)
		} else {
			fmt.Fprintf(&buf, "%s", result)
		}
	}
	writer.Header().Add("Content-Type", "text/html; charset=utf-8")
	writer.Write(buf.Bytes())
	return nil
}

var RespondDetectContentType ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	if len(resultVals) != 1 {
		return errors.Errorf("RespondDetectContentType needs 1 result, got %d", len(resultVals))
	}
	data, ok := resultVals[0].Interface().([]byte)
	if !ok {
		return errors.Errorf("RespondDetectContentType needs []byte result, got %T", resultVals[0].Interface())
	}

	writer.Header().Add("Content-Type", DetectContentType(data))
	writer.Write(data)
	return nil
}

var RespondNothing ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	return resultErr
}

// DetectContentType tries to detect the MIME content-type of data,
// or returns "application/octet-stream" if none could be identified.
func DetectContentType(data []byte) string {
	kind, _ := filetype.Match(data)
	if kind == types.Unknown {
		return http.DetectContentType(data)
	}
	return kind.MIME.Value
}
<commit_msg>added gorillamux.RespondContentType<commit_after>package gorillamux

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"reflect"

	"github.com/domonda/errors"
	"github.com/h2non/filetype"
	"github.com/h2non/filetype/types"

	command "github.com/ungerik/go-command"
)

type ResultsWriter interface {
	WriteResults(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error
}

type ResultsWriterFunc func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error

func (f ResultsWriterFunc) WriteResults(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	return f(args, vars, resultVals, resultErr, writer, request)
}

func encodeJSON(response interface{}) ([]byte, error) {
	if PrettyPrint {
		return json.MarshalIndent(response, "", PrettyPrintIndent)
	}
	return json.Marshal(response)
}

var RespondJSON ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf []byte
	for _, resultVal := range resultVals {
		b, err := encodeJSON(resultVal.Interface())
		if err != nil {
			return err
		}
		buf = append(buf, b...)
	}
	writer.Header().Set("Content-Type", "application/json; charset=utf-8")
	writer.Write(buf)
	return nil
}

// RespondBinary responds with contentType using the binary data from results of type []byte, string, or io.Reader.
func RespondBinary(contentType string) ResultsWriterFunc {
	return func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) (err error) {
		if resultErr != nil {
			return resultErr
		}
		var buf bytes.Buffer
		for _, resultVal := range resultVals {
			switch data := resultVal.Interface().(type) {
			case []byte:
				_, err = buf.Write(data)
			case string:
				_, err = buf.WriteString(data)
			case io.Reader:
				_, err = io.Copy(&buf, data)
			default:
				return errors.Errorf("RespondBinary does not support result type %s", resultVal.Type())
			}
			if err != nil {
				return err
			}
		}
		writer.Header().Set("Content-Type", contentType)
		writer.Write(buf.Bytes())
		return nil
	}
}

func RespondJSONField(fieldName string) ResultsWriterFunc {
	return func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) (err error) {
		if resultErr != nil {
			return resultErr
		}
		var buf []byte
		m := make(map[string]interface{})
		if len(resultVals) > 0 {
			m[fieldName] = resultVals[0].Interface()
		}
		buf, err = encodeJSON(m)
		if err != nil {
			return err
		}
		writer.Header().Set("Content-Type", "application/json; charset=utf-8")
		writer.Write(buf)
		return nil
	}
}

func encodeXML(response interface{}) ([]byte, error) {
	if PrettyPrint {
		return xml.MarshalIndent(response, "", PrettyPrintIndent)
	}
	return xml.Marshal(response)
}

var RespondXML ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf []byte
	for _, resultVal := range resultVals {
		b, err := encodeXML(resultVal.Interface())
		if err != nil {
			return err
		}
		buf = append(buf, b...)
	}
	writer.Header().Set("Content-Type", "application/xml; charset=utf-8")
	writer.Write(buf)
	return nil
}

var RespondPlaintext ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf bytes.Buffer
	for _, resultVal := range resultVals {
		result := resultVal.Interface()
		if b, ok := result.([]byte); ok {
			buf.Write(b)
		} else {
			fmt.Fprintf(&buf, "%s", result)
		}
	}
	writer.Header().Add("Content-Type", "text/plain; charset=utf-8")
	writer.Write(buf.Bytes())
	return nil
}

var RespondHTML ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	var buf bytes.Buffer
	for _, resultVal := range resultVals {
		result := resultVal.Interface()
		if b, ok := result.([]byte); ok {
			buf.Write(b)
		} else {
			fmt.Fprintf(&buf, "%s", result)
		}
	}
	writer.Header().Add("Content-Type", "text/html; charset=utf-8")
	writer.Write(buf.Bytes())
	return nil
}

var RespondDetectContentType ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	if resultErr != nil {
		return resultErr
	}
	if len(resultVals) != 1 {
		return errors.Errorf("RespondDetectContentType needs 1 result, got %d", len(resultVals))
	}
	data, ok := resultVals[0].Interface().([]byte)
	if !ok {
		return errors.Errorf("RespondDetectContentType needs []byte result, got %s", resultVals[0].Type())
	}

	writer.Header().Add("Content-Type", DetectContentType(data))
	writer.Write(data)
	return nil
}

func RespondContentType(contentType string) ResultsWriter {
	return ResultsWriterFunc(func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
		if resultErr != nil {
			return resultErr
		}
		if len(resultVals) != 1 {
			return errors.Errorf("RespondDetectContentType needs 1 result, got %d", len(resultVals))
		}
		data, ok := resultVals[0].Interface().([]byte)
		if !ok {
			return errors.Errorf("RespondDetectContentType needs []byte result, got %s", resultVals[0].Type())
		}

		writer.Header().Add("Content-Type", contentType)
		writer.Write(data)
		return nil
	})
}

var RespondNothing ResultsWriterFunc = func(args command.Args, vars map[string]string, resultVals []reflect.Value, resultErr error, writer http.ResponseWriter, request *http.Request) error {
	return resultErr
}

// DetectContentType tries to detect the MIME content-type of data,
// or returns "application/octet-stream" if none could be identified.
func DetectContentType(data []byte) string {
	kind, _ := filetype.Match(data)
	if kind == types.Unknown {
		return http.DetectContentType(data)
	}
	return kind.MIME.Value
}
<|endoftext|>