Add concept of platform IDs

b1naryth1ef · Jul 25, 2017 · 9060dc1 · 9060dc1
1 parent 3c0f113
commit 9060dc1
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 24 deletions.
diff --git a/lib/config.go b/lib/config.go
@@ -18,12 +18,13 @@ type AccountMFA struct {
 
 // Accounts represent individual users (auth keys) that can login
 type Account struct {
-	Username   string     `json:"username"`
-	Password   string     `json:"password"`
-	SSHKeysRaw []string   `json:"ssh-keys"`
-	MFA        AccountMFA `json:"mfa,omitempty"`
-	Whitelist  string     `json:"whitelist"`
-	Blacklist  string     `json:"blacklist"`
+	Username    string            `json:"username"`
+	Password    string            `json:"password"`
+	SSHKeysRaw  []string          `json:"ssh-keys"`
+	MFA         AccountMFA        `json:"mfa,omitempty"`
+	Whitelist   string            `json:"whitelist"`
+	Blacklist   string            `json:"blacklist"`
+	PlatformIDs map[string]string `json:"platform_ids"`
 
 	whitelistRe *regexp.Regexp
 	blacklistRe *regexp.Regexp

diff --git a/lib/session.go b/lib/session.go
@@ -264,7 +264,8 @@ func (s *SSHSession) handleChannelForward(newChannel ssh.NewChannel) {
 	}
 
 	for _, wp := range s.State.WebhookProviders {
-		wp.NotifySessionStart(s.Conn.User(), s.UUID, msg.RAddr, fmt.Sprintf("%s", s.Conn.RemoteAddr()))
+		platformID := s.Account.PlatformIDs[wp.PlatformName()]
+		wp.NotifySessionStart(platformID, s.Conn.User(), s.UUID, msg.RAddr, fmt.Sprintf("%s", s.Conn.RemoteAddr()))
 	}
 
 	conn, err := net.Dial("tcp", address)

diff --git a/lib/webhook.go b/lib/webhook.go
@@ -19,14 +19,18 @@ type Embed struct {
 }
 
 type WebhookProvider interface {
-	NotifySessionStart(username, sessionID, proxyHost, sourceHost string) error
-	NotifySessionEnd(username, sessionID, proxyHost, sourceHost string) error
+	NotifySessionStart(platformID, username, sessionID, proxyHost, sourceHost string) error
+	PlatformName() string
 }
 
 type DiscordWebhookProvider struct {
 	URL string
 }
 
+func (d DiscordWebhookProvider) PlatformName() string {
+	return "discord"
+}
+
 func (d DiscordWebhookProvider) send(payload MessagePayload) (err error) {
 	data, err := json.Marshal(payload)
 	if err != nil {
@@ -46,24 +50,20 @@ func (d DiscordWebhookProvider) send(payload MessagePayload) (err error) {
 	return err
 }
 
-func (d DiscordWebhookProvider) NotifySessionStart(username, sessionID, proxyHost, sourceHost string) error {
-	return d.send(MessagePayload{Embeds: []Embed{Embed{
-		Title: fmt.Sprintf("SSH session started by %s", username),
-		Description: fmt.Sprintf(
-			"**Host:** %s\n**Source:** %s\n**Session:** `%s`\n",
-			proxyHost,
-			sourceHost,
-			sessionID,
-		),
-		Color: 7855479,
-	}}})
-}
+func (d DiscordWebhookProvider) NotifySessionStart(platformID, username, sessionID, proxyHost, sourceHost string) error {
+	var title string
+
+	if platformID != "" {
+		title = fmt.Sprintf("<@%s>@%s", platformID, proxyHost)
+	} else {
+		title = fmt.Sprintf("%s@%s", username, proxyHost)
+	}
 
-func (d DiscordWebhookProvider) NotifySessionEnd(username, sessionID, proxyHost, sourceHost string) error {
 	return d.send(MessagePayload{Embeds: []Embed{Embed{
-		Title: fmt.Sprintf("SSH session ended by %s", username),
+		Title: title,
 		Description: fmt.Sprintf(
-			"**Host:** %s\n**Source:** %s\n**Session:** `%s`\n",
+			"**User:** %s\n**Host:** %s\n**Source:** %s\n**Session:** `%s`\n",
+			username,
 			proxyHost,
 			sourceHost,
 			sessionID,