<commit_before>package enmime

import (
	"fmt"
	"log"
	//"net/textproto"
	//"net/mail"
	"strings"

	"github.com/cention-sany/mime"
	"github.com/cention-sany/net/mail"
	"github.com/cention-sany/net/textproto"
	"github.com/jaytaylor/html2text"
)

// MIMEBody is the outer wrapper for MIME messages.
type MIMEBody struct {
	Text           string      // The plain text portion of the message
	HTML           string      // The HTML portion of the message
	IsTextFromHTML bool        // Plain text was empty; down-converted HTML
	Root           MIMEPart    // The top-level MIMEPart
	Attachments    []MIMEPart  // All parts having a Content-Disposition of attachment
	Inlines        []MIMEPart  // All parts having a Content-Disposition of inline
	OtherParts     []MIMEPart  // All parts not in Attachments and Inlines
	header         mail.Header // Header from original message
}

var AddressHeaders = []string{"From", "To", "Delivered-To", "Cc", "Bcc", "Reply-To"}

// IsMultipartMessage returns true if the message has a recognized multipart Content-Type
// header.  You don't need to check this before calling ParseMIMEBody, it can handle
// non-multipart messages.
func IsMultipartMessage(mailMsg *mail.Message) bool {
	// Parse top-level multipart
	ctype := mailMsg.Header.Get("Content-Type")
	mediatype, _, err := mime.ParseMediaType(ctype)
	if err != nil && mime.IsOkPMTError(err) != nil {
		return false
	}
	switch mediatype {
	case "multipart/alternative",
		"multipart/mixed",
		"multipart/related",
		"multipart/signed":
		return true
	default:
		if strings.HasPrefix(mediatype, "multipart/") {
			// according to rfc2046#section-5.1.7 all other multipart should
			// be treated as multipart/mixed
			return true
		}
	}
	return false
}

// IsAttachment returns true, if the given header defines an attachment.
// First it checks, if the Content-Disposition header defines an attachement.
// If this test is false, the Content-Type header is checked.
//
// Valid Attachment-Headers:
//
//    Content-Disposition: attachment; filename="frog.jpg"
//    Content-Type: attachment; filename="frog.jpg"
//
func IsAttachment(header mail.Header) bool {
	mediatype, _, _ := mime.ParseMediaType(header.Get("Content-Disposition"))
	if strings.ToLower(mediatype) == "attachment" {
		return true
	}

	mediatype, _, _ = mime.ParseMediaType(header.Get("Content-Type"))
	if strings.ToLower(mediatype) == "attachment" {
		return true
	}

	return false
}

// IsPlain returns true, if the the mime headers define a valid
// 'text/plain' or 'text/html part'. Ff emptyContentTypeIsPlain is set
// to true, a missing Content-Type header will result in a positive
// plain part detection.
func IsPlain(header mail.Header, emptyContentTypeIsPlain bool) bool {
	ctype := header.Get("Content-Type")
	if ctype == "" && emptyContentTypeIsPlain {
		return true
	}

	mediatype, _, err := mime.ParseMediaType(ctype)
	if err != nil {
		return false
	}
	switch mediatype {
	case "text/plain",
		"text/html":
		return true
	}

	return false

}

// IsBinaryBody returns true, if the mail header defines a binary body.
func IsBinaryBody(mailMsg *mail.Message) bool {
	if IsAttachment(mailMsg.Header) == true {
		return true
	}

	return !IsPlain(mailMsg.Header, true)
}

// Returns a MIME message with only one Attachment, the parsed original mail body.
func binMIME(mailMsg *mail.Message) (*MIMEBody, error) {
	// Root Node of our tree
	ctype := mailMsg.Header.Get("Content-Type")
	mediatype, mparams, err := mime.ParseMediaType(ctype)
	if err != nil {
		mediatype = "attachment"
	}

	m := &MIMEBody{
		header:         mailMsg.Header,
		Root:           NewMIMEPart(nil, mediatype),
		IsTextFromHTML: false,
	}

	p := NewMIMEPart(nil, mediatype)
	p.content, err = decodeSection(mailMsg.Header.Get("Content-Transfer-Encoding"), mailMsg.Body)
	if err != nil {
		return nil, err
	}

	// get set headers
	p.header = make(textproto.MIMEHeader, 4)
	// Figure out our disposition, filename
	disposition, dparams, err := mime.ParseMediaType(mailMsg.Header.Get("Content-Disposition"))
	if err == nil {
		// Disposition is optional
		p.disposition = disposition
		p.fileName = DecodeHeader(dparams["filename"])
	}
	if p.fileName == "" && mparams["name"] != "" {
		p.fileName = DecodeHeader(mparams["name"])
	}
	if p.fileName == "" && mparams["file"] != "" {
		p.fileName = DecodeHeader(mparams["file"])
	}
	if p.charset == "" {
		p.charset = mparams["charset"]
	}

	p.header.Set("Content-Type", mailMsg.Header.Get("Content-Type"))
	p.header.Set("Content-Disposition", mailMsg.Header.Get("Content-Disposition"))

	m.Attachments = append(m.Attachments, p)
	return m, err
}

// ParseMIMEBody parses the body of the message object into a  tree of MIMEPart objects,
// each of which is aware of its content type, filename and headers.  If the part was
// encoded in quoted-printable or base64, it is decoded before being stored in the
// MIMEPart object.
func ParseMIMEBody(mailMsg *mail.Message) (*MIMEBody, error) {
	var gerr error
	mimeMsg := &MIMEBody{
		IsTextFromHTML: false,
		header:         mailMsg.Header,
	}

	if !IsMultipartMessage(mailMsg) {
		// Attachment only?
		if IsBinaryBody(mailMsg) {
			return binMIME(mailMsg)
		}

		// Parse as text only
		bodyBytes, err := decodeSection(mailMsg.Header.Get("Content-Transfer-Encoding"),
			mailMsg.Body)
		if err != nil {
			return nil, fmt.Errorf("Error decoding text-only message: %v", err)
		}
		// Handle plain ASCII text, content-type unspecified
		mimeMsg.Text = string(bodyBytes)

		// Check for HTML at top-level, eat errors quietly
		ctype := mailMsg.Header.Get("Content-Type")
		if ctype != "" {
			if mediatype, mparams, err := mime.ParseMediaType(ctype); err == nil || mime.IsOkPMTError(err) == nil {
				/*
				 *Content-Type: text/plain;\t charset="hz-gb-2312"
				 */
				if mparams["charset"] != "" {
					// Convert plain text to UTF8 if content type specified a charset
					newStr, err := ConvertToUTF8String(mparams["charset"], bodyBytes)
					if err != nil {
						if newStr == "" {
							return nil, err
						} else {
							gerr = err
						}
					} else if newStr != "" {
						mimeMsg.Text = newStr
					}
				} else if mediatype == "text/html" { // charset is empty, look in html body for charset
					charset, err := charsetFromHTMLString(mimeMsg.Text)

					if charset != "" && err == nil {
						newStr, err := ConvertToUTF8String(charset, bodyBytes)
						if err != nil && newStr == "" {
							return nil, err
						} else {
							if err != nil {
								gerr = err
							}
							mimeMsg.Text = newStr
						}
					}
				}
				if mediatype == "text/html" {
					mimeMsg.HTML = mimeMsg.Text
					// Empty Text will trigger html2text conversion below
					mimeMsg.Text = ""
				}
			}
		}
	} else {
		// Parse top-level multipart
		ctype := mailMsg.Header.Get("Content-Type")
		mediatype, params, err := mime.ParseMediaType(ctype)
		if err != nil && mime.IsOkPMTError(err) != nil {
			return nil, fmt.Errorf("Unable to parse media type: %v", err)
		}
		if !strings.HasPrefix(mediatype, "multipart/") {
			return nil, fmt.Errorf("Unknown mediatype: %v", mediatype)
		}
		boundary := params["boundary"]
		if boundary == "" {
			return nil, fmt.Errorf("Unable to locate boundary param in Content-Type header")
		}
		// Root Node of our tree
		root := NewMIMEPart(nil, mediatype)
		mimeMsg.Root = root
		err = parseParts(root, mailMsg.Body, boundary)
		if err != nil {
			return nil, err
		}

		// Locate text body
		if mediatype == "multipart/altern" {
			log.Println("surPrise: should not be here go.enmime/mime.go")
			match := BreadthMatchFirst(root, func(p MIMEPart) bool {
				return p.ContentType() == "text/plain" && p.Disposition() != "attachment"
			})
			if match != nil {
				if match.Charset() != "" {
					newStr, err := ConvertToUTF8String(match.Charset(), match.Content())
					if err != nil {
						if newStr == "" {
							return nil, err
						} else {
							gerr = err
						}
					}
					mimeMsg.Text += newStr
				} else {
					mimeMsg.Text += string(match.Content())
				}
			}
		} else {
			// multipart is of a mixed type
			match := DepthMatchAll(root, func(p MIMEPart) bool {
				return p.ContentType() == "text/plain" && p.Disposition() != "attachment"
			})
			for i, m := range match {
				if i > 0 {
					mimeMsg.Text += "\n--\n"
				}
				if m.Charset() != "" {
					newStr, err := ConvertToUTF8String(m.Charset(), m.Content())
					if err != nil {
						if newStr == "" {
							return nil, err
						} else {
							gerr = err
						}
					}
					mimeMsg.Text += newStr
				} else {
					mimeMsg.Text += string(m.Content())
				}
			}
		}

		// Locate HTML body
		match := BreadthMatchFirst(root, func(p MIMEPart) bool {
			return p.ContentType() == "text/html" && p.Disposition() != "attachment"
		})
		if match != nil {
			if match.Charset() != "" {
				newStr, err := ConvertToUTF8String(match.Charset(), match.Content())
				if err != nil {
					if newStr == "" {
						return nil, err
					} else {
						gerr = err
					}
				}
				mimeMsg.HTML += newStr
			} else {
				mimeMsg.HTML = string(match.Content())
			}

		}

		// Locate attachments
		mimeMsg.Attachments = BreadthMatchAll(root, func(p MIMEPart) bool {
			return p.Disposition() == "attachment" || p.ContentType() == "application/octet-stream"
		})

		// Locate inlines
		mimeMsg.Inlines = BreadthMatchAll(root, func(p MIMEPart) bool {
			return p.Disposition() == "inline"
		})

		// Locate others parts not handled in "Attachments" and "inlines"
		mimeMsg.OtherParts = BreadthMatchAll(root, func(p MIMEPart) bool {
			if strings.HasPrefix(p.ContentType(), "multipart/") {
				return false
			}

			if p.Disposition() != "" {
				return false
			}

			if p.ContentType() == "application/octet-stream" {
				return false
			}

			return p.ContentType() != "text/plain" && p.ContentType() != "text/html"
		})
	}

	// Down-convert HTML to text if necessary
	if mimeMsg.Text == "" && mimeMsg.HTML != "" {
		mimeMsg.IsTextFromHTML = true
		var err error
		if mimeMsg.Text, err = html2text.FromString(mimeMsg.HTML); err != nil {
			// Fail gently
			mimeMsg.Text = ""
			return mimeMsg, err
		}
	}
	return mimeMsg, gerr
}

// Process the specified header for RFC 2047 encoded words and return the result
func (m *MIMEBody) GetHeader(name string) string {
	return DecodeHeader(m.header.Get(name))
}

// Return AddressList with RFC 2047 encoded encoded names.
func (m *MIMEBody) AddressList(key string) ([]*mail.Address, error) {
	isAddrHeader := false
	for _, hkey := range AddressHeaders {
		if strings.ToLower(hkey) == strings.ToLower(key) {
			isAddrHeader = true
			break
		}
	}
	if !isAddrHeader {
		return nil, fmt.Errorf("%s is not address header", key)
	}

	str := DecodeToUTF8Base64Header(m.header.Get(key))
	if str == "" {
		return nil, mail.ErrHeaderNotPresent
	}
	ret, err := mail.ParseAddressList(str)
	if err != nil {
		return nil, err
	}
	return ret, nil
}
<commit_msg>Better handling the non-multipart email body<commit_after>package enmime

import (
	"fmt"
	"log"
	//"net/textproto"
	//"net/mail"
	"strings"

	"github.com/cention-sany/mime"
	"github.com/cention-sany/net/mail"
	"github.com/cention-sany/net/textproto"
	"github.com/jaytaylor/html2text"
)

// MIMEBody is the outer wrapper for MIME messages.
type MIMEBody struct {
	Text           string      // The plain text portion of the message
	HTML           string      // The HTML portion of the message
	IsTextFromHTML bool        // Plain text was empty; down-converted HTML
	Root           MIMEPart    // The top-level MIMEPart
	Attachments    []MIMEPart  // All parts having a Content-Disposition of attachment
	Inlines        []MIMEPart  // All parts having a Content-Disposition of inline
	OtherParts     []MIMEPart  // All parts not in Attachments and Inlines
	header         mail.Header // Header from original message
}

var AddressHeaders = []string{"From", "To", "Delivered-To", "Cc", "Bcc", "Reply-To"}

// IsMultipartMessage returns true if the message has a recognized multipart Content-Type
// header.  You don't need to check this before calling ParseMIMEBody, it can handle
// non-multipart messages.
func IsMultipartMessage(mailMsg *mail.Message) bool {
	// Parse top-level multipart
	ctype := mailMsg.Header.Get("Content-Type")
	mediatype, _, err := mime.ParseMediaType(ctype)
	if err != nil && mime.IsOkPMTError(err) != nil {
		return false
	}
	switch mediatype {
	case "multipart/alternative",
		"multipart/mixed",
		"multipart/related",
		"multipart/signed":
		return true
	default:
		if strings.HasPrefix(mediatype, "multipart/") {
			// according to rfc2046#section-5.1.7 all other multipart should
			// be treated as multipart/mixed
			return true
		}
	}
	return false
}

// IsAttachment returns true, if the given header defines an attachment.
// First it checks, if the Content-Disposition header defines an attachement.
// If this test is false, the Content-Type header is checked.
//
// Valid Attachment-Headers:
//
//    Content-Disposition: attachment; filename="frog.jpg"
//    Content-Type: attachment; filename="frog.jpg"
//
func IsAttachment(header mail.Header) bool {
	mediatype, _, _ := mime.ParseMediaType(header.Get("Content-Disposition"))
	if strings.ToLower(mediatype) == "attachment" {
		return true
	}

	mediatype, _, _ = mime.ParseMediaType(header.Get("Content-Type"))
	if strings.ToLower(mediatype) == "attachment" {
		return true
	}

	return false
}

// IsPlain returns true, if the the mime headers define a valid
// 'text/plain' or 'text/html part'. Ff emptyContentTypeIsPlain is set
// to true, a missing Content-Type header will result in a positive
// plain part detection.
func IsPlain(header mail.Header, emptyContentTypeIsPlain bool) bool {
	ctype := header.Get("Content-Type")
	if ctype == "" && emptyContentTypeIsPlain {
		return true
	}

	mediatype, _, err := mime.ParseMediaType(ctype)
	if err != nil {
		return false
	}
	switch mediatype {
	case "text/plain",
		"text/html":
		return true
	}

	return false

}

// IsBinaryBody returns true, if the mail header defines a binary body.
func IsBinaryBody(mailMsg *mail.Message) bool {
	if IsAttachment(mailMsg.Header) == true {
		return true
	}

	return !IsPlain(mailMsg.Header, true)
}

// Returns a MIME message with only one Attachment, the parsed original mail body.
func binMIME(mailMsg *mail.Message) (*MIMEBody, error) {
	// Root Node of our tree
	ctype := mailMsg.Header.Get("Content-Type")
	mediatype, mparams, err := mime.ParseMediaType(ctype)
	if err != nil {
		mediatype = "attachment"
	}

	m := &MIMEBody{
		header:         mailMsg.Header,
		Root:           NewMIMEPart(nil, mediatype),
		IsTextFromHTML: false,
	}

	p := NewMIMEPart(nil, mediatype)
	p.content, err = decodeSection(mailMsg.Header.Get("Content-Transfer-Encoding"), mailMsg.Body)
	if err != nil {
		return nil, err
	}

	// get set headers
	p.header = make(textproto.MIMEHeader, 4)
	// Figure out our disposition, filename
	disposition, dparams, err := mime.ParseMediaType(mailMsg.Header.Get("Content-Disposition"))
	if err == nil {
		// Disposition is optional
		p.disposition = disposition
		p.fileName = DecodeHeader(dparams["filename"])
	}
	if p.fileName == "" && mparams["name"] != "" {
		p.fileName = DecodeHeader(mparams["name"])
	}
	if p.fileName == "" && mparams["file"] != "" {
		p.fileName = DecodeHeader(mparams["file"])
	}
	if p.charset == "" {
		p.charset = mparams["charset"]
	}

	p.header.Set("Content-Type", mailMsg.Header.Get("Content-Type"))
	p.header.Set("Content-Disposition", mailMsg.Header.Get("Content-Disposition"))

	m.Attachments = append(m.Attachments, p)
	return m, err
}

// ParseMIMEBody parses the body of the message object into a  tree of MIMEPart objects,
// each of which is aware of its content type, filename and headers.  If the part was
// encoded in quoted-printable or base64, it is decoded before being stored in the
// MIMEPart object.
func ParseMIMEBody(mailMsg *mail.Message) (*MIMEBody, error) {
	var gerr error
	mimeMsg := &MIMEBody{
		IsTextFromHTML: false,
		header:         mailMsg.Header,
	}

	if !IsMultipartMessage(mailMsg) {
		// Attachment only?
		if IsBinaryBody(mailMsg) {
			return binMIME(mailMsg)
		}

		// Parse as text only
		bodyBytes, err := decodeSection(mailMsg.Header.Get("Content-Transfer-Encoding"),
			mailMsg.Body)
		if err != nil {
			return nil, fmt.Errorf("Error decoding text-only message: %v", err)
		}
		// Handle plain ASCII text, content-type unspecified
		mimeMsg.Text = string(bodyBytes)

		// Check for HTML at top-level, eat errors quietly
		ctype := mailMsg.Header.Get("Content-Type")
		if ctype != "" {
			if mediatype, mparams, err := mime.ParseMediaType(ctype); err == nil || mime.IsOkPMTError(err) == nil {
				/*
				 *Content-Type: text/plain;\t charset="hz-gb-2312"
				 */
				if mparams["charset"] != "" {
					// Convert plain text to UTF8 if content type specified a charset
					newStr, err := ConvertToUTF8String(mparams["charset"], bodyBytes)
					if err != nil && newStr == "" {
						return nil, err
					} else {
						if err != nil {
							gerr = err
						}
						mimeMsg.Text = newStr
					}
				} else if mediatype == "text/html" { // charset is empty, look in html body for charset
					charset, err := charsetFromHTMLString(mimeMsg.Text)

					if charset != "" && err == nil {
						newStr, err := ConvertToUTF8String(charset, bodyBytes)
						if err != nil && newStr == "" {
							return nil, err
						} else {
							if err != nil {
								gerr = err
							}
							mimeMsg.Text = newStr
						}
					}
				}
				if mediatype == "text/html" {
					mimeMsg.HTML = mimeMsg.Text
					// Empty Text will trigger html2text conversion below
					mimeMsg.Text = ""
				}
			}
		}
	} else {
		// Parse top-level multipart
		ctype := mailMsg.Header.Get("Content-Type")
		mediatype, params, err := mime.ParseMediaType(ctype)
		if err != nil && mime.IsOkPMTError(err) != nil {
			return nil, fmt.Errorf("Unable to parse media type: %v", err)
		}
		if !strings.HasPrefix(mediatype, "multipart/") {
			return nil, fmt.Errorf("Unknown mediatype: %v", mediatype)
		}
		boundary := params["boundary"]
		if boundary == "" {
			return nil, fmt.Errorf("Unable to locate boundary param in Content-Type header")
		}
		// Root Node of our tree
		root := NewMIMEPart(nil, mediatype)
		mimeMsg.Root = root
		err = parseParts(root, mailMsg.Body, boundary)
		if err != nil {
			return nil, err
		}

		// Locate text body
		if mediatype == "multipart/altern" {
			log.Println("surPrise: should not be here go.enmime/mime.go")
			match := BreadthMatchFirst(root, func(p MIMEPart) bool {
				return p.ContentType() == "text/plain" && p.Disposition() != "attachment"
			})
			if match != nil {
				if match.Charset() != "" {
					newStr, err := ConvertToUTF8String(match.Charset(), match.Content())
					if err != nil {
						if newStr == "" {
							return nil, err
						} else {
							gerr = err
						}
					}
					mimeMsg.Text += newStr
				} else {
					mimeMsg.Text += string(match.Content())
				}
			}
		} else {
			// multipart is of a mixed type
			match := DepthMatchAll(root, func(p MIMEPart) bool {
				return p.ContentType() == "text/plain" && p.Disposition() != "attachment"
			})
			for i, m := range match {
				if i > 0 {
					mimeMsg.Text += "\n--\n"
				}
				if m.Charset() != "" {
					newStr, err := ConvertToUTF8String(m.Charset(), m.Content())
					if err != nil {
						if newStr == "" {
							return nil, err
						} else {
							gerr = err
						}
					}
					mimeMsg.Text += newStr
				} else {
					mimeMsg.Text += string(m.Content())
				}
			}
		}

		// Locate HTML body
		match := BreadthMatchFirst(root, func(p MIMEPart) bool {
			return p.ContentType() == "text/html" && p.Disposition() != "attachment"
		})
		if match != nil {
			if match.Charset() != "" {
				newStr, err := ConvertToUTF8String(match.Charset(), match.Content())
				if err != nil {
					if newStr == "" {
						return nil, err
					} else {
						gerr = err
					}
				}
				mimeMsg.HTML += newStr
			} else {
				mimeMsg.HTML = string(match.Content())
			}

		}

		// Locate attachments
		mimeMsg.Attachments = BreadthMatchAll(root, func(p MIMEPart) bool {
			return p.Disposition() == "attachment" || p.ContentType() == "application/octet-stream"
		})

		// Locate inlines
		mimeMsg.Inlines = BreadthMatchAll(root, func(p MIMEPart) bool {
			return p.Disposition() == "inline"
		})

		// Locate others parts not handled in "Attachments" and "inlines"
		mimeMsg.OtherParts = BreadthMatchAll(root, func(p MIMEPart) bool {
			if strings.HasPrefix(p.ContentType(), "multipart/") {
				return false
			}

			if p.Disposition() != "" {
				return false
			}

			if p.ContentType() == "application/octet-stream" {
				return false
			}

			return p.ContentType() != "text/plain" && p.ContentType() != "text/html"
		})
	}

	// Down-convert HTML to text if necessary
	if mimeMsg.Text == "" && mimeMsg.HTML != "" {
		mimeMsg.IsTextFromHTML = true
		var err error
		if mimeMsg.Text, err = html2text.FromString(mimeMsg.HTML); err != nil {
			// Fail gently
			mimeMsg.Text = ""
			return mimeMsg, err
		}
	}
	return mimeMsg, gerr
}

// Process the specified header for RFC 2047 encoded words and return the result
func (m *MIMEBody) GetHeader(name string) string {
	return DecodeHeader(m.header.Get(name))
}

// Return AddressList with RFC 2047 encoded encoded names.
func (m *MIMEBody) AddressList(key string) ([]*mail.Address, error) {
	isAddrHeader := false
	for _, hkey := range AddressHeaders {
		if strings.ToLower(hkey) == strings.ToLower(key) {
			isAddrHeader = true
			break
		}
	}
	if !isAddrHeader {
		return nil, fmt.Errorf("%s is not address header", key)
	}

	str := DecodeToUTF8Base64Header(m.header.Get(key))
	if str == "" {
		return nil, mail.ErrHeaderNotPresent
	}
	ret, err := mail.ParseAddressList(str)
	if err != nil {
		return nil, err
	}
	return ret, nil
}
<|endoftext|>


<commit_before>// vim: tabstop=2 shiftwidth=2

package main

import (
	"os"
	"os/exec"
	"strings"
	"net/mail"
	"bytes"
	"fmt"
	"path"
	"encoding/hex"
	"net/smtp"
	"io/ioutil"
	"crypto/sha256"
)

func assemble(msg mail.Message) []byte {
	buf := new(bytes.Buffer)
	for h := range msg.Header {
		buf.WriteString(h + ": " + msg.Header.Get(h) + "\n")
	}
	buf.WriteString("\n")
	buf.ReadFrom(msg.Body)
	return buf.Bytes()
}

// testMail vets outbound messages to final recipients
func testMail(b []byte) (sendTo []string, err error) {
	f := bytes.NewReader(b)
	msg, err := mail.ReadMessage(f)
	if err != nil {
		Info.Printf("Outbound read failure: %s", err)
		return
	}
	sendTo = append(sendTo, headToAddy(msg.Header, "To")...)
	sendTo = append(sendTo, headToAddy(msg.Header, "Cc")...)
	return
}

// headToAddy parses a header containing email addresses
func headToAddy(h mail.Header, header string) (addys []string) {
	_, exists := h[header]
	if ! exists {
		return
	}
	addyList, err := h.AddressList(header)
	if err != nil {
		Warn.Printf("Failed to parse header: %s", header)
	}
	for _, addy := range addyList {
		addys = append(addys, addy.Address)
	}
	return
}

// splitAddress splits an email address into its component parts
func splitAddress(addy string) (name, domain string, err error) {
	// Email addresses must have '@' signs in them.
	if ! strings.Contains(addy, "@") {
		err = fmt.Errorf("%s: Email address contains no '@'", addy)
		return
	}
	components := strings.Split(addy, "@")
	if len(components) != 2 {
		err = fmt.Errorf("%s: Malformed email address", addy)
		return
	}
	name = components[0]
	domain = components[1]
	return
}

// Read a file from the outbound pool and mail it
func mailPoolFile(filename string) error {
	var err error
	var payload []byte
	payload, err = ioutil.ReadFile(filename)
	if err != nil {
		Error.Printf("Failed to read file for mailing: %s", err)
		return err
	}
	// Split recipients from payload
	unjoined := bytes.SplitN(payload, []byte{0}, 2)
	if len(unjoined) != 2 {
		Warn.Println("Failed to split recipients from payload")
		return nil
	}
	sendTo := strings.Split(string(unjoined[0]), ",")
	return mailBytes(unjoined[1], sendTo)
}

// Mail a byte payload to a given address
func mailBytes(payload []byte, sendTo []string) (err error) {
	// Test if the message is destined for the local remailer
	Trace.Printf("Message recipient is: %s", sendTo)
	if cfg.Mail.Outfile {
		var f *os.File
		digest := sha256.New()
		digest.Write(payload)
		filename := "outfile-" + hex.EncodeToString(digest.Sum(nil)[:16])
		f, err = os.Create(path.Join(cfg.Files.Pooldir, filename))
		defer f.Close()
		_, err = f.WriteString(string(payload))
		if err != nil {
			Warn.Printf("Outfile write failed: %s\n", err)
			return
		}
	} else if cfg.Mail.Pipe != "" {
		execSend(payload, cfg.Mail.Pipe)
	} else if cfg.Mail.Sendmail {
		err = sendmail(payload, sendTo)
		if err != nil {
			Warn.Println("Sendmail failed")
			return
		}
	} else {
		err = SMTPRelay(payload, sendTo)
		if err != nil {
			Warn.Println("SMTP relay failed")
			return
		}
	}
	return
}

// Pipe mail to an external command (E.g. sendmail -t)
func execSend(payload []byte, execCmd string) {
	sendmail := new(exec.Cmd)
	sendmail.Args = strings.Fields(execCmd)
	sendmail.Path = sendmail.Args[0]

	stdin, err := sendmail.StdinPipe()
	if err != nil {
		panic(err)
	}
	defer stdin.Close()
	sendmail.Stdout = os.Stdout
	sendmail.Stderr = os.Stderr
	err = sendmail.Start()
	if err != nil {
		panic(err)
	}
	stdin.Write(payload)
	stdin.Close()
	err = sendmail.Wait()
	if err != nil {
		Warn.Printf("%s: %s", execCmd, err)
	}
}

func SMTPRelay(payload []byte, sendTo []string) (err error) {
	c, err := smtp.Dial(
		fmt.Sprintf("%s:%d", cfg.Mail.SMTPRelay, cfg.Mail.SMTPPort))
	if err != nil {
		Warn.Println(err)
		return
	}
	err = c.Mail(cfg.Mail.EnvelopeSender)
	if err != nil {
		Warn.Println(err)
		return
	}
	for _, s := range sendTo {
		err = c.Rcpt(s)
		if err != nil {
			Warn.Println(err)
			continue
		}
		wc, err := c.Data()
		if err != nil {
			Warn.Println(err)
			continue
		}
		_, err = fmt.Fprintf(wc, string(payload))
		if err != nil {
			Warn.Println(err)
			continue
		}
		err = wc.Close()
		if err != nil {
			Warn.Println(err)
			continue
		}
	}
	err = c.Quit()
	if err != nil {
		Warn.Println(err)
		return
	}
	return
}

// sendmail invokes go's sendmail method
func sendmail(payload []byte, sendTo []string) (err error) {
	auth := smtp.PlainAuth(
		"",
		cfg.Mail.SMTPUsername,
		cfg.Mail.SMTPPassword,
		cfg.Mail.SMTPRelay)
	relay := fmt.Sprintf("%s:%d", cfg.Mail.SMTPRelay, cfg.Mail.SMTPPort)
	err = smtp.SendMail(relay, auth, cfg.Mail.EnvelopeSender, sendTo, payload)
	if err != nil {
		Warn.Println(err)
		return
	}
	return
}
<commit_msg>Use a random filename for outfile messages<commit_after>// vim: tabstop=2 shiftwidth=2

package main

import (
	"os"
	"os/exec"
	"strings"
	"net/mail"
	"bytes"
	"fmt"
	"path"
	"net/smtp"
	"io/ioutil"
)

func assemble(msg mail.Message) []byte {
	buf := new(bytes.Buffer)
	for h := range msg.Header {
		buf.WriteString(h + ": " + msg.Header.Get(h) + "\n")
	}
	buf.WriteString("\n")
	buf.ReadFrom(msg.Body)
	return buf.Bytes()
}

// testMail vets outbound messages to final recipients
func testMail(b []byte) (sendTo []string, err error) {
	f := bytes.NewReader(b)
	msg, err := mail.ReadMessage(f)
	if err != nil {
		Info.Printf("Outbound read failure: %s", err)
		return
	}
	sendTo = append(sendTo, headToAddy(msg.Header, "To")...)
	sendTo = append(sendTo, headToAddy(msg.Header, "Cc")...)
	return
}

// headToAddy parses a header containing email addresses
func headToAddy(h mail.Header, header string) (addys []string) {
	_, exists := h[header]
	if ! exists {
		return
	}
	addyList, err := h.AddressList(header)
	if err != nil {
		Warn.Printf("Failed to parse header: %s", header)
	}
	for _, addy := range addyList {
		addys = append(addys, addy.Address)
	}
	return
}

// splitAddress splits an email address into its component parts
func splitAddress(addy string) (name, domain string, err error) {
	// Email addresses must have '@' signs in them.
	if ! strings.Contains(addy, "@") {
		err = fmt.Errorf("%s: Email address contains no '@'", addy)
		return
	}
	components := strings.Split(addy, "@")
	if len(components) != 2 {
		err = fmt.Errorf("%s: Malformed email address", addy)
		return
	}
	name = components[0]
	domain = components[1]
	return
}

// Read a file from the outbound pool and mail it
func mailPoolFile(filename string) error {
	var err error
	var payload []byte
	payload, err = ioutil.ReadFile(filename)
	if err != nil {
		Error.Printf("Failed to read file for mailing: %s", err)
		return err
	}
	// Split recipients from payload
	unjoined := bytes.SplitN(payload, []byte{0}, 2)
	if len(unjoined) != 2 {
		Warn.Println("Failed to split recipients from payload")
		return nil
	}
	sendTo := strings.Split(string(unjoined[0]), ",")
	return mailBytes(unjoined[1], sendTo)
}

// Mail a byte payload to a given address
func mailBytes(payload []byte, sendTo []string) (err error) {
	// Test if the message is destined for the local remailer
	Trace.Printf("Message recipients are: %s", strings.Join(sendTo, ","))
	if cfg.Mail.Outfile {
		var f *os.File
		filename := randPoolFilename("outfile-")
		f, err = os.Create(path.Join(cfg.Files.Pooldir, filename))
		defer f.Close()
		_, err = f.WriteString(string(payload))
		if err != nil {
			Warn.Printf("Outfile write failed: %s\n", err)
			return
		}
	} else if cfg.Mail.Pipe != "" {
		execSend(payload, cfg.Mail.Pipe)
	} else if cfg.Mail.Sendmail {
		err = sendmail(payload, sendTo)
		if err != nil {
			Warn.Println("Sendmail failed")
			return
		}
	} else {
		err = SMTPRelay(payload, sendTo)
		if err != nil {
			Warn.Println("SMTP relay failed")
			return
		}
	}
	return
}

// Pipe mail to an external command (E.g. sendmail -t)
func execSend(payload []byte, execCmd string) {
	sendmail := new(exec.Cmd)
	sendmail.Args = strings.Fields(execCmd)
	sendmail.Path = sendmail.Args[0]

	stdin, err := sendmail.StdinPipe()
	if err != nil {
		panic(err)
	}
	defer stdin.Close()
	sendmail.Stdout = os.Stdout
	sendmail.Stderr = os.Stderr
	err = sendmail.Start()
	if err != nil {
		panic(err)
	}
	stdin.Write(payload)
	stdin.Close()
	err = sendmail.Wait()
	if err != nil {
		Warn.Printf("%s: %s", execCmd, err)
	}
}

func SMTPRelay(payload []byte, sendTo []string) (err error) {
	c, err := smtp.Dial(
		fmt.Sprintf("%s:%d", cfg.Mail.SMTPRelay, cfg.Mail.SMTPPort))
	if err != nil {
		Warn.Println(err)
		return
	}
	err = c.Mail(cfg.Mail.EnvelopeSender)
	if err != nil {
		Warn.Println(err)
		return
	}
	for _, s := range sendTo {
		err = c.Rcpt(s)
		if err != nil {
			Warn.Println(err)
			continue
		}
		wc, err := c.Data()
		if err != nil {
			Warn.Println(err)
			continue
		}
		_, err = fmt.Fprintf(wc, string(payload))
		if err != nil {
			Warn.Println(err)
			continue
		}
		err = wc.Close()
		if err != nil {
			Warn.Println(err)
			continue
		}
	}
	err = c.Quit()
	if err != nil {
		Warn.Println(err)
		return
	}
	return
}

// sendmail invokes go's sendmail method
func sendmail(payload []byte, sendTo []string) (err error) {
	auth := smtp.PlainAuth(
		"",
		cfg.Mail.SMTPUsername,
		cfg.Mail.SMTPPassword,
		cfg.Mail.SMTPRelay)
	relay := fmt.Sprintf("%s:%d", cfg.Mail.SMTPRelay, cfg.Mail.SMTPPort)
	err = smtp.SendMail(relay, auth, cfg.Mail.EnvelopeSender, sendTo, payload)
	if err != nil {
		Warn.Println(err)
		return
	}
	return
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package mapper

import (
	"github.com/ernestio/aws-definition-mapper/definition"
	"github.com/ernestio/aws-definition-mapper/output"
)

// ConvertPayload will build an FSMMessage based on an input definition
func ConvertPayload(p *definition.Payload) *output.FSMMessage {
	m := output.FSMMessage{
		ID:          p.ServiceID,
		Service:     p.ServiceID,
		ServiceName: p.Service.Name,
		ClientName:  p.Client.Name,
		Type:        p.Datacenter.Type,
	}

	// Map datacenters
	m.Datacenters.Items = MapDatacenters(p.Datacenter)

	// Map VPCs
	m.VPCs.Items = MapVPCs(p)

	// Map networks
	m.Networks.Items = MapNetworks(p.Service)

	// Map instances
	m.Instances.Items = MapInstances(p.Service)

	// Map firewalls
	m.Firewalls.Items = MapSecurityGroups(p.Service)

	// Map nats/port forwarding
	m.Nats.Items = MapNats(p.Service)

	// Map ELB's
	m.ELBs.Items = MapELBs(p.Service)

	return &m
}

// MapProviderData will map any information generated by a provider that is not
// deductible from the input definition
func MapProviderData(m, om *output.FSMMessage) {
	// Map network ID's
	for i, network := range m.Networks.Items {
		nw := om.FindNetwork(network.Name)
		if nw != nil {
			m.Networks.Items[i].NetworkAWSID = nw.NetworkAWSID
		}
	}

	// Map instance ID's
	for i, instance := range m.Instances.Items {
		in := om.FindInstance(instance.Name)
		if in != nil {
			m.Instances.Items[i].InstanceAWSID = in.InstanceAWSID
			m.Instances.Items[i].PublicIP = in.PublicIP
		}
	}

	// Map firewall ID's
	for i, firewall := range m.Firewalls.Items {
		fw := om.FindFirewall(firewall.Name)
		if fw != nil {
			m.Firewalls.Items[i].SecurityGroupAWSID = fw.SecurityGroupAWSID
		}
	}

	// Map nat ID's
	for i, nat := range m.Nats.Items {
		nt := om.FindNat(nat.Name)
		if nt != nil {
			m.Nats.Items[i].NatGatewayAWSID = nt.NatGatewayAWSID
		}
	}

	// Map elb data
	for i, elb := range m.ELBs.Items {
		lb := om.FindELB(elb.Name)
		if lb != nil {
			m.ELBs.Items[i].DNSName = lb.DNSName
		}
	}
}
<commit_msg>mapping az for network<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package mapper

import (
	"github.com/ernestio/aws-definition-mapper/definition"
	"github.com/ernestio/aws-definition-mapper/output"
)

// ConvertPayload will build an FSMMessage based on an input definition
func ConvertPayload(p *definition.Payload) *output.FSMMessage {
	m := output.FSMMessage{
		ID:          p.ServiceID,
		Service:     p.ServiceID,
		ServiceName: p.Service.Name,
		ClientName:  p.Client.Name,
		Type:        p.Datacenter.Type,
	}

	// Map datacenters
	m.Datacenters.Items = MapDatacenters(p.Datacenter)

	// Map VPCs
	m.VPCs.Items = MapVPCs(p)

	// Map networks
	m.Networks.Items = MapNetworks(p.Service)

	// Map instances
	m.Instances.Items = MapInstances(p.Service)

	// Map firewalls
	m.Firewalls.Items = MapSecurityGroups(p.Service)

	// Map nats/port forwarding
	m.Nats.Items = MapNats(p.Service)

	// Map ELB's
	m.ELBs.Items = MapELBs(p.Service)

	return &m
}

// MapProviderData will map any information generated by a provider that is not
// deductible from the input definition
func MapProviderData(m, om *output.FSMMessage) {
	// Map network ID's
	for i, network := range m.Networks.Items {
		nw := om.FindNetwork(network.Name)
		if nw != nil {
			m.Networks.Items[i].NetworkAWSID = nw.NetworkAWSID
			m.Networks.Items[i].AvailabilityZone = nw.AvailabilityZone
		}
	}

	// Map instance ID's
	for i, instance := range m.Instances.Items {
		in := om.FindInstance(instance.Name)
		if in != nil {
			m.Instances.Items[i].InstanceAWSID = in.InstanceAWSID
			m.Instances.Items[i].PublicIP = in.PublicIP
		}
	}

	// Map firewall ID's
	for i, firewall := range m.Firewalls.Items {
		fw := om.FindFirewall(firewall.Name)
		if fw != nil {
			m.Firewalls.Items[i].SecurityGroupAWSID = fw.SecurityGroupAWSID
		}
	}

	// Map nat ID's
	for i, nat := range m.Nats.Items {
		nt := om.FindNat(nat.Name)
		if nt != nil {
			m.Nats.Items[i].NatGatewayAWSID = nt.NatGatewayAWSID
		}
	}

	// Map elb data
	for i, elb := range m.ELBs.Items {
		lb := om.FindELB(elb.Name)
		if lb != nil {
			m.ELBs.Items[i].DNSName = lb.DNSName
		}
	}
}
<|endoftext|>


<commit_before>package margelet_test

import (
	"github.com/Syfaro/telegram-bot-api"
	. "github.com/smartystreets/goconvey/convey"
	"github.com/zhulik/margelet"
	"testing"
	"time"
)

func TestMargelet(t *testing.T) {
	Convey("When given margelet", t, func() {
		m := getMargelet()

		Convey("When adding new message responder", func() {
			m.AddMessageResponder(margelet.EchoResponder{})

			Convey("It should be aded to message responders array", func() {
				So(m.MessageResponders, ShouldNotBeEmpty)
			})
		})

		Convey("When adding new command responder", func() {
			m.AddCommandHandler("/test", margelet.HelpResponder{})

			Convey("It should be aded to command responders hash", func() {
				So(m.CommandResponders, ShouldNotBeEmpty)
			})

		})

		Convey("When adding new session handler", func() {
			m.AddSessionHandler("/test", margelet.SumSession{})

			Convey("It should be aded to command responders array", func() {
				So(m.SessionHandlers, ShouldNotBeEmpty)
			})

		})

		Convey("When sending message", func() {
			_, err := m.Send(tgbotapi.NewMessage(0, "TEST"))

			Convey("Error should be nil", func() {
				So(err, ShouldBeNil)
			})
		})

		Convey("When getting file direct URL", func() {
			_, err := m.GetFileDirectURL("test")

			Convey("Error should be nil", func() {
				So(err, ShouldBeNil)
			})
		})

		Convey("When asking is message sent to me", func() {
			msg := tgbotapi.Message{}
			res := m.IsMessageToMe(msg)
			Convey("It should return false", func() {
				So(res, ShouldBeFalse)
			})
		})

		Convey("Given configured margelet", func() {
			m.AddMessageResponder(margelet.EchoResponder{})
			m.AddSessionHandler("/sum", margelet.SumSession{})

			Convey("When running should handle message without panic", func() {
				go m.Run()

				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "Test"}}
				time.Sleep(10 * time.Millisecond)
				m.Stop()
			})

			Convey("When running should handle command without panic", func() {
				go m.Run()

				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "/test"}}

				time.Sleep(10 * time.Millisecond)
				m.Stop()
			})

			Convey("When running should handle session without panic", func() {
				go m.Run()

				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "/sum"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "10"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "20"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "test"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "end"}}

				time.Sleep(100 * time.Millisecond)
				m.Stop()
			})
		})
	})
}
<commit_msg>Test updated<commit_after>package margelet_test

import (
	"github.com/Syfaro/telegram-bot-api"
	. "github.com/smartystreets/goconvey/convey"
	"github.com/zhulik/margelet"
	"testing"
	"time"
)

func TestMargelet(t *testing.T) {
	Convey("When given margelet", t, func() {
		m := getMargelet()

		Convey("When adding new message responder", func() {
			m.AddMessageResponder(margelet.EchoResponder{})

			Convey("It should be aded to message responders array", func() {
				So(m.MessageResponders, ShouldNotBeEmpty)
			})
		})

		Convey("When adding new command responder", func() {
			m.AddCommandHandler("/test", margelet.HelpResponder{})

			Convey("It should be aded to command responders hash", func() {
				So(m.CommandResponders, ShouldNotBeEmpty)
			})

		})

		Convey("When adding new session handler", func() {
			m.AddSessionHandler("/test", margelet.SumSession{})

			Convey("It should be aded to command responders array", func() {
				So(m.SessionHandlers, ShouldNotBeEmpty)
			})

		})

		Convey("When sending message", func() {
			_, err := m.Send(tgbotapi.NewMessage(0, "TEST"))

			Convey("Error should be nil", func() {
				So(err, ShouldBeNil)
			})
		})

		Convey("When getting file direct URL", func() {
			_, err := m.GetFileDirectURL("test")

			Convey("Error should be nil", func() {
				So(err, ShouldBeNil)
			})
		})

		Convey("When asking is message sent to me", func() {
			msg := tgbotapi.Message{}
			res := m.IsMessageToMe(msg)
			Convey("It should return false", func() {
				So(res, ShouldBeFalse)
			})
		})

		Convey("Given configured margelet", func() {
			m.AddMessageResponder(margelet.EchoResponder{})
			m.AddSessionHandler("/sum", margelet.SumSession{})

			Convey("When running should handle message without panic", func() {
				go m.Run()

				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "Test"}}
				time.Sleep(10 * time.Millisecond)
				m.Stop()
			})

			Convey("When running should handle command without panic", func() {
				go m.Run()

				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "/help"}}

				time.Sleep(10 * time.Millisecond)
				m.Stop()
			})

			Convey("When running should handle session without panic", func() {
				go m.Run()

				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "/sum"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "10"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "20"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "test"}}
				botMock.Updates <- tgbotapi.Update{Message: tgbotapi.Message{Text: "end"}}

				time.Sleep(100 * time.Millisecond)
				m.Stop()
			})
		})
	})
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

func main() {
	fs := flag.NewFlagSet("mesos-exporter", flag.ExitOnError)
	addr := fs.String("addr", ":9110", "Address to listen on")
	masterURL := fs.String("url", "http://master.mesos.:5050/state.json", "Master URL")
	timeout := fs.Duration("timeout", 5*time.Second, "Master polling timeout")

	fs.Parse(os.Args[1:])

	url, err := url.Parse(*masterURL)
	if err != nil {
		log.Fatal(err)
	}

	c := newCollector(url, *timeout)
	if err := prometheus.Register(c); err != nil {
		log.Fatal(err)
	}

	http.Handle("/metrics", prometheus.Handler())
	if err := http.ListenAndServe(*addr, nil); err != nil {
		log.Fatal(err)
	}
}

type collector struct {
	*http.Client
	*url.URL
	metrics map[prometheus.Collector]func(*slave, prometheus.Collector)
}

func newCollector(url *url.URL, timeout time.Duration) *collector {
	labels := []string{"slave"}
	return &collector{
		Client: &http.Client{Timeout: timeout},
		URL:    url,
		metrics: map[prometheus.Collector]func(*slave, prometheus.Collector){
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave CPUs (fractional)",
				Name:      "cpus",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Total.CPUs)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave CPUs (fractional)",
				Name:      "cpus_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Used.CPUs)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave CPUs (fractional)",
				Name:      "cpus_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Unreserved.CPUs)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave memory in MB",
				Name:      "mem",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Total.Mem)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave memory in MB",
				Name:      "mem_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Used.Mem)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave memory in MB",
				Name:      "mem_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Unreserved.Mem)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave disk in MB",
				Name:      "disk",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Total.Disk)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave disk in MB",
				Name:      "disk_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Used.Disk)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave disk in MB",
				Name:      "disk_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Unreserved.Disk)
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave ports",
				Name:      "ports",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				size := s.Total.Ports.size()
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(float64(size))
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave ports",
				Name:      "ports_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				size := s.Used.Ports.size()
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(float64(size))
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave ports",
				Name:      "ports_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(s *slave, c prometheus.Collector) {
				size := s.Unreserved.Ports.size()
				c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(float64(size))
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Slave tasks",
				Name:      "tasks",
				Namespace: "mesos",
				Subsystem: "slave",
			}, nil): func(s *slave, c prometheus.Collector) {
				for _, task := range s.Tasks {
					labels := prometheus.Labels{
						"slave":     s.PID,
						"executor":  task.ExecutorID,
						"name":      task.Name,
						"framework": task.FrameworkID,
					}
					for _, status := range task.Statuses {
						labels["timestamp"] = strconv.FormatFloat(status.Timestamp, 'f', -1, 64)
						labels["state"] = status.State
						c.(*prometheus.GaugeVec).With(labels).Set(float64(len(s.Tasks)))
					}
				}
			},
		},
	}
}

func (c *collector) Collect(ch chan<- prometheus.Metric) {
	res, err := c.Do(&http.Request{Method: "GET", URL: c.URL})
	if err != nil {
		log.Print(err)
		return
	}
	defer res.Body.Close()

	var s state
	if err := json.NewDecoder(res.Body).Decode(&s); err != nil {
		log.Print(err)
		return
	}

	for _, slave := range s.Slaves {
		for c, set := range c.metrics {
			set(&slave, c)
			c.Collect(ch)
		}
	}
}

func (c *collector) Describe(ch chan<- *prometheus.Desc) {
	for metric := range c.metrics {
		metric.Describe(ch)
	}
}

type ranges [][2]uint64

func (rs *ranges) UnmarshalJSON(data []byte) (err error) {
	if data = bytes.Trim(data, `[]"`); len(data) == 0 {
		return nil
	}

	var rng [2]uint64
	for _, r := range bytes.Split(data, []byte(",")) {
		ps := bytes.SplitN(r, []byte("-"), 2)
		if len(ps) != 2 {
			return fmt.Errorf("bad range: %s", r)
		}

		rng[0], err = strconv.ParseUint(string(bytes.TrimSpace(ps[0])), 10, 64)
		if err != nil {
			return err
		}

		rng[1], err = strconv.ParseUint(string(bytes.TrimSpace(ps[1])), 10, 64)
		if err != nil {
			return err
		}

		*rs = append(*rs, rng)
	}

	return nil
}

func (rs ranges) size() uint64 {
	var sz uint64
	for i := range rs {
		sz += 1 + (rs[i][1] - rs[i][0])
	}
	return sz
}

type (
	resources struct {
		CPUs  float64 `json:"cpus"`
		Disk  float64 `json:"disk"`
		Mem   float64 `json:"mem"`
		Ports ranges  `json:"ports"`
	}

	task struct {
		Name        string    `json:"name"`
		ID          string    `json:"id"`
		ExecutorID  string    `json:"executor_id"`
		FrameworkID string    `json:"framework_id"`
		SlaveID     string    `json:"slave_id"`
		State       string    `json:"state"`
		Labels      []string  `json:"labels"`
		Resources   resources `json:"resources"`
		Statuses    []status  `json:"statuses"`
	}

	status struct {
		State     string  `json:"state"`
		Timestamp float64 `json:"timestamp"`
	}

	slave struct {
		PID        string    `json:"pid"`
		Used       resources `json:"used_resources"`
		Unreserved resources `json:"unreserved_resources"`
		Total      resources `json:"resources"`
		Tasks      []task    `json:"task"`
	}

	state struct {
		Slaves []slave `json:"slaves"`
	}
)
<commit_msg>Mesos slave tasks<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

func main() {
	fs := flag.NewFlagSet("mesos-exporter", flag.ExitOnError)
	addr := fs.String("addr", ":9110", "Address to listen on")
	masterURL := fs.String("url", "http://master.mesos.:5050/state.json", "Master URL")
	timeout := fs.Duration("timeout", 5*time.Second, "Master polling timeout")

	fs.Parse(os.Args[1:])

	url, err := url.Parse(*masterURL)
	if err != nil {
		log.Fatal(err)
	}

	c := newCollector(url, *timeout)
	if err := prometheus.Register(c); err != nil {
		log.Fatal(err)
	}

	http.Handle("/metrics", prometheus.Handler())
	if err := http.ListenAndServe(*addr, nil); err != nil {
		log.Fatal(err)
	}
}

type collector struct {
	*http.Client
	*url.URL
	metrics map[prometheus.Collector]func(*state, prometheus.Collector)
}

func newCollector(url *url.URL, timeout time.Duration) *collector {
	labels := []string{"slave"}
	return &collector{
		Client: &http.Client{Timeout: timeout},
		URL:    url,
		metrics: map[prometheus.Collector]func(*state, prometheus.Collector){
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave CPUs (fractional)",
				Name:      "cpus",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Total.CPUs)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave CPUs (fractional)",
				Name:      "cpus_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Used.CPUs)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave CPUs (fractional)",
				Name:      "cpus_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Unreserved.CPUs)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave memory in MB",
				Name:      "mem",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Total.Mem)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave memory in MB",
				Name:      "mem_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Used.Mem)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave memory in MB",
				Name:      "mem_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Unreserved.Mem)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave disk in MB",
				Name:      "disk",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Total.Disk)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave disk in MB",
				Name:      "disk_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Used.Disk)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave disk in MB",
				Name:      "disk_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(s.Unreserved.Disk)
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Total slave ports",
				Name:      "ports",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					size := s.Total.Ports.size()
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(float64(size))
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Used slave ports",
				Name:      "ports_used",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					size := s.Used.Ports.size()
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(float64(size))
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Unreserved slave ports",
				Name:      "ports_unreserved",
				Namespace: "mesos",
				Subsystem: "slave",
			}, labels): func(st *state, c prometheus.Collector) {
				for _, s := range st.Slaves {
					size := s.Unreserved.Ports.size()
					c.(*prometheus.GaugeVec).WithLabelValues(s.PID).Set(float64(size))
				}
			},
			prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Help:      "Framework tasks",
				Name:      "tasks",
				Namespace: "mesos",
				Subsystem: "slave",
			}, []string{"slave", "executor", "name", "framework", "timestamp", "state"}): func(st *state, c prometheus.Collector) {
				for _, f := range st.Frameworks {
					if !f.Active {
						continue
					}
					tasks := f.tasks()
					for _, task := range tasks {
						for _, status := range task.Statuses {
							values := []string{
								task.SlaveID,
								task.ExecutorID,
								task.Name,
								task.FrameworkID,
								strconv.FormatFloat(status.Timestamp, 'f', -1, 64),
								status.State,
							}
							c.(*prometheus.GaugeVec).WithLabelValues(values...).Set(float64(len(tasks)))
						}
					}
				}
			},
		},
	}
}

func (c *collector) Collect(ch chan<- prometheus.Metric) {
	res, err := c.Do(&http.Request{Method: "GET", URL: c.URL})
	if err != nil {
		log.Print(err)
		return
	}
	defer res.Body.Close()

	var s state
	if err := json.NewDecoder(res.Body).Decode(&s); err != nil {
		log.Print(err)
		return
	}

	for c, set := range c.metrics {
		set(&s, c)
		c.Collect(ch)
	}
}

func (c *collector) Describe(ch chan<- *prometheus.Desc) {
	for metric := range c.metrics {
		metric.Describe(ch)
	}
}

type ranges [][2]uint64

func (rs *ranges) UnmarshalJSON(data []byte) (err error) {
	if data = bytes.Trim(data, `[]"`); len(data) == 0 {
		return nil
	}

	var rng [2]uint64
	for _, r := range bytes.Split(data, []byte(",")) {
		ps := bytes.SplitN(r, []byte("-"), 2)
		if len(ps) != 2 {
			return fmt.Errorf("bad range: %s", r)
		}

		rng[0], err = strconv.ParseUint(string(bytes.TrimSpace(ps[0])), 10, 64)
		if err != nil {
			return err
		}

		rng[1], err = strconv.ParseUint(string(bytes.TrimSpace(ps[1])), 10, 64)
		if err != nil {
			return err
		}

		*rs = append(*rs, rng)
	}

	return nil
}

func (rs ranges) size() uint64 {
	var sz uint64
	for i := range rs {
		sz += 1 + (rs[i][1] - rs[i][0])
	}
	return sz
}

type (
	resources struct {
		CPUs  float64 `json:"cpus"`
		Disk  float64 `json:"disk"`
		Mem   float64 `json:"mem"`
		Ports ranges  `json:"ports"`
	}

	task struct {
		Name        string    `json:"name"`
		ID          string    `json:"id"`
		ExecutorID  string    `json:"executor_id"`
		FrameworkID string    `json:"framework_id"`
		SlaveID     string    `json:"slave_id"`
		State       string    `json:"state"`
		Labels      []label   `json:"labels"`
		Resources   resources `json:"resources"`
		Statuses    []status  `json:"statuses"`
	}

	label struct {
		Key   string `json:"key"`
		Value string `json:"value"`
	}

	status struct {
		State     string  `json:"state"`
		Timestamp float64 `json:"timestamp"`
	}

	slave struct {
		PID        string    `json:"pid"`
		Used       resources `json:"used_resources"`
		Unreserved resources `json:"unreserved_resources"`
		Total      resources `json:"resources"`
	}

	framework struct {
		Active    bool   `json:"active"`
		Tasks     []task `json:"tasks"`
		Completed []task `json:"completed_tasks"`
	}

	state struct {
		Slaves     []slave     `json:"slaves"`
		Frameworks []framework `json:"frameworks"`
	}
)

func (f framework) tasks() []task {
	tasks := make([]task, len(f.Tasks)+len(f.Completed))
	tasks = append(tasks, f.Tasks...)
	tasks = append(tasks, f.Completed...)
	return tasks
}
<|endoftext|>


<commit_before>// Copyright 2015 realglobe, 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 main

import (
	"github.com/realglobe-Inc/edo-id-provider/database/account"
	"github.com/realglobe-Inc/edo-id-provider/database/authcode"
	"github.com/realglobe-Inc/edo-id-provider/database/consent"
	"github.com/realglobe-Inc/edo-id-provider/database/coopcode"
	jtidb "github.com/realglobe-Inc/edo-id-provider/database/jti"
	keydb "github.com/realglobe-Inc/edo-id-provider/database/key"
	"github.com/realglobe-Inc/edo-id-provider/database/pairwise"
	"github.com/realglobe-Inc/edo-id-provider/database/sector"
	"github.com/realglobe-Inc/edo-id-provider/database/session"
	"github.com/realglobe-Inc/edo-id-provider/database/token"
	idpdb "github.com/realglobe-Inc/edo-idp-selector/database/idp"
	tadb "github.com/realglobe-Inc/edo-idp-selector/database/ta"
	webdb "github.com/realglobe-Inc/edo-idp-selector/database/web"
	idperr "github.com/realglobe-Inc/edo-idp-selector/error"
	"github.com/realglobe-Inc/edo-lib/driver"
	logutil "github.com/realglobe-Inc/edo-lib/log"
	"github.com/realglobe-Inc/edo-lib/server"
	"github.com/realglobe-Inc/go-lib/erro"
	"github.com/realglobe-Inc/go-lib/rglog"
	"net/http"
	"os"
	"strings"
)

func main() {
	var exitCode = 0
	defer func() {
		if exitCode != 0 {
			os.Exit(exitCode)
		}
	}()

	defer rglog.Flush()

	logutil.InitConsole("github.com/realglobe-Inc")

	param, err := parseParameters(os.Args...)
	if err != nil {
		log.Err(erro.Unwrap(err))
		log.Debug(erro.Wrap(err))
		exitCode = 1
		return
	}

	logutil.SetupConsole("github.com/realglobe-Inc", param.consLv)
	if err := logutil.Setup("github.com/realglobe-Inc", param.logType, param.logLv, param); err != nil {
		log.Err(erro.Unwrap(err))
		log.Debug(erro.Wrap(err))
		exitCode = 1
		return
	}

	if err := serve(param); err != nil {
		log.Err(erro.Unwrap(err))
		log.Debug(erro.Wrap(err))
		exitCode = 1
		return
	}

	log.Info("Shut down")
}

func serve(param *parameters) (err error) {

	// バックエンドの準備。

	redPools := driver.NewRedisPoolSet(param.redTimeout, param.redPoolSize, param.redPoolExpIn)
	defer redPools.Close()
	monPools := driver.NewMongoPoolSet(param.monTimeout)
	defer monPools.Close()

	// 鍵。
	var keyDb keydb.Db
	switch param.keyDbType {
	case "file":
		keyDb = keydb.NewFileDb(param.keyDbPath)
		log.Info("Use keys in directory " + param.keyDbPath)
	case "redis":
		keyDb = keydb.NewRedisCache(keydb.NewFileDb(param.keyDbPath), redPools.Get(param.keyDbAddr), param.keyDbTag, param.keyDbExpIn)
		log.Info("Use keys in directory " + param.keyDbPath + " with redis " + param.keyDbAddr + ": " + param.keyDbTag)
	default:
		return erro.New("invalid key DB type " + param.keyDbType)
	}

	// アカウント情報。
	var acntDb account.Db
	switch param.acntDbType {
	case "mongo":
		pool, err := monPools.Get(param.acntDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		acntDb = account.NewMongoDb(pool, param.acntDbTag, param.acntDbTag2)
		log.Info("Use account info in mongodb " + param.acntDbAddr + ": " + param.acntDbTag + "." + param.acntDbTag2)
	default:
		return erro.New("invalid account DB type " + param.acntDbType)
	}

	// スコープ・属性許可情報。
	var consDb consent.Db
	switch param.consDbType {
	case "memory":
		consDb = consent.NewMemoryDb()
		log.Info("Save consent info in memory")
	case "mongo":
		pool, err := monPools.Get(param.consDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		consDb = consent.NewMongoDb(pool, param.consDbTag, param.consDbTag2)
		log.Info("Save consent info in mongodb " + param.consDbAddr + ": " + param.consDbTag + "." + param.consDbTag2)
	default:
		return erro.New("invalid consent DB type " + param.consDbType)
	}

	// web データ。
	var webDb webdb.Db
	switch param.webDbType {
	case "direct":
		webDb = webdb.NewDirectDb()
		log.Info("Get web data directly")
	case "redis":
		webDb = webdb.NewRedisCache(webdb.NewDirectDb(), redPools.Get(param.webDbAddr), param.webDbTag, param.webDbExpIn)
		log.Info("Get web data with redis " + param.webDbAddr + ": " + param.webDbTag)
	default:
		return erro.New("invalid web data DB type " + param.webDbType)
	}

	// TA 情報。
	var taDb tadb.Db
	switch param.taDbType {
	case "mongo":
		pool, err := monPools.Get(param.taDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		taDb = tadb.NewMongoDb(pool, param.taDbTag, param.taDbTag2, webDb)
		log.Info("Use TA info in mongodb " + param.taDbAddr + ": " + param.taDbTag + "." + param.taDbTag2)
	default:
		return erro.New("invalid TA DB type " + param.taDbType)
	}

	// セクタ固有のアカウント ID の計算に使う情報。
	var sectDb sector.Db
	switch param.sectDbType {
	case "memory":
		sectDb = sector.NewMemoryDb()
		log.Info("Save pairwise account ID calculation info in memory")
	case "mongo":
		pool, err := monPools.Get(param.sectDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		sectDb = sector.NewMongoDb(pool, param.sectDbTag, param.sectDbTag2)
		log.Info("Save pairwise account ID calculation info in mongodb " + param.sectDbAddr + ": " + param.sectDbTag + "." + param.sectDbTag2)
	default:
		return erro.New("invalid pairwise account ID calculation info DB type " + param.sectDbType)
	}

	// セクタ固有のアカウント ID 情報。
	var pwDb pairwise.Db
	switch param.pwDbType {
	case "memory":
		pwDb = pairwise.NewMemoryDb()
		log.Info("Save pairwise account IDs in memory")
	case "mongo":
		pool, err := monPools.Get(param.pwDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		pwDb = pairwise.NewMongoDb(pool, param.pwDbTag, param.pwDbTag2)
		log.Info("Save pairwise account IDs in mongodb " + param.pwDbAddr + ": " + param.pwDbTag + "." + param.pwDbTag2)
	default:
		return erro.New("invalid pairwise account ID DB type " + param.pwDbType)
	}

	// IdP 情報。
	var idpDb idpdb.Db
	switch param.idpDbType {
	case "mongo":
		pool, err := monPools.Get(param.idpDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		idpDb = idpdb.NewMongoDb(pool, param.idpDbTag, param.idpDbTag2, webDb)
		log.Info("Use IdP info in mongodb " + param.idpDbAddr + ": " + param.idpDbTag + "." + param.idpDbTag2)
	default:
		return erro.New("invalid IdP DB type " + param.idpDbType)
	}

	// セッション。
	var sessDb session.Db
	switch param.sessDbType {
	case "memory":
		sessDb = session.NewMemoryDb()
		log.Info("Save sessions in memory")
	case "redis":
		sessDb = session.NewRedisDb(redPools.Get(param.sessDbAddr), param.sessDbTag)
		log.Info("Save sessions in redis " + param.sessDbAddr + ": " + param.sessDbTag)
	default:
		return erro.New("invalid session DB type " + param.sessDbType)
	}

	// 認可コード。
	var acodDb authcode.Db
	switch param.acodDbType {
	case "memory":
		acodDb = authcode.NewMemoryDb()
		log.Info("Save authorization codes in memory")
	case "redis":
		acodDb = authcode.NewRedisDb(redPools.Get(param.acodDbAddr), param.acodDbTag)
		log.Info("Save authorization codes in redis " + param.acodDbAddr + ": " + param.acodDbTag)
	default:
		return erro.New("invalid authorization code DB type " + param.acodDbType)
	}

	// アクセストークン。
	var tokDb token.Db
	switch param.tokDbType {
	case "memory":
		tokDb = token.NewMemoryDb()
		log.Info("Save access tokens in memory")
	case "redis":
		tokDb = token.NewRedisDb(redPools.Get(param.tokDbAddr), param.tokDbTag)
		log.Info("Save access tokens in redis " + param.tokDbAddr + ": " + param.tokDbTag)
	default:
		return erro.New("invalid access token DB type " + param.tokDbType)
	}

	// 仲介コード。
	var ccodDb coopcode.Db
	switch param.ccodDbType {
	case "memory":
		ccodDb = coopcode.NewMemoryDb()
		log.Info("Save cooperation codes in memory")
	case "redis":
		ccodDb = coopcode.NewRedisDb(redPools.Get(param.ccodDbAddr), param.ccodDbTag)
		log.Info("Save cooperation codes in redis " + param.ccodDbAddr + ": " + param.ccodDbTag)
	default:
		return erro.New("invalid cooperation code DB type " + param.ccodDbType)
	}

	// JWT の ID。
	var jtiDb jtidb.Db
	switch param.jtiDbType {
	case "memory":
		jtiDb = jtidb.NewMemoryDb()
		log.Info("Save JWT IDs in memory")
	case "redis":
		jtiDb = jtidb.NewRedisDb(redPools.Get(param.jtiDbAddr), param.jtiDbTag)
		log.Info("Save JWT IDs in redis " + param.jtiDbAddr + ": " + param.jtiDbTag)
	default:
		return erro.New("invalid JWT ID DB type " + param.jtiDbType)
	}

	sys := &system{
		param.selfId,
		param.sigAlg,
		param.sigKid,

		param.pathTok,
		param.pathTa,
		param.pathSelUi,
		param.pathLginUi,
		param.pathConsUi,
		param.pathErrUi,

		param.pwSaltLen,
		param.sessLen,
		param.sessExpIn,
		param.sessRefDelay,
		param.sessDbExpIn,
		param.acodLen,
		param.acodExpIn,
		param.acodDbExpIn,
		param.tokLen,
		param.tokExpIn,
		param.tokDbExpIn,
		param.ccodLen,
		param.ccodExpIn,
		param.ccodDbExpIn,
		param.jtiLen,
		param.jtiExpIn,
		param.jtiDbExpIn,
		param.ticLen,

		keyDb,
		webDb,
		acntDb,
		consDb,
		taDb,
		sectDb,
		pwDb,
		idpDb,
		sessDb,
		acodDb,
		tokDb,
		ccodDb,
		jtiDb,

		param.cookPath,
		param.cookSec,
	}

	// バックエンドの準備完了。

	s := server.NewStopper()
	defer func() {
		// 処理の終了待ち。
		s.Lock()
		defer s.Unlock()
		for s.Stopped() {
			s.Wait()
		}
	}()

	mux := http.NewServeMux()

	mux.HandleFunc("/", panicErrorWrapper(s, func(w http.ResponseWriter, r *http.Request) error {
		return erro.Wrap(idperr.New(idperr.Invalid_request, "invalid endpoint", http.StatusNotFound, nil))
	}))
	mux.HandleFunc(param.pathOk, panicErrorWrapper(s, func(w http.ResponseWriter, r *http.Request) error {
		return nil
	}))
	mux.HandleFunc(param.pathAuth, panicErrorWrapper(s, sys.authPage))
	mux.HandleFunc(param.pathSel, panicErrorWrapper(s, sys.selectPage))
	mux.HandleFunc(param.pathLgin, panicErrorWrapper(s, sys.lginPage))
	mux.HandleFunc(param.pathCons, panicErrorWrapper(s, sys.consentPage))
	mux.HandleFunc(param.pathTa, panicErrorWrapper(s, sys.taApiHandler().ServeHTTP))
	mux.HandleFunc(param.pathTok, panicErrorWrapper(s, sys.tokenApi))
	mux.HandleFunc(param.pathAcnt, panicErrorWrapper(s, sys.accountApi))
	mux.HandleFunc(param.pathCoopFr, panicErrorWrapper(s, sys.cooperateFromApi))
	mux.HandleFunc(param.pathCoopTo, panicErrorWrapper(s, sys.cooperateToApi))
	if param.uiDir != "" {
		// ファイル配信も自前でやる。
		pathUi := strings.TrimRight(param.pathUi, "/") + "/"
		filer := http.StripPrefix(pathUi, http.FileServer(http.Dir(param.uiDir)))
		mux.Handle(pathUi, filer)
	}

	return server.Serve(param, mux)
}
<commit_msg>自分の鍵の redis キャッシュのキーに自分の ID を付けるようにした<commit_after>// Copyright 2015 realglobe, 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 main

import (
	"github.com/realglobe-Inc/edo-id-provider/database/account"
	"github.com/realglobe-Inc/edo-id-provider/database/authcode"
	"github.com/realglobe-Inc/edo-id-provider/database/consent"
	"github.com/realglobe-Inc/edo-id-provider/database/coopcode"
	jtidb "github.com/realglobe-Inc/edo-id-provider/database/jti"
	keydb "github.com/realglobe-Inc/edo-id-provider/database/key"
	"github.com/realglobe-Inc/edo-id-provider/database/pairwise"
	"github.com/realglobe-Inc/edo-id-provider/database/sector"
	"github.com/realglobe-Inc/edo-id-provider/database/session"
	"github.com/realglobe-Inc/edo-id-provider/database/token"
	idpdb "github.com/realglobe-Inc/edo-idp-selector/database/idp"
	tadb "github.com/realglobe-Inc/edo-idp-selector/database/ta"
	webdb "github.com/realglobe-Inc/edo-idp-selector/database/web"
	idperr "github.com/realglobe-Inc/edo-idp-selector/error"
	"github.com/realglobe-Inc/edo-lib/driver"
	logutil "github.com/realglobe-Inc/edo-lib/log"
	"github.com/realglobe-Inc/edo-lib/server"
	"github.com/realglobe-Inc/go-lib/erro"
	"github.com/realglobe-Inc/go-lib/rglog"
	"net/http"
	"os"
	"strings"
)

func main() {
	var exitCode = 0
	defer func() {
		if exitCode != 0 {
			os.Exit(exitCode)
		}
	}()

	defer rglog.Flush()

	logutil.InitConsole("github.com/realglobe-Inc")

	param, err := parseParameters(os.Args...)
	if err != nil {
		log.Err(erro.Unwrap(err))
		log.Debug(erro.Wrap(err))
		exitCode = 1
		return
	}

	logutil.SetupConsole("github.com/realglobe-Inc", param.consLv)
	if err := logutil.Setup("github.com/realglobe-Inc", param.logType, param.logLv, param); err != nil {
		log.Err(erro.Unwrap(err))
		log.Debug(erro.Wrap(err))
		exitCode = 1
		return
	}

	if err := serve(param); err != nil {
		log.Err(erro.Unwrap(err))
		log.Debug(erro.Wrap(err))
		exitCode = 1
		return
	}

	log.Info("Shut down")
}

func serve(param *parameters) (err error) {

	// バックエンドの準備。

	redPools := driver.NewRedisPoolSet(param.redTimeout, param.redPoolSize, param.redPoolExpIn)
	defer redPools.Close()
	monPools := driver.NewMongoPoolSet(param.monTimeout)
	defer monPools.Close()

	// 鍵。
	var keyDb keydb.Db
	switch param.keyDbType {
	case "file":
		keyDb = keydb.NewFileDb(param.keyDbPath)
		log.Info("Use keys in directory " + param.keyDbPath)
	case "redis":
		keyDb = keydb.NewRedisCache(keydb.NewFileDb(param.keyDbPath), redPools.Get(param.keyDbAddr), param.keyDbTag+"."+param.selfId, param.keyDbExpIn)
		log.Info("Use keys in directory " + param.keyDbPath + " with redis " + param.keyDbAddr + ": " + param.keyDbTag + "." + param.selfId)
	default:
		return erro.New("invalid key DB type " + param.keyDbType)
	}

	// アカウント情報。
	var acntDb account.Db
	switch param.acntDbType {
	case "mongo":
		pool, err := monPools.Get(param.acntDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		acntDb = account.NewMongoDb(pool, param.acntDbTag, param.acntDbTag2)
		log.Info("Use account info in mongodb " + param.acntDbAddr + ": " + param.acntDbTag + "." + param.acntDbTag2)
	default:
		return erro.New("invalid account DB type " + param.acntDbType)
	}

	// スコープ・属性許可情報。
	var consDb consent.Db
	switch param.consDbType {
	case "memory":
		consDb = consent.NewMemoryDb()
		log.Info("Save consent info in memory")
	case "mongo":
		pool, err := monPools.Get(param.consDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		consDb = consent.NewMongoDb(pool, param.consDbTag, param.consDbTag2)
		log.Info("Save consent info in mongodb " + param.consDbAddr + ": " + param.consDbTag + "." + param.consDbTag2)
	default:
		return erro.New("invalid consent DB type " + param.consDbType)
	}

	// web データ。
	var webDb webdb.Db
	switch param.webDbType {
	case "direct":
		webDb = webdb.NewDirectDb()
		log.Info("Get web data directly")
	case "redis":
		webDb = webdb.NewRedisCache(webdb.NewDirectDb(), redPools.Get(param.webDbAddr), param.webDbTag, param.webDbExpIn)
		log.Info("Get web data with redis " + param.webDbAddr + ": " + param.webDbTag)
	default:
		return erro.New("invalid web data DB type " + param.webDbType)
	}

	// TA 情報。
	var taDb tadb.Db
	switch param.taDbType {
	case "mongo":
		pool, err := monPools.Get(param.taDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		taDb = tadb.NewMongoDb(pool, param.taDbTag, param.taDbTag2, webDb)
		log.Info("Use TA info in mongodb " + param.taDbAddr + ": " + param.taDbTag + "." + param.taDbTag2)
	default:
		return erro.New("invalid TA DB type " + param.taDbType)
	}

	// セクタ固有のアカウント ID の計算に使う情報。
	var sectDb sector.Db
	switch param.sectDbType {
	case "memory":
		sectDb = sector.NewMemoryDb()
		log.Info("Save pairwise account ID calculation info in memory")
	case "mongo":
		pool, err := monPools.Get(param.sectDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		sectDb = sector.NewMongoDb(pool, param.sectDbTag, param.sectDbTag2)
		log.Info("Save pairwise account ID calculation info in mongodb " + param.sectDbAddr + ": " + param.sectDbTag + "." + param.sectDbTag2)
	default:
		return erro.New("invalid pairwise account ID calculation info DB type " + param.sectDbType)
	}

	// セクタ固有のアカウント ID 情報。
	var pwDb pairwise.Db
	switch param.pwDbType {
	case "memory":
		pwDb = pairwise.NewMemoryDb()
		log.Info("Save pairwise account IDs in memory")
	case "mongo":
		pool, err := monPools.Get(param.pwDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		pwDb = pairwise.NewMongoDb(pool, param.pwDbTag, param.pwDbTag2)
		log.Info("Save pairwise account IDs in mongodb " + param.pwDbAddr + ": " + param.pwDbTag + "." + param.pwDbTag2)
	default:
		return erro.New("invalid pairwise account ID DB type " + param.pwDbType)
	}

	// IdP 情報。
	var idpDb idpdb.Db
	switch param.idpDbType {
	case "mongo":
		pool, err := monPools.Get(param.idpDbAddr)
		if err != nil {
			return erro.Wrap(err)
		}
		idpDb = idpdb.NewMongoDb(pool, param.idpDbTag, param.idpDbTag2, webDb)
		log.Info("Use IdP info in mongodb " + param.idpDbAddr + ": " + param.idpDbTag + "." + param.idpDbTag2)
	default:
		return erro.New("invalid IdP DB type " + param.idpDbType)
	}

	// セッション。
	var sessDb session.Db
	switch param.sessDbType {
	case "memory":
		sessDb = session.NewMemoryDb()
		log.Info("Save sessions in memory")
	case "redis":
		sessDb = session.NewRedisDb(redPools.Get(param.sessDbAddr), param.sessDbTag)
		log.Info("Save sessions in redis " + param.sessDbAddr + ": " + param.sessDbTag)
	default:
		return erro.New("invalid session DB type " + param.sessDbType)
	}

	// 認可コード。
	var acodDb authcode.Db
	switch param.acodDbType {
	case "memory":
		acodDb = authcode.NewMemoryDb()
		log.Info("Save authorization codes in memory")
	case "redis":
		acodDb = authcode.NewRedisDb(redPools.Get(param.acodDbAddr), param.acodDbTag)
		log.Info("Save authorization codes in redis " + param.acodDbAddr + ": " + param.acodDbTag)
	default:
		return erro.New("invalid authorization code DB type " + param.acodDbType)
	}

	// アクセストークン。
	var tokDb token.Db
	switch param.tokDbType {
	case "memory":
		tokDb = token.NewMemoryDb()
		log.Info("Save access tokens in memory")
	case "redis":
		tokDb = token.NewRedisDb(redPools.Get(param.tokDbAddr), param.tokDbTag)
		log.Info("Save access tokens in redis " + param.tokDbAddr + ": " + param.tokDbTag)
	default:
		return erro.New("invalid access token DB type " + param.tokDbType)
	}

	// 仲介コード。
	var ccodDb coopcode.Db
	switch param.ccodDbType {
	case "memory":
		ccodDb = coopcode.NewMemoryDb()
		log.Info("Save cooperation codes in memory")
	case "redis":
		ccodDb = coopcode.NewRedisDb(redPools.Get(param.ccodDbAddr), param.ccodDbTag)
		log.Info("Save cooperation codes in redis " + param.ccodDbAddr + ": " + param.ccodDbTag)
	default:
		return erro.New("invalid cooperation code DB type " + param.ccodDbType)
	}

	// JWT の ID。
	var jtiDb jtidb.Db
	switch param.jtiDbType {
	case "memory":
		jtiDb = jtidb.NewMemoryDb()
		log.Info("Save JWT IDs in memory")
	case "redis":
		jtiDb = jtidb.NewRedisDb(redPools.Get(param.jtiDbAddr), param.jtiDbTag)
		log.Info("Save JWT IDs in redis " + param.jtiDbAddr + ": " + param.jtiDbTag)
	default:
		return erro.New("invalid JWT ID DB type " + param.jtiDbType)
	}

	sys := &system{
		param.selfId,
		param.sigAlg,
		param.sigKid,

		param.pathTok,
		param.pathTa,
		param.pathSelUi,
		param.pathLginUi,
		param.pathConsUi,
		param.pathErrUi,

		param.pwSaltLen,
		param.sessLen,
		param.sessExpIn,
		param.sessRefDelay,
		param.sessDbExpIn,
		param.acodLen,
		param.acodExpIn,
		param.acodDbExpIn,
		param.tokLen,
		param.tokExpIn,
		param.tokDbExpIn,
		param.ccodLen,
		param.ccodExpIn,
		param.ccodDbExpIn,
		param.jtiLen,
		param.jtiExpIn,
		param.jtiDbExpIn,
		param.ticLen,

		keyDb,
		webDb,
		acntDb,
		consDb,
		taDb,
		sectDb,
		pwDb,
		idpDb,
		sessDb,
		acodDb,
		tokDb,
		ccodDb,
		jtiDb,

		param.cookPath,
		param.cookSec,
	}

	// バックエンドの準備完了。

	s := server.NewStopper()
	defer func() {
		// 処理の終了待ち。
		s.Lock()
		defer s.Unlock()
		for s.Stopped() {
			s.Wait()
		}
	}()

	mux := http.NewServeMux()

	mux.HandleFunc("/", panicErrorWrapper(s, func(w http.ResponseWriter, r *http.Request) error {
		return erro.Wrap(idperr.New(idperr.Invalid_request, "invalid endpoint", http.StatusNotFound, nil))
	}))
	mux.HandleFunc(param.pathOk, panicErrorWrapper(s, func(w http.ResponseWriter, r *http.Request) error {
		return nil
	}))
	mux.HandleFunc(param.pathAuth, panicErrorWrapper(s, sys.authPage))
	mux.HandleFunc(param.pathSel, panicErrorWrapper(s, sys.selectPage))
	mux.HandleFunc(param.pathLgin, panicErrorWrapper(s, sys.lginPage))
	mux.HandleFunc(param.pathCons, panicErrorWrapper(s, sys.consentPage))
	mux.HandleFunc(param.pathTa, panicErrorWrapper(s, sys.taApiHandler().ServeHTTP))
	mux.HandleFunc(param.pathTok, panicErrorWrapper(s, sys.tokenApi))
	mux.HandleFunc(param.pathAcnt, panicErrorWrapper(s, sys.accountApi))
	mux.HandleFunc(param.pathCoopFr, panicErrorWrapper(s, sys.cooperateFromApi))
	mux.HandleFunc(param.pathCoopTo, panicErrorWrapper(s, sys.cooperateToApi))
	if param.uiDir != "" {
		// ファイル配信も自前でやる。
		pathUi := strings.TrimRight(param.pathUi, "/") + "/"
		filer := http.StripPrefix(pathUi, http.FileServer(http.Dir(param.uiDir)))
		mux.Handle(pathUi, filer)
	}

	return server.Serve(param, mux)
}
<|endoftext|>


<commit_before>package zenbot

import (
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/aybabtme/log"
	"github.com/nlopes/slack"
)

// Slack contains the Slack client and RTM object.
type Slack struct {
	Bot *slack.Client
	RTM *slack.RTM
}

// Config contains all the necessary configs for a
// zenbot instance.
type Config struct {
	Slack            *Slack
	Debug            bool
	Log              *log.Log
	TimeoutDuration  time.Duration
	ChannelWhitelist StringList
}

// A Zen is a zen time period for a user
type Zen struct {
	User, Name, Channel, Reason string
	EndsAt, Timeout             time.Time
}

// A Bot is an instance of zenbot.
type Bot struct {
	Config       *Config
	zens         []*Zen
	zensMutex    sync.RWMutex
	channelNames map[string]string
}

var regexps = struct {
	Zen, ZenArgs, Cancel *regexp.Regexp
}{
	Zen:     regexp.MustCompile(`^\.\/zen`),
	ZenArgs: regexp.MustCompile(`^\.\/zen +t?((?:\d+h)?(?:\d+m)?(?:\d+s)?)(?: (.*)?)$`),
	Cancel:  regexp.MustCompile(`^\.\/zen cancel(?: (.*))?$`),
}

func New(config *Config) *Bot {
	return &Bot{
		Config:       config,
		channelNames: make(map[string]string),
	}
}

// Zen starts listening for Slack messages.
func (b *Bot) Zen() {
	go b.ExpireZens()

	for {
		select {
		case msg := <-b.Config.Slack.RTM.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.MessageEvent:
				go b.handleMessageEvent(msg.Data.(*slack.MessageEvent))
			case *slack.ConnectedEvent:
				b.Config.Log.Info("connected to slack")

				if b.Config.Debug {
					b.Config.Log.KV("info", ev.Info).Info("got slack info")
					b.Config.Log.KV("connections", ev.ConnectionCount).Info("got connection count")
				}
			case *slack.RTMError:
				b.Config.Log.Err(ev).Error("slack rtm error")
			case *slack.InvalidAuthEvent:
				b.Config.Log.Fatal("invalid slack token")
			// user activity events
			case *slack.UserTypingEvent:
				go b.enforceZen(ev.User, "typing")
			case *slack.ReactionAddedEvent:
				go b.enforceZen(ev.User, "using reactjis")
			case *slack.ReactionRemovedEvent:
				go b.enforceZen(ev.User, "using reactjis")
			case *slack.StarRemovedEvent:
				go b.enforceZen(ev.User, "starring messages")
			case *slack.StarAddedEvent:
				go b.enforceZen(ev.User, "starring messages")
			case *slack.PinRemovedEvent:
				go b.enforceZen(ev.User, "pinning messages")
			case *slack.PinAddedEvent:
				go b.enforceZen(ev.User, "pinning messages")
			default:
			}
		}
	}
}

// SendMessage sends a message to a Slack channel.
func (b *Bot) SendMessage(message, channel string) {
	b.Config.Slack.RTM.SendMessage(b.Config.Slack.RTM.NewOutgoingMessage(message, channel))
}

func (b *Bot) handleError(err error, channel string) bool {
	if err == nil {
		return false
	}

	b.Config.Log.Err(err).Error("error")

	b.SendMessage(err.Error(), channel)
	return true
}

func (b *Bot) handleMessageEvent(ev *slack.MessageEvent) {
	if ev.Type != "message" {
		return
	}

	switch {
	case regexps.Cancel.MatchString(ev.Text):
		if b.isWhitelisted(ev) {
			b.cancelZen(ev)
		}
	case regexps.Zen.MatchString(ev.Text):
		if b.isWhitelisted(ev) {
			b.startZen(ev)
		}
	default:
	}
}

func (b *Bot) isWhitelisted(ev *slack.MessageEvent) bool {
	if len(b.Config.ChannelWhitelist) == 0 {
		return true
	}

	name, err := b.getChannelName(ev.Channel)
	if err != nil {
		b.Config.Log.KV("id", ev.Channel).Err(err).Error("could not look up channel name")
		return false
	}

	_, whitelisted := b.Config.ChannelWhitelist[name]
	if b.Config.Debug {
		b.Config.Log.KV("channel", name).Info("channel is not whitelisted, command ignored")
	}
	return whitelisted
}

func (b *Bot) startZen(ev *slack.MessageEvent) {
	match := regexps.ZenArgs.FindStringSubmatch(ev.Text)
	if len(match) == 0 {
		b.SendMessage("Usage: `./zen <duration e.g. 1h30m> [reason - optional]`", ev.Channel)
		return
	}

	durationString, reason := match[1], match[2]
	duration, err := time.ParseDuration(durationString)

	if b.handleError(err, ev.Channel) {
		return
	}

	name, err := b.getUserName(ev.User)
	if b.handleError(err, ev.Channel) {
		return
	}

	zen := &Zen{
		User:    ev.User,
		Name:    name,
		Channel: ev.Channel,
		Reason:  reason,
		EndsAt:  time.Now().Add(duration),
		Timeout: time.Now().Add(5 * time.Second),
	}

	b.zensMutex.Lock()
	b.zens = append(b.zens, zen)
	b.zensMutex.Unlock()

	b.SendMessage(fmt.Sprintf("Added a zen for %s (%s), ends at [%s].", durationString, reason, zen.EndsAt), ev.Channel)
}

func (b *Bot) cancelZen(ev *slack.MessageEvent) {
	match := regexps.Cancel.FindStringSubmatch(ev.Text)
	if len(match) == 0 {
		b.SendMessage("Usage: `./zen cancel [reason - optional]`", ev.Channel)
		return
	}

	var (
		reason  = match[1]
		message = ""
		count   = 0
	)

	b.zensMutex.RLock()
	for i := 0; i < len(b.zens); i++ {
		zen := b.zens[i]
		if zen.User == ev.User && (reason == "" || strings.ToLower(zen.Reason) == strings.ToLower(reason)) {
			b.zensMutex.RUnlock()

			b.zensMutex.Lock()
			b.zens = append(b.zens[:i], b.zens[i+1:]...)
			i--
			b.zensMutex.Unlock()

			message += fmt.Sprintf("(%s) zen canceled.\n", zen.Reason)
			count++

			b.zensMutex.RLock()
		}
	}
	b.zensMutex.RUnlock()

	if message == "" {
		if reason != "" {
			message = "you do not have any such running zens"
		} else {
			message = "you do not have any running zens"
		}
	}
	b.SendMessage(message, ev.Channel)

	if count > 0 {
		b.SendMessage(fmt.Sprintf("<@%s>-%s for canceling %d zens.", ev.User, strings.Repeat("-", count), count), ev.Channel)
	}

}

func (b *Bot) enforceZen(user, activity string) {
	b.zensMutex.RLock()
	defer b.zensMutex.RUnlock()

	for i := 0; i < len(b.zens); i++ {
		zen := b.zens[i]
		if zen.User == user && time.Now().After(zen.Timeout) {
			b.SendMessage(fmt.Sprintf("<@%s>-- for %s during your zen period (%s).", zen.User, activity, zen.Reason), zen.Channel)

			b.zensMutex.RUnlock()
			b.zensMutex.Lock()
			zen.Timeout = time.Now().Add(b.Config.TimeoutDuration)
			b.zensMutex.Unlock()
			b.zensMutex.RLock()
			break
		}
	}
}

// ExpireZens removes zens that have ended
func (b *Bot) ExpireZens() {
	var wg sync.WaitGroup

	for {
		wg.Wait()
		<-time.After(1 * time.Second)

		wg.Add(1)
		go func() {
			b.zensMutex.RLock()

			now := time.Now()
			for i, zen := range b.zens {
				if now.After(zen.EndsAt) {
					b.zensMutex.RUnlock()

					b.zensMutex.Lock()
					b.zens = append(b.zens[:i], b.zens[i+1:]...)
					b.zensMutex.Unlock()

					b.SendMessage(fmt.Sprintf("<@%s>: Be free, for your zen (%s) has ended!", zen.User, zen.Reason), zen.Channel)

					b.zensMutex.RLock()
				}
			}

			b.zensMutex.RUnlock()
			wg.Done()
		}()
	}
}

func (b *Bot) getUserName(id string) (string, error) {
	userInfo, err := b.Config.Slack.Bot.GetUserInfo(id)
	if err != nil {
		return "", err
	}

	return userInfo.Name, nil
}

func (b *Bot) getChannelName(id string) (string, error) {
	if name, ok := b.channelNames[id]; ok {
		return name, nil
	}

	channelInfo, err := b.Config.Slack.Bot.GetChannelInfo(id)
	if err != nil {
		return "", err
	}

	b.channelNames[id] = channelInfo.Name
	return channelInfo.Name, nil
}
<commit_msg>Update time formatting to be more readable (#4)<commit_after>package zenbot

import (
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/aybabtme/log"
	"github.com/nlopes/slack"
)

// Slack contains the Slack client and RTM object.
type Slack struct {
	Bot *slack.Client
	RTM *slack.RTM
}

// Config contains all the necessary configs for a
// zenbot instance.
type Config struct {
	Slack            *Slack
	Debug            bool
	Log              *log.Log
	TimeoutDuration  time.Duration
	ChannelWhitelist StringList
}

// A Zen is a zen time period for a user
type Zen struct {
	User, Name, Channel, Reason string
	EndsAt, Timeout             time.Time
}

// A Bot is an instance of zenbot.
type Bot struct {
	Config       *Config
	zens         []*Zen
	zensMutex    sync.RWMutex
	channelNames map[string]string
}

var regexps = struct {
	Zen, ZenArgs, Cancel *regexp.Regexp
}{
	Zen:     regexp.MustCompile(`^\.\/zen`),
	ZenArgs: regexp.MustCompile(`^\.\/zen +t?((?:\d+h)?(?:\d+m)?(?:\d+s)?)(?: (.*)?)$`),
	Cancel:  regexp.MustCompile(`^\.\/zen cancel(?: (.*))?$`),
}

func New(config *Config) *Bot {
	return &Bot{
		Config:       config,
		channelNames: make(map[string]string),
	}
}

// Zen starts listening for Slack messages.
func (b *Bot) Zen() {
	go b.ExpireZens()

	for {
		select {
		case msg := <-b.Config.Slack.RTM.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.MessageEvent:
				go b.handleMessageEvent(msg.Data.(*slack.MessageEvent))
			case *slack.ConnectedEvent:
				b.Config.Log.Info("connected to slack")

				if b.Config.Debug {
					b.Config.Log.KV("info", ev.Info).Info("got slack info")
					b.Config.Log.KV("connections", ev.ConnectionCount).Info("got connection count")
				}
			case *slack.RTMError:
				b.Config.Log.Err(ev).Error("slack rtm error")
			case *slack.InvalidAuthEvent:
				b.Config.Log.Fatal("invalid slack token")
			// user activity events
			case *slack.UserTypingEvent:
				go b.enforceZen(ev.User, "typing")
			case *slack.ReactionAddedEvent:
				go b.enforceZen(ev.User, "using reactjis")
			case *slack.ReactionRemovedEvent:
				go b.enforceZen(ev.User, "using reactjis")
			case *slack.StarRemovedEvent:
				go b.enforceZen(ev.User, "starring messages")
			case *slack.StarAddedEvent:
				go b.enforceZen(ev.User, "starring messages")
			case *slack.PinRemovedEvent:
				go b.enforceZen(ev.User, "pinning messages")
			case *slack.PinAddedEvent:
				go b.enforceZen(ev.User, "pinning messages")
			default:
			}
		}
	}
}

// SendMessage sends a message to a Slack channel.
func (b *Bot) SendMessage(message, channel string) {
	b.Config.Slack.RTM.SendMessage(b.Config.Slack.RTM.NewOutgoingMessage(message, channel))
}

func (b *Bot) handleError(err error, channel string) bool {
	if err == nil {
		return false
	}

	b.Config.Log.Err(err).Error("error")

	b.SendMessage(err.Error(), channel)
	return true
}

func (b *Bot) handleMessageEvent(ev *slack.MessageEvent) {
	if ev.Type != "message" {
		return
	}

	switch {
	case regexps.Cancel.MatchString(ev.Text):
		if b.isWhitelisted(ev) {
			b.cancelZen(ev)
		}
	case regexps.Zen.MatchString(ev.Text):
		if b.isWhitelisted(ev) {
			b.startZen(ev)
		}
	default:
	}
}

func (b *Bot) isWhitelisted(ev *slack.MessageEvent) bool {
	if len(b.Config.ChannelWhitelist) == 0 {
		return true
	}

	name, err := b.getChannelName(ev.Channel)
	if err != nil {
		b.Config.Log.KV("id", ev.Channel).Err(err).Error("could not look up channel name")
		return false
	}

	_, whitelisted := b.Config.ChannelWhitelist[name]
	if b.Config.Debug {
		b.Config.Log.KV("channel", name).Info("channel is not whitelisted, command ignored")
	}
	return whitelisted
}

func (b *Bot) startZen(ev *slack.MessageEvent) {
	match := regexps.ZenArgs.FindStringSubmatch(ev.Text)
	if len(match) == 0 {
		b.SendMessage("Usage: `./zen <duration e.g. 1h30m> [reason - optional]`", ev.Channel)
		return
	}

	durationString, reason := match[1], match[2]
	duration, err := time.ParseDuration(durationString)

	if b.handleError(err, ev.Channel) {
		return
	}

	name, err := b.getUserName(ev.User)
	if b.handleError(err, ev.Channel) {
		return
	}

	zen := &Zen{
		User:    ev.User,
		Name:    name,
		Channel: ev.Channel,
		Reason:  reason,
		EndsAt:  time.Now().Add(duration),
		Timeout: time.Now().Add(5 * time.Second),
	}

	b.zensMutex.Lock()
	b.zens = append(b.zens, zen)
	b.zensMutex.Unlock()

	b.SendMessage(fmt.Sprintf("Added a zen for %s (%s), ends at %s.", durationString, reason, zen.EndsAt.Format("15:04:05 (_2 Jan, MST)")), ev.Channel)
}

func (b *Bot) cancelZen(ev *slack.MessageEvent) {
	match := regexps.Cancel.FindStringSubmatch(ev.Text)
	if len(match) == 0 {
		b.SendMessage("Usage: `./zen cancel [reason - optional]`", ev.Channel)
		return
	}

	var (
		reason  = match[1]
		message = ""
		count   = 0
	)

	b.zensMutex.RLock()
	for i := 0; i < len(b.zens); i++ {
		zen := b.zens[i]
		if zen.User == ev.User && (reason == "" || strings.ToLower(zen.Reason) == strings.ToLower(reason)) {
			b.zensMutex.RUnlock()

			b.zensMutex.Lock()
			b.zens = append(b.zens[:i], b.zens[i+1:]...)
			i--
			b.zensMutex.Unlock()

			message += fmt.Sprintf("(%s) zen canceled.\n", zen.Reason)
			count++

			b.zensMutex.RLock()
		}
	}
	b.zensMutex.RUnlock()

	if message == "" {
		if reason != "" {
			message = "you do not have any such running zens"
		} else {
			message = "you do not have any running zens"
		}
	}
	b.SendMessage(message, ev.Channel)

	if count > 0 {
		b.SendMessage(fmt.Sprintf("<@%s>-%s for canceling %d zens.", ev.User, strings.Repeat("-", count), count), ev.Channel)
	}

}

func (b *Bot) enforceZen(user, activity string) {
	b.zensMutex.RLock()
	defer b.zensMutex.RUnlock()

	for i := 0; i < len(b.zens); i++ {
		zen := b.zens[i]
		if zen.User == user && time.Now().After(zen.Timeout) {
			b.SendMessage(fmt.Sprintf("<@%s>-- for %s during your zen period (%s).", zen.User, activity, zen.Reason), zen.Channel)

			b.zensMutex.RUnlock()
			b.zensMutex.Lock()
			zen.Timeout = time.Now().Add(b.Config.TimeoutDuration)
			b.zensMutex.Unlock()
			b.zensMutex.RLock()
			break
		}
	}
}

// ExpireZens removes zens that have ended
func (b *Bot) ExpireZens() {
	var wg sync.WaitGroup

	for {
		wg.Wait()
		<-time.After(1 * time.Second)

		wg.Add(1)
		go func() {
			b.zensMutex.RLock()

			now := time.Now()
			for i, zen := range b.zens {
				if now.After(zen.EndsAt) {
					b.zensMutex.RUnlock()

					b.zensMutex.Lock()
					b.zens = append(b.zens[:i], b.zens[i+1:]...)
					b.zensMutex.Unlock()

					b.SendMessage(fmt.Sprintf("<@%s>: Be free, for your zen (%s) has ended!", zen.User, zen.Reason), zen.Channel)

					b.zensMutex.RLock()
				}
			}

			b.zensMutex.RUnlock()
			wg.Done()
		}()
	}
}

func (b *Bot) getUserName(id string) (string, error) {
	userInfo, err := b.Config.Slack.Bot.GetUserInfo(id)
	if err != nil {
		return "", err
	}

	return userInfo.Name, nil
}

func (b *Bot) getChannelName(id string) (string, error) {
	if name, ok := b.channelNames[id]; ok {
		return name, nil
	}

	channelInfo, err := b.Config.Slack.Bot.GetChannelInfo(id)
	if err != nil {
		return "", err
	}

	b.channelNames[id] = channelInfo.Name
	return channelInfo.Name, nil
}
<|endoftext|>