<commit_before>package providers

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/admin/directory/v1"
	"google.golang.org/api/googleapi"
)

type GoogleProvider struct {
	*ProviderData
	RedeemRefreshURL *url.URL
	// GroupValidator is a function that determines if the passed email is in
	// the configured Google group.
	GroupValidator func(string) bool
}

func NewGoogleProvider(p *ProviderData) *GoogleProvider {
	p.ProviderName = "Google"
	if p.LoginURL.String() == "" {
		p.LoginURL = &url.URL{Scheme: "https",
			Host: "accounts.google.com",
			Path: "/o/oauth2/auth",
			// to get a refresh token. see https://developers.google.com/identity/protocols/OAuth2WebServer#offline
			RawQuery: "access_type=offline",
		}
	}
	if p.RedeemURL.String() == "" {
		p.RedeemURL = &url.URL{Scheme: "https",
			Host: "www.googleapis.com",
			Path: "/oauth2/v3/token"}
	}
	if p.ValidateURL.String() == "" {
		p.ValidateURL = &url.URL{Scheme: "https",
			Host: "www.googleapis.com",
			Path: "/oauth2/v1/tokeninfo"}
	}
	if p.Scope == "" {
		p.Scope = "profile email"
	}

	return &GoogleProvider{
		ProviderData: p,
		// Set a default GroupValidator to just always return valid (true), it will
		// be overwritten if we configured a Google group restriction.
		GroupValidator: func(email string) bool {
			return true
		},
	}
}

func emailFromIdToken(idToken string) (string, error) {

	// id_token is a base64 encode ID token payload
	// https://developers.google.com/accounts/docs/OAuth2Login#obtainuserinfo
	jwt := strings.Split(idToken, ".")
	b, err := base64.RawURLEncoding.DecodeString(jwt[1])
	if err != nil {
		return "", err
	}

	var email struct {
		Email         string `json:"email"`
		EmailVerified bool   `json:"email_verified"`
	}
	err = json.Unmarshal(b, &email)
	if err != nil {
		return "", err
	}
	if email.Email == "" {
		return "", errors.New("missing email")
	}
	if !email.EmailVerified {
		return "", fmt.Errorf("email %s not listed as verified", email.Email)
	}
	return email.Email, nil
}

func (p *GoogleProvider) Redeem(redirectURL, code string) (s *SessionState, err error) {
	if code == "" {
		err = errors.New("missing code")
		return
	}

	params := url.Values{}
	params.Add("redirect_uri", redirectURL)
	params.Add("client_id", p.ClientID)
	params.Add("client_secret", p.ClientSecret)
	params.Add("code", code)
	params.Add("grant_type", "authorization_code")
	var req *http.Request
	req, err = http.NewRequest("POST", p.RedeemURL.String(), bytes.NewBufferString(params.Encode()))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	var body []byte
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("got %d from %q %s", resp.StatusCode, p.RedeemURL.String(), body)
		return
	}

	var jsonResponse struct {
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
		ExpiresIn    int64  `json:"expires_in"`
		IdToken      string `json:"id_token"`
	}
	err = json.Unmarshal(body, &jsonResponse)
	if err != nil {
		return
	}
	var email string
	email, err = emailFromIdToken(jsonResponse.IdToken)
	if err != nil {
		return
	}
	s = &SessionState{
		AccessToken:  jsonResponse.AccessToken,
		ExpiresOn:    time.Now().Add(time.Duration(jsonResponse.ExpiresIn) * time.Second).Truncate(time.Second),
		RefreshToken: jsonResponse.RefreshToken,
		Email:        email,
	}
	return
}

// SetGroupRestriction configures the GoogleProvider to restrict access to the
// specified group(s). AdminEmail has to be an administrative email on the domain that is
// checked. CredentialsFile is the path to a json file containing a Google service
// account credentials.
func (p *GoogleProvider) SetGroupRestriction(groups []string, adminEmail string, credentialsReader io.Reader) {
	adminService := getAdminService(adminEmail, credentialsReader)
	p.GroupValidator = func(email string) bool {
		return userInGroup(adminService, groups, email)
	}
}

func getAdminService(adminEmail string, credentialsReader io.Reader) *admin.Service {
	data, err := ioutil.ReadAll(credentialsReader)
	if err != nil {
		log.Fatal("can't read Google credentials file:", err)
	}
	conf, err := google.JWTConfigFromJSON(data, admin.AdminDirectoryUserReadonlyScope, admin.AdminDirectoryGroupReadonlyScope)
	if err != nil {
		log.Fatal("can't load Google credentials file:", err)
	}
	conf.Subject = adminEmail

	client := conf.Client(oauth2.NoContext)
	adminService, err := admin.New(client)
	if err != nil {
		log.Fatal(err)
	}
	return adminService
}

func userInGroup(service *admin.Service, groups []string, email string) bool {
	user, err := fetchUser(service, email)
	if err != nil {
		log.Printf("error fetching user: %v", err)
		return false
	}
	id := user.Id
	custID := user.CustomerId

	for _, group := range groups {
		members, err := fetchGroupMembers(service, group)
		if err != nil {
			if err, ok := err.(*googleapi.Error); ok && err.Code == 404 {
				log.Printf("error fetching members for group %s: group does not exist", group)
			} else {
				log.Printf("error fetching group members: %v", err)
				return false
			}
		}

		for _, member := range members {
			switch member.Type {
			case "CUSTOMER":
				if member.Id == custID {
					return true
				}
			case "USER":
				if member.Id == id {
					return true
				}
			}
		}
	}
	return false
}

func fetchUser(service *admin.Service, email string) (*admin.User, error) {
	user, err := service.Users.Get(email).Do()
	return user, err
}

func fetchGroupMembers(service *admin.Service, group string) ([]*admin.Member, error) {
	members := []*admin.Member{}
	pageToken := ""
	for {
		req := service.Members.List(group)
		if pageToken != "" {
			req.PageToken(pageToken)
		}
		r, err := req.Do()
		if err != nil {
			return nil, err
		}
		for _, member := range r.Members {
			members = append(members, member)
		}
		if r.NextPageToken == "" {
			break
		}
		pageToken = r.NextPageToken
	}
	return members, nil
}

// ValidateGroup validates that the provided email exists in the configured Google
// group(s).
func (p *GoogleProvider) ValidateGroup(email string) bool {
	return p.GroupValidator(email)
}

func (p *GoogleProvider) RefreshSessionIfNeeded(s *SessionState) (bool, error) {
	if s == nil || s.ExpiresOn.After(time.Now()) || s.RefreshToken == "" {
		return false, nil
	}

	newToken, duration, err := p.redeemRefreshToken(s.RefreshToken)
	if err != nil {
		return false, err
	}

	// re-check that the user is in the proper google group(s)
	if !p.ValidateGroup(s.Email) {
		return false, fmt.Errorf("%s is no longer in the group(s)", s.Email)
	}

	origExpiration := s.ExpiresOn
	s.AccessToken = newToken
	s.ExpiresOn = time.Now().Add(duration).Truncate(time.Second)
	log.Printf("refreshed access token %s (expired on %s)", s, origExpiration)
	return true, nil
}

func (p *GoogleProvider) redeemRefreshToken(refreshToken string) (token string, expires time.Duration, err error) {
	// https://developers.google.com/identity/protocols/OAuth2WebServer#refresh
	params := url.Values{}
	params.Add("client_id", p.ClientID)
	params.Add("client_secret", p.ClientSecret)
	params.Add("refresh_token", refreshToken)
	params.Add("grant_type", "refresh_token")
	var req *http.Request
	req, err = http.NewRequest("POST", p.RedeemURL.String(), bytes.NewBufferString(params.Encode()))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	var body []byte
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("got %d from %q %s", resp.StatusCode, p.RedeemURL.String(), body)
		return
	}

	var data struct {
		AccessToken string `json:"access_token"`
		ExpiresIn   int64  `json:"expires_in"`
	}
	err = json.Unmarshal(body, &data)
	if err != nil {
		return
	}
	token = data.AccessToken
	expires = time.Duration(data.ExpiresIn) * time.Second
	return
}
<commit_msg>Strip JWT base64 padding before parsing. #560<commit_after>package providers

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/admin/directory/v1"
	"google.golang.org/api/googleapi"
)

type GoogleProvider struct {
	*ProviderData
	RedeemRefreshURL *url.URL
	// GroupValidator is a function that determines if the passed email is in
	// the configured Google group.
	GroupValidator func(string) bool
}

func NewGoogleProvider(p *ProviderData) *GoogleProvider {
	p.ProviderName = "Google"
	if p.LoginURL.String() == "" {
		p.LoginURL = &url.URL{Scheme: "https",
			Host: "accounts.google.com",
			Path: "/o/oauth2/auth",
			// to get a refresh token. see https://developers.google.com/identity/protocols/OAuth2WebServer#offline
			RawQuery: "access_type=offline",
		}
	}
	if p.RedeemURL.String() == "" {
		p.RedeemURL = &url.URL{Scheme: "https",
			Host: "www.googleapis.com",
			Path: "/oauth2/v3/token"}
	}
	if p.ValidateURL.String() == "" {
		p.ValidateURL = &url.URL{Scheme: "https",
			Host: "www.googleapis.com",
			Path: "/oauth2/v1/tokeninfo"}
	}
	if p.Scope == "" {
		p.Scope = "profile email"
	}

	return &GoogleProvider{
		ProviderData: p,
		// Set a default GroupValidator to just always return valid (true), it will
		// be overwritten if we configured a Google group restriction.
		GroupValidator: func(email string) bool {
			return true
		},
	}
}

func emailFromIdToken(idToken string) (string, error) {

	// id_token is a base64 encode ID token payload
	// https://developers.google.com/accounts/docs/OAuth2Login#obtainuserinfo
	jwt := strings.Split(idToken, ".")
	jwtData := strings.TrimSuffix(jwt[1], "=")
	b, err := base64.RawURLEncoding.DecodeString(jwtData)
	if err != nil {
		return "", err
	}

	var email struct {
		Email         string `json:"email"`
		EmailVerified bool   `json:"email_verified"`
	}
	err = json.Unmarshal(b, &email)
	if err != nil {
		return "", err
	}
	if email.Email == "" {
		return "", errors.New("missing email")
	}
	if !email.EmailVerified {
		return "", fmt.Errorf("email %s not listed as verified", email.Email)
	}
	return email.Email, nil
}

func (p *GoogleProvider) Redeem(redirectURL, code string) (s *SessionState, err error) {
	if code == "" {
		err = errors.New("missing code")
		return
	}

	params := url.Values{}
	params.Add("redirect_uri", redirectURL)
	params.Add("client_id", p.ClientID)
	params.Add("client_secret", p.ClientSecret)
	params.Add("code", code)
	params.Add("grant_type", "authorization_code")
	var req *http.Request
	req, err = http.NewRequest("POST", p.RedeemURL.String(), bytes.NewBufferString(params.Encode()))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	var body []byte
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("got %d from %q %s", resp.StatusCode, p.RedeemURL.String(), body)
		return
	}

	var jsonResponse struct {
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
		ExpiresIn    int64  `json:"expires_in"`
		IdToken      string `json:"id_token"`
	}
	err = json.Unmarshal(body, &jsonResponse)
	if err != nil {
		return
	}
	var email string
	email, err = emailFromIdToken(jsonResponse.IdToken)
	if err != nil {
		return
	}
	s = &SessionState{
		AccessToken:  jsonResponse.AccessToken,
		ExpiresOn:    time.Now().Add(time.Duration(jsonResponse.ExpiresIn) * time.Second).Truncate(time.Second),
		RefreshToken: jsonResponse.RefreshToken,
		Email:        email,
	}
	return
}

// SetGroupRestriction configures the GoogleProvider to restrict access to the
// specified group(s). AdminEmail has to be an administrative email on the domain that is
// checked. CredentialsFile is the path to a json file containing a Google service
// account credentials.
func (p *GoogleProvider) SetGroupRestriction(groups []string, adminEmail string, credentialsReader io.Reader) {
	adminService := getAdminService(adminEmail, credentialsReader)
	p.GroupValidator = func(email string) bool {
		return userInGroup(adminService, groups, email)
	}
}

func getAdminService(adminEmail string, credentialsReader io.Reader) *admin.Service {
	data, err := ioutil.ReadAll(credentialsReader)
	if err != nil {
		log.Fatal("can't read Google credentials file:", err)
	}
	conf, err := google.JWTConfigFromJSON(data, admin.AdminDirectoryUserReadonlyScope, admin.AdminDirectoryGroupReadonlyScope)
	if err != nil {
		log.Fatal("can't load Google credentials file:", err)
	}
	conf.Subject = adminEmail

	client := conf.Client(oauth2.NoContext)
	adminService, err := admin.New(client)
	if err != nil {
		log.Fatal(err)
	}
	return adminService
}

func userInGroup(service *admin.Service, groups []string, email string) bool {
	user, err := fetchUser(service, email)
	if err != nil {
		log.Printf("error fetching user: %v", err)
		return false
	}
	id := user.Id
	custID := user.CustomerId

	for _, group := range groups {
		members, err := fetchGroupMembers(service, group)
		if err != nil {
			if err, ok := err.(*googleapi.Error); ok && err.Code == 404 {
				log.Printf("error fetching members for group %s: group does not exist", group)
			} else {
				log.Printf("error fetching group members: %v", err)
				return false
			}
		}

		for _, member := range members {
			switch member.Type {
			case "CUSTOMER":
				if member.Id == custID {
					return true
				}
			case "USER":
				if member.Id == id {
					return true
				}
			}
		}
	}
	return false
}

func fetchUser(service *admin.Service, email string) (*admin.User, error) {
	user, err := service.Users.Get(email).Do()
	return user, err
}

func fetchGroupMembers(service *admin.Service, group string) ([]*admin.Member, error) {
	members := []*admin.Member{}
	pageToken := ""
	for {
		req := service.Members.List(group)
		if pageToken != "" {
			req.PageToken(pageToken)
		}
		r, err := req.Do()
		if err != nil {
			return nil, err
		}
		for _, member := range r.Members {
			members = append(members, member)
		}
		if r.NextPageToken == "" {
			break
		}
		pageToken = r.NextPageToken
	}
	return members, nil
}

// ValidateGroup validates that the provided email exists in the configured Google
// group(s).
func (p *GoogleProvider) ValidateGroup(email string) bool {
	return p.GroupValidator(email)
}

func (p *GoogleProvider) RefreshSessionIfNeeded(s *SessionState) (bool, error) {
	if s == nil || s.ExpiresOn.After(time.Now()) || s.RefreshToken == "" {
		return false, nil
	}

	newToken, duration, err := p.redeemRefreshToken(s.RefreshToken)
	if err != nil {
		return false, err
	}

	// re-check that the user is in the proper google group(s)
	if !p.ValidateGroup(s.Email) {
		return false, fmt.Errorf("%s is no longer in the group(s)", s.Email)
	}

	origExpiration := s.ExpiresOn
	s.AccessToken = newToken
	s.ExpiresOn = time.Now().Add(duration).Truncate(time.Second)
	log.Printf("refreshed access token %s (expired on %s)", s, origExpiration)
	return true, nil
}

func (p *GoogleProvider) redeemRefreshToken(refreshToken string) (token string, expires time.Duration, err error) {
	// https://developers.google.com/identity/protocols/OAuth2WebServer#refresh
	params := url.Values{}
	params.Add("client_id", p.ClientID)
	params.Add("client_secret", p.ClientSecret)
	params.Add("refresh_token", refreshToken)
	params.Add("grant_type", "refresh_token")
	var req *http.Request
	req, err = http.NewRequest("POST", p.RedeemURL.String(), bytes.NewBufferString(params.Encode()))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	var body []byte
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		err = fmt.Errorf("got %d from %q %s", resp.StatusCode, p.RedeemURL.String(), body)
		return
	}

	var data struct {
		AccessToken string `json:"access_token"`
		ExpiresIn   int64  `json:"expires_in"`
	}
	err = json.Unmarshal(body, &data)
	if err != nil {
		return
	}
	token = data.AccessToken
	expires = time.Duration(data.ExpiresIn) * time.Second
	return
}
<|endoftext|>


<commit_before>package builder

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	stdlog "log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/Symantec/Dominator/lib/filesystem/util"
	"github.com/Symantec/Dominator/lib/filter"
	"github.com/Symantec/Dominator/lib/format"
	"github.com/Symantec/Dominator/lib/fsutil"
	"github.com/Symantec/Dominator/lib/image"
	objectclient "github.com/Symantec/Dominator/lib/objectserver/client"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/lib/triggers"
	proto "github.com/Symantec/Dominator/proto/imaginator"
)

func (stream *imageStreamType) build(b *Builder, client *srpc.Client,
	request proto.BuildImageRequest, buildLog buildLogger) (
	*image.Image, error) {
	manifestDirectory, err := stream.getManifest(b, request.StreamName,
		request.GitBranch, request.Variables, buildLog)
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(manifestDirectory)
	img, err := buildImageFromManifest(client, manifestDirectory, request,
		b.bindMounts, stream, buildLog)
	if err != nil {
		return nil, err
	}
	return img, nil
}

func (stream *imageStreamType) getenv() map[string]string {
	envTable := make(map[string]string, 1)
	envTable["IMAGE_STREAM"] = stream.name
	envTable["IMAGE_STREAM_LEAF_NAME"] = filepath.Base(stream.name)
	return envTable
}

func (stream *imageStreamType) getManifest(b *Builder, streamName string,
	gitBranch string, variables map[string]string,
	buildLog io.Writer) (string, error) {
	if gitBranch == "" {
		gitBranch = "master"
	}
	variableFunc := b.getVariableFunc(stream.getenv(), variables)
	manifestRoot, err := makeTempDirectory("",
		strings.Replace(streamName, "/", "_", -1)+".manifest")
	if err != nil {
		return "", err
	}
	doCleanup := true
	defer func() {
		if doCleanup {
			os.RemoveAll(manifestRoot)
		}
	}()
	manifestDirectory := os.Expand(stream.ManifestDirectory, variableFunc)
	manifestUrl := os.Expand(stream.ManifestUrl, variableFunc)
	if parsedUrl, err := url.Parse(manifestUrl); err == nil {
		if parsedUrl.Scheme == "dir" {
			if parsedUrl.Path[0] != '/' {
				return "", fmt.Errorf("missing leading slash: %s",
					parsedUrl.Path)
			}
			if gitBranch != "master" {
				return "", fmt.Errorf("branch: %s is not master", gitBranch)
			}
			sourceTree := filepath.Join(parsedUrl.Path, manifestDirectory)
			fmt.Fprintf(buildLog, "Copying manifest tree: %s\n", sourceTree)
			if err := fsutil.CopyTree(manifestRoot, sourceTree); err != nil {
				return "", fmt.Errorf("error copying manifest: %s", err)
			}
			doCleanup = false
			return manifestRoot, nil
		}
	}
	fmt.Fprintf(buildLog, "Cloning repository: %s branch: %s\n",
		stream.ManifestUrl, gitBranch)
	err = runCommand(buildLog, "", "git", "init", manifestRoot)
	if err != nil {
		return "", err
	}
	err = runCommand(buildLog, manifestRoot, "git", "remote", "add", "origin",
		manifestUrl)
	if err != nil {
		return "", err
	}
	err = runCommand(buildLog, manifestRoot, "git", "config",
		"core.sparsecheckout", "true")
	if err != nil {
		return "", err
	}
	directorySelector := "*\n"
	if manifestDirectory != "" {
		directorySelector = manifestDirectory + "/*\n"
	}
	err = ioutil.WriteFile(
		path.Join(manifestRoot, ".git", "info", "sparse-checkout"),
		[]byte(directorySelector), 0644)
	if err != nil {
		return "", err
	}
	startTime := time.Now()
	err = runCommand(buildLog, manifestRoot, "git", "pull", "--depth=1",
		"origin", gitBranch)
	if err != nil {
		return "", err
	}
	if gitBranch != "master" {
		err = runCommand(buildLog, manifestRoot, "git", "checkout", gitBranch)
		if err != nil {
			return "", err
		}
	}
	loadTime := time.Since(startTime)
	repoSize, err := getTreeSize(manifestRoot)
	if err != nil {
		return "", err
	}
	speed := float64(repoSize) / loadTime.Seconds()
	fmt.Fprintf(buildLog,
		"Downloaded partial repository in %s, size: %s (%s/s)\n",
		format.Duration(loadTime), format.FormatBytes(repoSize),
		format.FormatBytes(uint64(speed)))
	gitDirectory := path.Join(manifestRoot, ".git")
	if err := os.RemoveAll(gitDirectory); err != nil {
		return "", err
	}
	if manifestDirectory != "" {
		// Move manifestDirectory into manifestRoot, remove anything else.
		err := os.Rename(path.Join(manifestRoot, manifestDirectory),
			gitDirectory)
		if err != nil {
			return "", err
		}
		filenames, err := listDirectory(manifestRoot)
		if err != nil {
			return "", err
		}
		for _, filename := range filenames {
			if filename == ".git" {
				continue
			}
			err := os.RemoveAll(path.Join(manifestRoot, filename))
			if err != nil {
				return "", err
			}
		}
		filenames, err = listDirectory(gitDirectory)
		if err != nil {
			return "", err
		}
		for _, filename := range filenames {
			err := os.Rename(path.Join(gitDirectory, filename),
				path.Join(manifestRoot, filename))
			if err != nil {
				return "", err
			}
		}
		if err := os.Remove(gitDirectory); err != nil {
			return "", err
		}
	}
	doCleanup = false
	return manifestRoot, nil
}

func getTreeSize(dirname string) (uint64, error) {
	var size uint64
	err := filepath.Walk(dirname,
		func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			size += uint64(info.Size())
			return nil
		})
	if err != nil {
		return 0, err
	}
	return size, nil
}

func listDirectory(directoryName string) ([]string, error) {
	directory, err := os.Open(directoryName)
	if err != nil {
		return nil, err
	}
	defer directory.Close()
	filenames, err := directory.Readdirnames(-1)
	if err != nil {
		return nil, err
	}
	return filenames, nil
}

func runCommand(buildLog io.Writer, cwd string, args ...string) error {
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = cwd
	cmd.Stdout = buildLog
	cmd.Stderr = buildLog
	return cmd.Run()
}

func buildImageFromManifest(client *srpc.Client, manifestDir string,
	request proto.BuildImageRequest, bindMounts []string,
	envGetter environmentGetter, buildLog buildLogger) (*image.Image, error) {
	// First load all the various manifest files (fail early on error).
	computedFilesList, err := util.LoadComputedFiles(
		path.Join(manifestDir, "computed-files.json"))
	if os.IsNotExist(err) {
		computedFilesList, err = util.LoadComputedFiles(
			path.Join(manifestDir, "computed-files"))
	}
	if err != nil && !os.IsNotExist(err) {
		return nil, fmt.Errorf("error loading computed files: %s", err)
	}
	imageFilter, addFilter, err := loadFilter(manifestDir)
	if err != nil {
		return nil, err
	}
	imageTriggers, addTriggers, err := loadTriggers(manifestDir)
	if err != nil {
		return nil, err
	}
	rootDir, err := makeTempDirectory("",
		strings.Replace(request.StreamName, "/", "_", -1)+".root")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(rootDir)
	fmt.Fprintf(buildLog, "Created image working directory: %s\n", rootDir)
	manifest, err := unpackImageAndProcessManifest(client, manifestDir,
		rootDir, bindMounts, false, envGetter, buildLog)
	if err != nil {
		return nil, err
	}
	if fi, err := os.Lstat(filepath.Join(manifestDir, "tests")); err == nil {
		if fi.IsDir() {
			testsDir := filepath.Join(rootDir, "tests", request.StreamName)
			if err := os.MkdirAll(testsDir, fsutil.DirPerms); err != nil {
				return nil, err
			}
			err := copyFiles(manifestDir, "tests", testsDir, buildLog)
			if err != nil {
				return nil, err
			}
		}
	}
	if addFilter {
		mergeableFilter := &filter.MergeableFilter{}
		mergeableFilter.Merge(manifest.sourceImageInfo.filter)
		mergeableFilter.Merge(imageFilter)
		imageFilter = mergeableFilter.ExportFilter()
	}
	if addTriggers {
		mergeableTriggers := &triggers.MergeableTriggers{}
		mergeableTriggers.Merge(manifest.sourceImageInfo.triggers)
		mergeableTriggers.Merge(imageTriggers)
		imageTriggers = mergeableTriggers.ExportTriggers()
	}
	return packImage(client, request, rootDir, manifest.filter,
		computedFilesList, imageFilter, imageTriggers, buildLog)
}

func buildImageFromManifestAndUpload(client *srpc.Client, manifestDir string,
	request proto.BuildImageRequest, bindMounts []string,
	envGetter environmentGetter,
	buildLog buildLogger) (*image.Image, string, error) {
	img, err := buildImageFromManifest(client, manifestDir, request, bindMounts,
		envGetter, buildLog)
	if err != nil {
		return nil, "", err
	}
	name, err := addImage(client, request, img)
	if err != nil {
		return nil, "", err
	}
	return img, name, nil
}

func buildTreeFromManifest(client *srpc.Client, manifestDir string,
	bindMounts []string, envGetter environmentGetter,
	buildLog io.Writer) (string, error) {
	rootDir, err := makeTempDirectory("", "tree")
	if err != nil {
		return "", err
	}
	_, err = unpackImageAndProcessManifest(client, manifestDir, rootDir,
		bindMounts, true, envGetter, buildLog)
	if err != nil {
		os.RemoveAll(rootDir)
		return "", err
	}
	return rootDir, nil
}

func loadFilter(manifestDir string) (*filter.Filter, bool, error) {
	imageFilter, err := filter.Load(path.Join(manifestDir, "filter"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	addFilter, err := filter.Load(path.Join(manifestDir, "filter.add"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	if imageFilter == nil && addFilter == nil {
		return nil, false, nil
	} else if imageFilter != nil && addFilter != nil {
		return nil, false, errors.New(
			"filter and filter.add files both present")
	} else if imageFilter != nil {
		return imageFilter, false, nil
	} else {
		return addFilter, true, nil
	}
}

func loadTriggers(manifestDir string) (*triggers.Triggers, bool, error) {
	imageTriggers, err := triggers.Load(path.Join(manifestDir, "triggers"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	addTriggers, err := triggers.Load(path.Join(manifestDir, "triggers.add"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	if imageTriggers == nil && addTriggers == nil {
		return nil, false, nil
	} else if imageTriggers != nil && addTriggers != nil {
		return nil, false, errors.New(
			"triggers and triggers.add files both present")
	} else if imageTriggers != nil {
		return imageTriggers, false, nil
	} else {
		return addTriggers, true, nil
	}
}

func unpackImage(client *srpc.Client, streamName string,
	maxSourceAge, expiresIn time.Duration, rootDir string,
	buildLog io.Writer) (*sourceImageInfoType, error) {
	imageName, sourceImage, err := getLatestImage(client, streamName, buildLog)
	if err != nil {
		return nil, err
	}
	if sourceImage == nil {
		return nil, errors.New(errNoSourceImage + streamName)
	}
	if maxSourceAge > 0 && time.Since(sourceImage.CreatedOn) > maxSourceAge {
		return nil, errors.New(errNoSourceImage + streamName)
	}
	objClient := objectclient.AttachObjectClient(client)
	defer objClient.Close()
	err = util.Unpack(sourceImage.FileSystem, objClient, rootDir,
		stdlog.New(buildLog, "", 0))
	if err != nil {
		return nil, err
	}
	fmt.Fprintf(buildLog, "Source image: %s\n", imageName)
	return &sourceImageInfoType{sourceImage.Filter, sourceImage.Triggers}, nil
}
<commit_msg>imaginator: add IMAGE_STREAM_DIRECTORY_NAME to build variables.<commit_after>package builder

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	stdlog "log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/Symantec/Dominator/lib/filesystem/util"
	"github.com/Symantec/Dominator/lib/filter"
	"github.com/Symantec/Dominator/lib/format"
	"github.com/Symantec/Dominator/lib/fsutil"
	"github.com/Symantec/Dominator/lib/image"
	objectclient "github.com/Symantec/Dominator/lib/objectserver/client"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/lib/triggers"
	proto "github.com/Symantec/Dominator/proto/imaginator"
)

func (stream *imageStreamType) build(b *Builder, client *srpc.Client,
	request proto.BuildImageRequest, buildLog buildLogger) (
	*image.Image, error) {
	manifestDirectory, err := stream.getManifest(b, request.StreamName,
		request.GitBranch, request.Variables, buildLog)
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(manifestDirectory)
	img, err := buildImageFromManifest(client, manifestDirectory, request,
		b.bindMounts, stream, buildLog)
	if err != nil {
		return nil, err
	}
	return img, nil
}

func (stream *imageStreamType) getenv() map[string]string {
	envTable := make(map[string]string, 1)
	envTable["IMAGE_STREAM"] = stream.name
	envTable["IMAGE_STREAM_DIRECTORY_NAME"] = filepath.Dir(stream.name)
	envTable["IMAGE_STREAM_LEAF_NAME"] = filepath.Base(stream.name)
	return envTable
}

func (stream *imageStreamType) getManifest(b *Builder, streamName string,
	gitBranch string, variables map[string]string,
	buildLog io.Writer) (string, error) {
	if gitBranch == "" {
		gitBranch = "master"
	}
	variableFunc := b.getVariableFunc(stream.getenv(), variables)
	manifestRoot, err := makeTempDirectory("",
		strings.Replace(streamName, "/", "_", -1)+".manifest")
	if err != nil {
		return "", err
	}
	doCleanup := true
	defer func() {
		if doCleanup {
			os.RemoveAll(manifestRoot)
		}
	}()
	manifestDirectory := os.Expand(stream.ManifestDirectory, variableFunc)
	manifestUrl := os.Expand(stream.ManifestUrl, variableFunc)
	if parsedUrl, err := url.Parse(manifestUrl); err == nil {
		if parsedUrl.Scheme == "dir" {
			if parsedUrl.Path[0] != '/' {
				return "", fmt.Errorf("missing leading slash: %s",
					parsedUrl.Path)
			}
			if gitBranch != "master" {
				return "", fmt.Errorf("branch: %s is not master", gitBranch)
			}
			sourceTree := filepath.Join(parsedUrl.Path, manifestDirectory)
			fmt.Fprintf(buildLog, "Copying manifest tree: %s\n", sourceTree)
			if err := fsutil.CopyTree(manifestRoot, sourceTree); err != nil {
				return "", fmt.Errorf("error copying manifest: %s", err)
			}
			doCleanup = false
			return manifestRoot, nil
		}
	}
	fmt.Fprintf(buildLog, "Cloning repository: %s branch: %s\n",
		stream.ManifestUrl, gitBranch)
	err = runCommand(buildLog, "", "git", "init", manifestRoot)
	if err != nil {
		return "", err
	}
	err = runCommand(buildLog, manifestRoot, "git", "remote", "add", "origin",
		manifestUrl)
	if err != nil {
		return "", err
	}
	err = runCommand(buildLog, manifestRoot, "git", "config",
		"core.sparsecheckout", "true")
	if err != nil {
		return "", err
	}
	directorySelector := "*\n"
	if manifestDirectory != "" {
		directorySelector = manifestDirectory + "/*\n"
	}
	err = ioutil.WriteFile(
		path.Join(manifestRoot, ".git", "info", "sparse-checkout"),
		[]byte(directorySelector), 0644)
	if err != nil {
		return "", err
	}
	startTime := time.Now()
	err = runCommand(buildLog, manifestRoot, "git", "pull", "--depth=1",
		"origin", gitBranch)
	if err != nil {
		return "", err
	}
	if gitBranch != "master" {
		err = runCommand(buildLog, manifestRoot, "git", "checkout", gitBranch)
		if err != nil {
			return "", err
		}
	}
	loadTime := time.Since(startTime)
	repoSize, err := getTreeSize(manifestRoot)
	if err != nil {
		return "", err
	}
	speed := float64(repoSize) / loadTime.Seconds()
	fmt.Fprintf(buildLog,
		"Downloaded partial repository in %s, size: %s (%s/s)\n",
		format.Duration(loadTime), format.FormatBytes(repoSize),
		format.FormatBytes(uint64(speed)))
	gitDirectory := path.Join(manifestRoot, ".git")
	if err := os.RemoveAll(gitDirectory); err != nil {
		return "", err
	}
	if manifestDirectory != "" {
		// Move manifestDirectory into manifestRoot, remove anything else.
		err := os.Rename(path.Join(manifestRoot, manifestDirectory),
			gitDirectory)
		if err != nil {
			return "", err
		}
		filenames, err := listDirectory(manifestRoot)
		if err != nil {
			return "", err
		}
		for _, filename := range filenames {
			if filename == ".git" {
				continue
			}
			err := os.RemoveAll(path.Join(manifestRoot, filename))
			if err != nil {
				return "", err
			}
		}
		filenames, err = listDirectory(gitDirectory)
		if err != nil {
			return "", err
		}
		for _, filename := range filenames {
			err := os.Rename(path.Join(gitDirectory, filename),
				path.Join(manifestRoot, filename))
			if err != nil {
				return "", err
			}
		}
		if err := os.Remove(gitDirectory); err != nil {
			return "", err
		}
	}
	doCleanup = false
	return manifestRoot, nil
}

func getTreeSize(dirname string) (uint64, error) {
	var size uint64
	err := filepath.Walk(dirname,
		func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			size += uint64(info.Size())
			return nil
		})
	if err != nil {
		return 0, err
	}
	return size, nil
}

func listDirectory(directoryName string) ([]string, error) {
	directory, err := os.Open(directoryName)
	if err != nil {
		return nil, err
	}
	defer directory.Close()
	filenames, err := directory.Readdirnames(-1)
	if err != nil {
		return nil, err
	}
	return filenames, nil
}

func runCommand(buildLog io.Writer, cwd string, args ...string) error {
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = cwd
	cmd.Stdout = buildLog
	cmd.Stderr = buildLog
	return cmd.Run()
}

func buildImageFromManifest(client *srpc.Client, manifestDir string,
	request proto.BuildImageRequest, bindMounts []string,
	envGetter environmentGetter, buildLog buildLogger) (*image.Image, error) {
	// First load all the various manifest files (fail early on error).
	computedFilesList, err := util.LoadComputedFiles(
		path.Join(manifestDir, "computed-files.json"))
	if os.IsNotExist(err) {
		computedFilesList, err = util.LoadComputedFiles(
			path.Join(manifestDir, "computed-files"))
	}
	if err != nil && !os.IsNotExist(err) {
		return nil, fmt.Errorf("error loading computed files: %s", err)
	}
	imageFilter, addFilter, err := loadFilter(manifestDir)
	if err != nil {
		return nil, err
	}
	imageTriggers, addTriggers, err := loadTriggers(manifestDir)
	if err != nil {
		return nil, err
	}
	rootDir, err := makeTempDirectory("",
		strings.Replace(request.StreamName, "/", "_", -1)+".root")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(rootDir)
	fmt.Fprintf(buildLog, "Created image working directory: %s\n", rootDir)
	manifest, err := unpackImageAndProcessManifest(client, manifestDir,
		rootDir, bindMounts, false, envGetter, buildLog)
	if err != nil {
		return nil, err
	}
	if fi, err := os.Lstat(filepath.Join(manifestDir, "tests")); err == nil {
		if fi.IsDir() {
			testsDir := filepath.Join(rootDir, "tests", request.StreamName)
			if err := os.MkdirAll(testsDir, fsutil.DirPerms); err != nil {
				return nil, err
			}
			err := copyFiles(manifestDir, "tests", testsDir, buildLog)
			if err != nil {
				return nil, err
			}
		}
	}
	if addFilter {
		mergeableFilter := &filter.MergeableFilter{}
		mergeableFilter.Merge(manifest.sourceImageInfo.filter)
		mergeableFilter.Merge(imageFilter)
		imageFilter = mergeableFilter.ExportFilter()
	}
	if addTriggers {
		mergeableTriggers := &triggers.MergeableTriggers{}
		mergeableTriggers.Merge(manifest.sourceImageInfo.triggers)
		mergeableTriggers.Merge(imageTriggers)
		imageTriggers = mergeableTriggers.ExportTriggers()
	}
	return packImage(client, request, rootDir, manifest.filter,
		computedFilesList, imageFilter, imageTriggers, buildLog)
}

func buildImageFromManifestAndUpload(client *srpc.Client, manifestDir string,
	request proto.BuildImageRequest, bindMounts []string,
	envGetter environmentGetter,
	buildLog buildLogger) (*image.Image, string, error) {
	img, err := buildImageFromManifest(client, manifestDir, request, bindMounts,
		envGetter, buildLog)
	if err != nil {
		return nil, "", err
	}
	name, err := addImage(client, request, img)
	if err != nil {
		return nil, "", err
	}
	return img, name, nil
}

func buildTreeFromManifest(client *srpc.Client, manifestDir string,
	bindMounts []string, envGetter environmentGetter,
	buildLog io.Writer) (string, error) {
	rootDir, err := makeTempDirectory("", "tree")
	if err != nil {
		return "", err
	}
	_, err = unpackImageAndProcessManifest(client, manifestDir, rootDir,
		bindMounts, true, envGetter, buildLog)
	if err != nil {
		os.RemoveAll(rootDir)
		return "", err
	}
	return rootDir, nil
}

func loadFilter(manifestDir string) (*filter.Filter, bool, error) {
	imageFilter, err := filter.Load(path.Join(manifestDir, "filter"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	addFilter, err := filter.Load(path.Join(manifestDir, "filter.add"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	if imageFilter == nil && addFilter == nil {
		return nil, false, nil
	} else if imageFilter != nil && addFilter != nil {
		return nil, false, errors.New(
			"filter and filter.add files both present")
	} else if imageFilter != nil {
		return imageFilter, false, nil
	} else {
		return addFilter, true, nil
	}
}

func loadTriggers(manifestDir string) (*triggers.Triggers, bool, error) {
	imageTriggers, err := triggers.Load(path.Join(manifestDir, "triggers"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	addTriggers, err := triggers.Load(path.Join(manifestDir, "triggers.add"))
	if err != nil && !os.IsNotExist(err) {
		return nil, false, err
	}
	if imageTriggers == nil && addTriggers == nil {
		return nil, false, nil
	} else if imageTriggers != nil && addTriggers != nil {
		return nil, false, errors.New(
			"triggers and triggers.add files both present")
	} else if imageTriggers != nil {
		return imageTriggers, false, nil
	} else {
		return addTriggers, true, nil
	}
}

func unpackImage(client *srpc.Client, streamName string,
	maxSourceAge, expiresIn time.Duration, rootDir string,
	buildLog io.Writer) (*sourceImageInfoType, error) {
	imageName, sourceImage, err := getLatestImage(client, streamName, buildLog)
	if err != nil {
		return nil, err
	}
	if sourceImage == nil {
		return nil, errors.New(errNoSourceImage + streamName)
	}
	if maxSourceAge > 0 && time.Since(sourceImage.CreatedOn) > maxSourceAge {
		return nil, errors.New(errNoSourceImage + streamName)
	}
	objClient := objectclient.AttachObjectClient(client)
	defer objClient.Close()
	err = util.Unpack(sourceImage.FileSystem, objClient, rootDir,
		stdlog.New(buildLog, "", 0))
	if err != nil {
		return nil, err
	}
	fmt.Fprintf(buildLog, "Source image: %s\n", imageName)
	return &sourceImageInfoType{sourceImage.Filter, sourceImage.Triggers}, nil
}
<|endoftext|>


<commit_before>package pop

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"

	_ "github.com/lib/pq"
	"github.com/markbates/going/defaults"
	. "github.com/markbates/pop/columns"
	"github.com/markbates/pop/fizz"
	"github.com/markbates/pop/fizz/translators"
	"github.com/pkg/errors"
)

var _ dialect = &postgresql{}

type postgresql struct {
	translateCache    map[string]string
	mu                sync.Mutex
	ConnectionDetails *ConnectionDetails
}

func (p *postgresql) Details() *ConnectionDetails {
	return p.ConnectionDetails
}

func (p *postgresql) Create(s store, model *Model, cols Columns) error {
	keyType := model.PrimaryKeyType()
	switch keyType {
	case "int", "int64":
		cols.Remove("id")
		id := struct {
			ID int `db:"id"`
		}{}
		w := cols.Writeable()
		query := fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s) returning id", model.TableName(), w.String(), w.SymbolizedString())
		Log(query)
		stmt, err := s.PrepareNamed(query)
		if err != nil {
			return errors.WithStack(err)
		}
		err = stmt.Get(&id, model.Value)
		if err != nil {
			return errors.WithStack(err)
		}
		model.setID(id.ID)
		return nil
	case "UUID":
		return genericCreate(s, model, cols)
	}
	return errors.Errorf("can not use %s as a primary key type!", keyType)
}

func (p *postgresql) Update(s store, model *Model, cols Columns) error {
	return genericUpdate(s, model, cols)
}

func (p *postgresql) Destroy(s store, model *Model) error {
	return genericDestroy(s, model)
}

func (p *postgresql) SelectOne(s store, model *Model, query Query) error {
	return genericSelectOne(s, model, query)
}

func (p *postgresql) SelectMany(s store, models *Model, query Query) error {
	return genericSelectMany(s, models, query)
}

func (p *postgresql) CreateDB() error {
	// createdb -h db -p 5432 -U postgres enterprise_development
	deets := p.ConnectionDetails
	cmd := exec.Command("createdb", "-e", "-h", deets.Host, "-p", deets.Port, "-U", deets.User, deets.Database)
	Log(strings.Join(cmd.Args, " "))
	comboOut, err := cmd.CombinedOutput()
	if err != nil {
		err = fmt.Errorf("%s: %s", err.Error(), string(comboOut))
		return errors.Wrapf(err, "error creating PostgreSQL database %s", deets.Database)
	}

	fmt.Printf("created database %s\n", deets.Database)
	return nil
}

func (p *postgresql) DropDB() error {
	deets := p.ConnectionDetails
	cmd := exec.Command("dropdb", "-e", "-h", deets.Host, "-p", deets.Port, "-U", deets.User, deets.Database)
	Log(strings.Join(cmd.Args, " "))
	comboOut, err := cmd.CombinedOutput()
	if err != nil {
		err = fmt.Errorf("%s: %s", err.Error(), string(comboOut))
		return errors.Wrapf(err, "error dropping PostgreSQL database %s", deets.Database)
	}
	fmt.Printf("dropped database %s\n", deets.Database)
	return nil
}

func (m *postgresql) URL() string {
	c := m.ConnectionDetails
	if c.URL != "" {
		return c.URL
	}
	ssl := defaults.String(c.Options["sslmode"], "disable")

	s := "postgres://%s:%s@%s:%s/%s?sslmode=%s"
	return fmt.Sprintf(s, c.User, c.Password, c.Host, c.Port, c.Database, ssl)
}

func (m *postgresql) MigrationURL() string {
	return m.URL()
}

func (p *postgresql) TranslateSQL(sql string) string {
	defer p.mu.Unlock()
	p.mu.Lock()

	if csql, ok := p.translateCache[sql]; ok {
		return csql
	}
	curr := 1
	out := make([]byte, 0, len(sql))
	for i := 0; i < len(sql); i++ {
		if sql[i] == '?' {
			str := "$" + strconv.Itoa(curr)
			for _, char := range str {
				out = append(out, byte(char))
			}
			curr += 1
		} else {
			out = append(out, sql[i])
		}
	}
	csql := string(out)
	p.translateCache[sql] = csql
	return csql
}

func (p *postgresql) FizzTranslator() fizz.Translator {
	return translators.NewPostgres()
}

func (p *postgresql) Lock(fn func() error) error {
	return fn()
}

func (p *postgresql) DumpSchema(w io.Writer) error {
	cmd := exec.Command("pg_dump", "-s", fmt.Sprintf("--dbname=%s", p.URL()))
	Log(strings.Join(cmd.Args, " "))
	cmd.Stdout = w
	cmd.Stderr = os.Stderr

	err := cmd.Run()
	if err != nil {
		return err
	}

	fmt.Printf("dumped schema for %s\n", p.Details().Database)
	return nil
}

func (p *postgresql) LoadSchema(r io.Reader) error {
	cmd := exec.Command("psql", p.URL())
	in, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	go func() {
		defer in.Close()
		io.Copy(in, r)
	}()
	Log(strings.Join(cmd.Args, " "))
	err = cmd.Start()
	if err != nil {
		return err
	}
	err = cmd.Wait()
	if err != nil {
		return err
	}

	fmt.Printf("loaded schema for %s\n", p.Details().Database)
	return nil
}

func (p *postgresql) TruncateAll(tx *Connection) error {
	return tx.RawQuery(pgTruncate).Exec()
}

func newPostgreSQL(deets *ConnectionDetails) dialect {
	cd := &postgresql{
		ConnectionDetails: deets,
		translateCache:    map[string]string{},
		mu:                sync.Mutex{},
	}
	return cd
}

const pgTruncate = `DO
$func$
BEGIN
   EXECUTE
  (SELECT 'TRUNCATE TABLE '
       || string_agg(quote_ident(schemaname) || '.' || quote_ident(tablename), ', ')
       || ' CASCADE'
   FROM   pg_tables
   WHERE  schemaname = 'public'
   );
END
$func$;`
<commit_msg>create/drop db via psql, trust auth not required<commit_after>package pop

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"

	_ "github.com/lib/pq"
	"github.com/markbates/going/defaults"
	. "github.com/markbates/pop/columns"
	"github.com/markbates/pop/fizz"
	"github.com/markbates/pop/fizz/translators"
	"github.com/pkg/errors"
)

var _ dialect = &postgresql{}

type postgresql struct {
	translateCache    map[string]string
	mu                sync.Mutex
	ConnectionDetails *ConnectionDetails
}

func (p *postgresql) Details() *ConnectionDetails {
	return p.ConnectionDetails
}

func (p *postgresql) Create(s store, model *Model, cols Columns) error {
	keyType := model.PrimaryKeyType()
	switch keyType {
	case "int", "int64":
		cols.Remove("id")
		id := struct {
			ID int `db:"id"`
		}{}
		w := cols.Writeable()
		query := fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s) returning id", model.TableName(), w.String(), w.SymbolizedString())
		Log(query)
		stmt, err := s.PrepareNamed(query)
		if err != nil {
			return errors.WithStack(err)
		}
		err = stmt.Get(&id, model.Value)
		if err != nil {
			return errors.WithStack(err)
		}
		model.setID(id.ID)
		return nil
	case "UUID":
		return genericCreate(s, model, cols)
	}
	return errors.Errorf("can not use %s as a primary key type!", keyType)
}

func (p *postgresql) Update(s store, model *Model, cols Columns) error {
	return genericUpdate(s, model, cols)
}

func (p *postgresql) Destroy(s store, model *Model) error {
	return genericDestroy(s, model)
}

func (p *postgresql) SelectOne(s store, model *Model, query Query) error {
	return genericSelectOne(s, model, query)
}

func (p *postgresql) SelectMany(s store, models *Model, query Query) error {
	return genericSelectMany(s, models, query)
}

func (p *postgresql) CreateDB() error {
	// createdb -h db -p 5432 -U postgres enterprise_development
	deets := p.ConnectionDetails
	cmd := exec.Command("psql", p.urlWithoutDb(), "-c", fmt.Sprintf("create database %s", deets.Database))
	Log(strings.Join(cmd.Args, " "))
	comboOut, err := cmd.CombinedOutput()
	if err != nil {
		err = fmt.Errorf("%s: %s", err.Error(), string(comboOut))
		return errors.Wrapf(err, "error creating PostgreSQL database %s", deets.Database)
	}

	fmt.Printf("created database %s\n", deets.Database)
	return nil
}

func (p *postgresql) DropDB() error {
	deets := p.ConnectionDetails
	cmd := exec.Command("psql", p.urlWithoutDb(), "-c", fmt.Sprintf("drop database %s", deets.Database))
	Log(strings.Join(cmd.Args, " "))
	comboOut, err := cmd.CombinedOutput()
	if err != nil {
		err = fmt.Errorf("%s: %s", err.Error(), string(comboOut))
		return errors.Wrapf(err, "error dropping PostgreSQL database %s", deets.Database)
	}
	fmt.Printf("dropped database %s\n", deets.Database)
	return nil
}

func (m *postgresql) URL() string {
	c := m.ConnectionDetails
	if c.URL != "" {
		return c.URL
	}
	ssl := defaults.String(c.Options["sslmode"], "disable")

	s := "postgres://%s:%s@%s:%s/%s?sslmode=%s"
	return fmt.Sprintf(s, c.User, c.Password, c.Host, c.Port, c.Database, ssl)
}

func (m *postgresql) urlWithoutDb() string {
	c := m.ConnectionDetails
	ssl := defaults.String(c.Options["sslmode"], "disable")

	s := "postgres://%s:%s@%s:%s/?sslmode=%s"
	return fmt.Sprintf(s, c.User, c.Password, c.Host, c.Port, ssl)
}

func (m *postgresql) MigrationURL() string {
	return m.URL()
}

func (p *postgresql) TranslateSQL(sql string) string {
	defer p.mu.Unlock()
	p.mu.Lock()

	if csql, ok := p.translateCache[sql]; ok {
		return csql
	}
	curr := 1
	out := make([]byte, 0, len(sql))
	for i := 0; i < len(sql); i++ {
		if sql[i] == '?' {
			str := "$" + strconv.Itoa(curr)
			for _, char := range str {
				out = append(out, byte(char))
			}
			curr += 1
		} else {
			out = append(out, sql[i])
		}
	}
	csql := string(out)
	p.translateCache[sql] = csql
	return csql
}

func (p *postgresql) FizzTranslator() fizz.Translator {
	return translators.NewPostgres()
}

func (p *postgresql) Lock(fn func() error) error {
	return fn()
}

func (p *postgresql) DumpSchema(w io.Writer) error {
	cmd := exec.Command("pg_dump", "-s", fmt.Sprintf("--dbname=%s", p.URL()))
	Log(strings.Join(cmd.Args, " "))
	cmd.Stdout = w
	cmd.Stderr = os.Stderr

	err := cmd.Run()
	if err != nil {
		return err
	}

	fmt.Printf("dumped schema for %s\n", p.Details().Database)
	return nil
}

func (p *postgresql) LoadSchema(r io.Reader) error {
	cmd := exec.Command("psql", p.URL())
	in, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	go func() {
		defer in.Close()
		io.Copy(in, r)
	}()
	Log(strings.Join(cmd.Args, " "))
	err = cmd.Start()
	if err != nil {
		return err
	}
	err = cmd.Wait()
	if err != nil {
		return err
	}

	fmt.Printf("loaded schema for %s\n", p.Details().Database)
	return nil
}

func (p *postgresql) TruncateAll(tx *Connection) error {
	return tx.RawQuery(pgTruncate).Exec()
}

func newPostgreSQL(deets *ConnectionDetails) dialect {
	cd := &postgresql{
		ConnectionDetails: deets,
		translateCache:    map[string]string{},
		mu:                sync.Mutex{},
	}
	return cd
}

const pgTruncate = `DO
$func$
BEGIN
   EXECUTE
  (SELECT 'TRUNCATE TABLE '
       || string_agg(quote_ident(schemaname) || '.' || quote_ident(tablename), ', ')
       || ' CASCADE'
   FROM   pg_tables
   WHERE  schemaname = 'public'
   );
END
$func$;`
<|endoftext|>


<commit_before>// Copyright 2016 Joel Scoble and The JoeFriday authors.
// 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 facts

import (
	//Flat "github.com/google/flatbuffers/go"

	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	joe "github.com/mohae/joefriday"
)

const ProcCPUInfo = "/proc/cpuinfo"

// Facter gathers CPUInfo facts.
type Profiler struct {
	joe.Proc
	Val []byte
}

func New() (prof *Profiler, err error) {
	f, err := os.Open(ProcCPUInfo)
	if err != nil {
		return nil, err
	}
	return &Profiler{Proc: joe.NewProc(f), Val: make([]byte, 0, 160)}, nil
}

func (prof *Profiler) Reset() {
	prof.Lock()
	prof.Val = prof.Val[:0]
	prof.Unlock()
	prof.Proc.Reset()
}

// GetFacts gets the processor information from /proc/cpuinfo
func (prof *Profiler) Get() (facts *Facts, err error) {
	var (
		cpuCnt, i, pos int
		v              byte
		name, value    string
		cpu            Fact
	)
	prof.Reset()
	facts = &Facts{Timestamp: time.Now().UTC().UnixNano()}
	for {
		prof.Line, err = prof.Buf.ReadSlice('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, fmt.Errorf("error reading output bytes: %s", err)
		}
		// First grab the attribute name; everything up to the ':'.  The key may have
		// spaces and has trailing spaces; that gets trimmed.
		for i, v = range prof.Line {
			if v == 0x3A {
				pos = i + 1
				break
			}
			prof.Val = append(prof.Val, v)
		}
		name = strings.TrimSpace(string(prof.Val[:]))
		prof.Val = prof.Val[:0]
		// if there's anything left, the value is everything else; trim spaces
		if pos < len(prof.Line) {
			value = strings.TrimSpace(string(prof.Line[pos:]))
		}
		// check to see if this is flat.Facts for a different processor
		if name == "processor" {
			if cpuCnt > 0 {
				facts.CPU = append(facts.CPU, cpu)
			}
			cpuCnt++
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "fact: processor", Err: err}
			}
			cpu = Fact{Processor: int16(i)}
			continue
		}
		if name == "vendor_id" {
			cpu.VendorID = value
			continue
		}
		if name == "cpu family" {
			cpu.CPUFamily = value
			continue
		}
		if name == "model" {
			cpu.Model = value
			continue
		}
		if name == "model name" {
			cpu.ModelName = value
			continue
		}
		if name == "stepping" {
			cpu.Stepping = value
			continue
		}
		if name == "microcode" {
			cpu.Microcode = value
			continue
		}
		if name == "cpu MHz" {
			f, err := strconv.ParseFloat(value, 32)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: cpu MHz", Err: err}
			}
			cpu.CPUMHz = float32(f)
			continue
		}
		if name == "cache size" {
			cpu.CacheSize = value
			continue
		}
		if name == "physical id" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: physical id", Err: err}
			}
			cpu.PhysicalID = int16(i)
			continue
		}
		if name == "siblings" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: siblings", Err: err}
			}
			cpu.Siblings = int16(i)
			continue
		}
		if name == "core id" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: core id", Err: err}
			}
			cpu.CoreID = int16(i)
			continue
		}
		if name == "cpu cores" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: cpu cores", Err: err}
			}
			cpu.CPUCores = int16(i)
			continue
		}
		if name == "apicid" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: apicid", Err: err}
			}
			cpu.ApicID = int16(i)
			continue
		}
		if name == "initial apicid" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: initial apicid", Err: err}
			}
			cpu.InitialApicID = int16(i)
			continue
		}
		if name == "fpu" {
			cpu.FPU = value
			continue
		}
		if name == "fpu_exception" {
			cpu.FPUException = value
			continue
		}
		if name == "cpuid level" {
			cpu.CPUIDLevel = value
			continue
		}
		if name == "WP" {
			cpu.WP = value
			continue
		}
		if name == "flags" {
			cpu.Flags = value
			continue
		}
		if name == "bogomips" {
			f, err := strconv.ParseFloat(value, 32)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: bogomips", Err: err}
			}
			cpu.BogoMIPS = float32(f)
			continue
		}
		if name == "clflush size" {
			cpu.CLFlushSize = value
			continue
		}
		if name == "cache_alignment" {
			cpu.CacheAlignment = value
			continue
		}
		if name == "address sizes" {
			cpu.AddressSizes = value
			continue
		}
		if name == "power management" {
			cpu.PowerManagement = value
		}
	}
	facts.CPU = append(facts.CPU, cpu)
	return facts, nil
}

// TODO: is it even worth it to have this as a global?  Should GetFacts
// just instantiate a local version and use that?
var stdProfiler *Profiler
var stdProfilerMu sync.Mutex

func Get() (facts *Facts, err error) {
	stdProfilerMu.Lock()
	defer stdProfilerMu.Unlock()
	if stdProfiler == nil {
		stdProfiler, err = New()
		if err != nil {
			return nil, err
		}
	}
	return stdProfiler.Get()
}

// Facts are a collection of facts, cpuinfo, about the system's cpus.
type Facts struct {
	Timestamp int64
	CPU       []Fact `json:"cpu"`
}

// Fact holds the /proc/cpuinfo for a single cpu
type Fact struct {
	Processor       int16   `json:"processor"`
	VendorID        string  `json:"vendor_id"`
	CPUFamily       string  `json:"cpu_family"`
	Model           string  `json:"model"`
	ModelName       string  `json:"model_name"`
	Stepping        string  `json:"stepping"`
	Microcode       string  `json:"microcode"`
	CPUMHz          float32 `json:"cpu_mhz"`
	CacheSize       string  `json:"cache_size"`
	PhysicalID      int16   `json:"physical_id"`
	Siblings        int16   `json:"siblings"`
	CoreID          int16   `json:"core_id"`
	CPUCores        int16   `json:"cpu_cores"`
	ApicID          int16   `json:"apicid"`
	InitialApicID   int16   `json:"initial_apicid"`
	FPU             string  `json:"fpu"`
	FPUException    string  `json:"fpu_exception"`
	CPUIDLevel      string  `json:"cpuid_level"`
	WP              string  `json:"wp"`
	Flags           string  `json:"flags"` // should this be a []string?
	BogoMIPS        float32 `json:"bogomips"`
	CLFlushSize     string  `json:"clflush_size"`
	CacheAlignment  string  `json:"cache_alignment"`
	AddressSizes    string  `json:"address_sizes"`
	PowerManagement string  `json:"power_management"`
}
<commit_msg>rename proc file name var<commit_after>// Copyright 2016 Joel Scoble and The JoeFriday authors.
// 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 facts

import (
	//Flat "github.com/google/flatbuffers/go"

	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	joe "github.com/mohae/joefriday"
)

const profFile = "/proc/cpuinfo"

// Facter gathers CPUInfo facts.
type Profiler struct {
	joe.Proc
	Val []byte
}

func New() (prof *Profiler, err error) {
	f, err := os.Open(profFile)
	if err != nil {
		return nil, err
	}
	return &Profiler{Proc: joe.NewProc(f), Val: make([]byte, 0, 160)}, nil
}

func (prof *Profiler) Reset() {
	prof.Lock()
	prof.Val = prof.Val[:0]
	prof.Unlock()
	prof.Proc.Reset()
}

// GetFacts gets the processor information from /proc/cpuinfo
func (prof *Profiler) Get() (facts *Facts, err error) {
	var (
		cpuCnt, i, pos int
		v              byte
		name, value    string
		cpu            Fact
	)
	prof.Reset()
	facts = &Facts{Timestamp: time.Now().UTC().UnixNano()}
	for {
		prof.Line, err = prof.Buf.ReadSlice('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, fmt.Errorf("error reading output bytes: %s", err)
		}
		// First grab the attribute name; everything up to the ':'.  The key may have
		// spaces and has trailing spaces; that gets trimmed.
		for i, v = range prof.Line {
			if v == 0x3A {
				pos = i + 1
				break
			}
			prof.Val = append(prof.Val, v)
		}
		name = strings.TrimSpace(string(prof.Val[:]))
		prof.Val = prof.Val[:0]
		// if there's anything left, the value is everything else; trim spaces
		if pos < len(prof.Line) {
			value = strings.TrimSpace(string(prof.Line[pos:]))
		}
		// check to see if this is flat.Facts for a different processor
		if name == "processor" {
			if cpuCnt > 0 {
				facts.CPU = append(facts.CPU, cpu)
			}
			cpuCnt++
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "fact: processor", Err: err}
			}
			cpu = Fact{Processor: int16(i)}
			continue
		}
		if name == "vendor_id" {
			cpu.VendorID = value
			continue
		}
		if name == "cpu family" {
			cpu.CPUFamily = value
			continue
		}
		if name == "model" {
			cpu.Model = value
			continue
		}
		if name == "model name" {
			cpu.ModelName = value
			continue
		}
		if name == "stepping" {
			cpu.Stepping = value
			continue
		}
		if name == "microcode" {
			cpu.Microcode = value
			continue
		}
		if name == "cpu MHz" {
			f, err := strconv.ParseFloat(value, 32)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: cpu MHz", Err: err}
			}
			cpu.CPUMHz = float32(f)
			continue
		}
		if name == "cache size" {
			cpu.CacheSize = value
			continue
		}
		if name == "physical id" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: physical id", Err: err}
			}
			cpu.PhysicalID = int16(i)
			continue
		}
		if name == "siblings" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: siblings", Err: err}
			}
			cpu.Siblings = int16(i)
			continue
		}
		if name == "core id" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: core id", Err: err}
			}
			cpu.CoreID = int16(i)
			continue
		}
		if name == "cpu cores" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: cpu cores", Err: err}
			}
			cpu.CPUCores = int16(i)
			continue
		}
		if name == "apicid" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: apicid", Err: err}
			}
			cpu.ApicID = int16(i)
			continue
		}
		if name == "initial apicid" {
			i, err = strconv.Atoi(value)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: initial apicid", Err: err}
			}
			cpu.InitialApicID = int16(i)
			continue
		}
		if name == "fpu" {
			cpu.FPU = value
			continue
		}
		if name == "fpu_exception" {
			cpu.FPUException = value
			continue
		}
		if name == "cpuid level" {
			cpu.CPUIDLevel = value
			continue
		}
		if name == "WP" {
			cpu.WP = value
			continue
		}
		if name == "flags" {
			cpu.Flags = value
			continue
		}
		if name == "bogomips" {
			f, err := strconv.ParseFloat(value, 32)
			if err != nil {
				return nil, joe.Error{Type: "cpu", Op: "facts: bogomips", Err: err}
			}
			cpu.BogoMIPS = float32(f)
			continue
		}
		if name == "clflush size" {
			cpu.CLFlushSize = value
			continue
		}
		if name == "cache_alignment" {
			cpu.CacheAlignment = value
			continue
		}
		if name == "address sizes" {
			cpu.AddressSizes = value
			continue
		}
		if name == "power management" {
			cpu.PowerManagement = value
		}
	}
	facts.CPU = append(facts.CPU, cpu)
	return facts, nil
}

// TODO: is it even worth it to have this as a global?  Should GetFacts
// just instantiate a local version and use that?
var stdProfiler *Profiler
var stdProfilerMu sync.Mutex

func Get() (facts *Facts, err error) {
	stdProfilerMu.Lock()
	defer stdProfilerMu.Unlock()
	if stdProfiler == nil {
		stdProfiler, err = New()
		if err != nil {
			return nil, err
		}
	}
	return stdProfiler.Get()
}

// Facts are a collection of facts, cpuinfo, about the system's cpus.
type Facts struct {
	Timestamp int64
	CPU       []Fact `json:"cpu"`
}

// Fact holds the /proc/cpuinfo for a single cpu
type Fact struct {
	Processor       int16   `json:"processor"`
	VendorID        string  `json:"vendor_id"`
	CPUFamily       string  `json:"cpu_family"`
	Model           string  `json:"model"`
	ModelName       string  `json:"model_name"`
	Stepping        string  `json:"stepping"`
	Microcode       string  `json:"microcode"`
	CPUMHz          float32 `json:"cpu_mhz"`
	CacheSize       string  `json:"cache_size"`
	PhysicalID      int16   `json:"physical_id"`
	Siblings        int16   `json:"siblings"`
	CoreID          int16   `json:"core_id"`
	CPUCores        int16   `json:"cpu_cores"`
	ApicID          int16   `json:"apicid"`
	InitialApicID   int16   `json:"initial_apicid"`
	FPU             string  `json:"fpu"`
	FPUException    string  `json:"fpu_exception"`
	CPUIDLevel      string  `json:"cpuid_level"`
	WP              string  `json:"wp"`
	Flags           string  `json:"flags"` // should this be a []string?
	BogoMIPS        float32 `json:"bogomips"`
	CLFlushSize     string  `json:"clflush_size"`
	CacheAlignment  string  `json:"cache_alignment"`
	AddressSizes    string  `json:"address_sizes"`
	PowerManagement string  `json:"power_management"`
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"os"
	"os/exec"
)

var options = struct {
	Output  string
	Name    string
	Plain   bool
	Preview bool
}{}

func init() {
	flag.StringVar(&options.Output, "output", "", "write output to file")
	flag.StringVar(&options.Name, "name", "", "set command name")
	flag.BoolVar(&options.Plain, "plain", false, "treat comments as plain text")
	flag.BoolVar(&options.Preview, "preview", false, "preview output in man")
}

func getReader() Reader {
	if options.Plain {
		return NewPlainReader()
	} else {
		return NewMarkupReader()
	}
}

func getWriter() Writer {
	return NewTroffWriter()
}

func main() {
	flag.Parse()

	builder := NewBuilder(getReader(), getWriter())
	for _, arg := range flag.Args() {
		file, err := NewFile(arg)
		if err != nil {
			panic(err)
		}
		if options.Name != "" {
			file.Name = options.Name
		}
		text, err := builder.Build(file)
		if err != nil {
			panic(err)
		}
		if options.Preview {
			cmd := exec.Command("groff", "-Wall", "-mtty-char", "-mandoc", "-Tascii")
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			inp, err := cmd.StdinPipe()
			if err != nil {
				panic(err)
			}
			if err := cmd.Start(); err != nil {
				panic(err)
			}
			inp.Write([]byte(text))
			inp.Close()
			cmd.Wait()
		} else {
			if options.Output == "" {
				os.Stdout.Write([]byte(text))
			} else {
				dst, err := os.Create(options.Output)
				if err != nil {
					panic(err)
				}
				dst.Write([]byte(text))
				dst.Close()
			}
		}
	}
}
<commit_msg>provide documentation<commit_after>// mango - generate manual pages from Go source code
//
// Description:
//
// mango generates manual pages from the source code of Go commands.
// It aims to generate full-fledged manual pages soley based on the comments
// and flag function calls found inside Go source code.
//
// See Also:
//
// man(1), man-pages(7)
package main

import (
	"flag"
	"os"
	"os/exec"
)

var options = struct {
	Output  string
	Name    string
	Plain   bool
	Preview bool
}{}

func init() {
	// Write the manual page to file.
	flag.StringVar(&options.Output, "output", "", "write to `file`")
	// Set the manual page title to name.
	flag.StringVar(&options.Name, "title", "", "set title to `name`")
	// Ignore markup inside comments.
	flag.BoolVar(&options.Plain, "plain", false, "treat comments as plain text")
	// Preview the manual page with the man command.
	flag.BoolVar(&options.Preview, "preview", false, "preview with man")
}

func getReader() Reader {
	if options.Plain {
		return NewPlainReader()
	} else {
		return NewMarkupReader()
	}
}

func getWriter() Writer {
	return NewTroffWriter()
}

func main() {
	flag.Parse()

	builder := NewBuilder(getReader(), getWriter())
	for _, arg := range flag.Args() {
		file, err := NewFile(arg)
		if err != nil {
			panic(err)
		}
		if options.Name != "" {
			file.Name = options.Name
		}
		text, err := builder.Build(file)
		if err != nil {
			panic(err)
		}
		if options.Preview {
			cmd := exec.Command("groff", "-Wall", "-mtty-char", "-mandoc", "-Tascii")
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			inp, err := cmd.StdinPipe()
			if err != nil {
				panic(err)
			}
			if err := cmd.Start(); err != nil {
				panic(err)
			}
			inp.Write([]byte(text))
			inp.Close()
			cmd.Wait()
		} else {
			if options.Output == "" {
				os.Stdout.Write([]byte(text))
			} else {
				dst, err := os.Create(options.Output)
				if err != nil {
					panic(err)
				}
				dst.Write([]byte(text))
				dst.Close()
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2022 The Cloudprober Authors.
//
// 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 resources provides webpages related resources.
package resources

var Style = `
<style type="text/css">
body {
  font-family: "Roboto","Helvetica","Arial",sans-serif;
	font-size: 14px;
  line-height: 1.3;
}

.debugging {
  font-size: 12px;
}

table.status-list {
  border-collapse: collapse;
  border-spacing: 0;
	margin-top: 10px;
	margin-bottom: 20px;
	font-family: monospace;
}
table.status-list td,th {
  border: 1px solid gray;
  padding: 0.25em 0.5em;
	max-width: 200px;
}
pre {
    white-space: pre-wrap;
		word-wrap: break-word;
}
</style>
`
<commit_msg>[probestatus] Minor style fix.<commit_after>// Copyright 2022 The Cloudprober Authors.
//
// 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 resources provides webpages related resources.
package resources

var Style = `
<style type="text/css">
body {
  font-family: "Roboto","Helvetica","Arial",sans-serif;
	font-size: 14px;
  line-height: 1.3;
}

.debugging {
  font-size: 12px;
}

table.status-list {
  border-collapse: collapse;
  border-spacing: 0;
	margin-top: 10px;
	margin-bottom: 20px;
	font-family: monospace;
	font-size: 14px;
}
table.status-list td,th {
  border: 1px solid gray;
  padding: 0.25em 0.5em;
	max-width: 200px;
}
pre {
    white-space: pre-wrap;
		word-wrap: break-word;
}
</style>
`
<|endoftext|>


<commit_before>package digraph

import (
	"container/list"
	"errors"
	"fmt"
)

var (
	// ErrCycle is returned when creating an edge between two vertices would result
	// in a cycle in the digraph
	ErrCycle = errors.New("digraph: cycle between edges")

	// ErrEdgeExists is returned when an edge between two vertices already exists
	ErrEdgeExists = errors.New("digraph: edge already exists")

	// ErrVertexExists is returned when a vertex with the same value already exists
	ErrVertexExists = errors.New("digraph: vertex already exists")
)

// Vertex represents a vertex or "node" in the digraph
type Vertex interface{}

// Digraph represents a "digraph", or directed graph data structure
type Digraph struct {
	adjList     map[Vertex]AdjacencyList
	edgeCount   int
	vertexCount int
}

// New creates a new acyclic Digraph, and initializes its adjacency list
func New() *Digraph {
	return &Digraph{
		adjList: map[Vertex]AdjacencyList{},
	}
}

// AddVertex tries to add a new vertex to the root of the adjacency list on the digraph,
// returning true if a new vertex is added, or false if it already exists
func (d *Digraph) AddVertex(vertex Vertex) error {
	// Check for a previous, identical vertex
	if _, found := d.adjList[vertex]; found {
		return ErrVertexExists
	}

	// Add the vertex to the adjacency list, initialize a new linked-list
	d.adjList[vertex] = AdjacencyList{list.New()}
	d.vertexCount++

	return nil
}

// AddEdge tries to add a new edge between two vertices on the adjacency list,
// returning true if a new edge is added, or false if it already exists
func (d *Digraph) AddEdge(source Vertex, target Vertex) error {
	// Ensure vertices are not identical
	if source == target {
		return ErrCycle
	}

	// Add both vertices to the graph, ignoring if they already exist
	d.AddVertex(source)
	d.AddVertex(target)

	// Check if this digraph already has this edge
	if d.HasEdge(source, target) {
		// Return false, edge already exists
		return ErrEdgeExists
	}

	// Do a depth-first search from the target to the source to determine if a cycle will
	// result if this edge is created
	if d.DepthFirstSearch(target, source) {
		// Return false, a cycle will be created
		return ErrCycle
	}

	// Retrieve adjacency list
	adjList := d.adjList[source]

	// Target was not found, so add an edge between source and target
	adjList.list.PushBack(target)
	d.edgeCount++

	// Store adjacency list
	d.adjList[source] = adjList

	return nil
}

// discovered maps out which vertices have been discovered using Depth-First Search
var discovered map[Vertex]bool

// DepthFirstSearch searches the digraph for the target vertex, using the Depth-First
// Search algorithm
func (d *Digraph) DepthFirstSearch(source Vertex, target Vertex) bool {
	// Clear discovery map
	discovered = map[Vertex]bool{}

	// Begin recursive Depth-First Search, looking for all vertices reachable from source
	d.dfs(source)

	// Check if target was discovered during Depth-First Search
	result := discovered[target]

	// Clear discovery map, return result
	discovered = map[Vertex]bool{}
	return result
}

// dfs implements a recursive Depth-First Search algorithm
func (d *Digraph) dfs(target Vertex) {
	// Get the adjacency list for this vertex
	adjList := d.adjList[target]

	// Check all adjacent vertices
	for _, v := range adjList.Adjacent() {
		// Check if vertex has not been discovered
		if !discovered[v] {
			// Mark it as discovered, recursively continue traversal
			discovered[v] = true
			d.dfs(v)
		}
	}
}

// EdgeCount returns the number of edges in the digraph
func (d *Digraph) EdgeCount() int {
	return d.edgeCount
}

// HasEdge determines if the digraph has an existing edge between source and target,
// returning true if it does, or false if it does not
func (d *Digraph) HasEdge(source Vertex, target Vertex) bool {
	// Retrieve adjacency list for this source
	adjList := d.adjList[source]

	// Search for target vertex
	if v := adjList.Search(target); v != nil {
		// Vertex is adjacent, edge exists
		return true
	}

	// No result, edge does not exist
	return false
}

// Print displays a printed "tree" of the digraph to the console
func (d *Digraph) Print(root Vertex) error {
	// Check if the vertex actually exists
	if _, ok := d.adjList[root]; !ok {
		return errors.New("digraph: root node does not exist, cannot print graph")
	}

	// Begin recursive printing at the specified root vertex
	d.printRecursive(root, "")
	return nil
}

// printRecursive handles the printing of each vertex in "tree" form
func (d *Digraph) printRecursive(vertex Vertex, prefix string) {
	// Print the current vertex
	fmt.Println(prefix, "-", vertex)

	// Get the current adjacency list, get adjacent vertices
	adjList := d.adjList[vertex]
	adjacent := adjList.Adjacent()

	// Iterate all adjacent vertices
	for i, v := range adjacent {
		// If last iteration, don't add a pipe character
		if i == len(adjacent)-1 {
			d.printRecursive(v, prefix+"    ")
		} else {
			// Add pipe character to show multiple items belong to same parent
			d.printRecursive(v, prefix+"   |")
		}
	}
}

// VertexCount returns the number of vertices in the digraph
func (d *Digraph) VertexCount() int {
	return d.vertexCount
}
<commit_msg>Fix some outdated comments on method descriptions<commit_after>package digraph

import (
	"container/list"
	"errors"
	"fmt"
)

var (
	// ErrCycle is returned when creating an edge between two vertices would result
	// in a cycle in the digraph
	ErrCycle = errors.New("digraph: cycle between edges")

	// ErrEdgeExists is returned when an edge between two vertices already exists
	ErrEdgeExists = errors.New("digraph: edge already exists")

	// ErrVertexExists is returned when a vertex with the same value already exists
	ErrVertexExists = errors.New("digraph: vertex already exists")
)

// Vertex represents a vertex or "node" in the digraph
type Vertex interface{}

// Digraph represents a "digraph", or directed graph data structure
type Digraph struct {
	adjList     map[Vertex]AdjacencyList
	edgeCount   int
	vertexCount int
}

// New creates a new acyclic Digraph, and initializes its adjacency list
func New() *Digraph {
	return &Digraph{
		adjList: map[Vertex]AdjacencyList{},
	}
}

// AddVertex tries to add a new vertex to the root of the adjacency list on the digraph
func (d *Digraph) AddVertex(vertex Vertex) error {
	// Check for a previous, identical vertex
	if _, found := d.adjList[vertex]; found {
		return ErrVertexExists
	}

	// Add the vertex to the adjacency list, initialize a new linked-list
	d.adjList[vertex] = AdjacencyList{list.New()}
	d.vertexCount++

	return nil
}

// AddEdge tries to add a new edge between two vertices on the adjacency list
func (d *Digraph) AddEdge(source Vertex, target Vertex) error {
	// Ensure vertices are not identical
	if source == target {
		return ErrCycle
	}

	// Add both vertices to the graph, ignoring if they already exist
	d.AddVertex(source)
	d.AddVertex(target)

	// Check if this digraph already has this edge
	if d.HasEdge(source, target) {
		// Return false, edge already exists
		return ErrEdgeExists
	}

	// Do a depth-first search from the target to the source to determine if a cycle will
	// result if this edge is created
	if d.DepthFirstSearch(target, source) {
		// Return false, a cycle will be created
		return ErrCycle
	}

	// Retrieve adjacency list
	adjList := d.adjList[source]

	// Target was not found, so add an edge between source and target
	adjList.list.PushBack(target)
	d.edgeCount++

	// Store adjacency list
	d.adjList[source] = adjList

	return nil
}

// discovered maps out which vertices have been discovered using Depth-First Search
var discovered map[Vertex]bool

// DepthFirstSearch searches the digraph for the target vertex, using the Depth-First
// Search algorithm, and returning true if a path to the target is found
func (d *Digraph) DepthFirstSearch(source Vertex, target Vertex) bool {
	// Clear discovery map
	discovered = map[Vertex]bool{}

	// Begin recursive Depth-First Search, looking for all vertices reachable from source
	d.dfs(source)

	// Check if target was discovered during Depth-First Search
	result := discovered[target]

	// Clear discovery map, return result
	discovered = map[Vertex]bool{}
	return result
}

// dfs implements a recursive Depth-First Search algorithm
func (d *Digraph) dfs(target Vertex) {
	// Get the adjacency list for this vertex
	adjList := d.adjList[target]

	// Check all adjacent vertices
	for _, v := range adjList.Adjacent() {
		// Check if vertex has not been discovered
		if !discovered[v] {
			// Mark it as discovered, recursively continue traversal
			discovered[v] = true
			d.dfs(v)
		}
	}
}

// EdgeCount returns the number of edges in the digraph
func (d *Digraph) EdgeCount() int {
	return d.edgeCount
}

// HasEdge determines if the digraph has an existing edge between source and target,
// returning true if it does, or false if it does not
func (d *Digraph) HasEdge(source Vertex, target Vertex) bool {
	// Retrieve adjacency list for this source
	adjList := d.adjList[source]

	// Search for target vertex
	if v := adjList.Search(target); v != nil {
		// Vertex is adjacent, edge exists
		return true
	}

	// No result, edge does not exist
	return false
}

// Print displays a printed "tree" of the digraph to the console
func (d *Digraph) Print(root Vertex) error {
	// Check if the vertex actually exists
	if _, ok := d.adjList[root]; !ok {
		return errors.New("digraph: root node does not exist, cannot print graph")
	}

	// Begin recursive printing at the specified root vertex
	d.printRecursive(root, "")
	return nil
}

// printRecursive handles the printing of each vertex in "tree" form
func (d *Digraph) printRecursive(vertex Vertex, prefix string) {
	// Print the current vertex
	fmt.Println(prefix, "-", vertex)

	// Get the current adjacency list, get adjacent vertices
	adjList := d.adjList[vertex]
	adjacent := adjList.Adjacent()

	// Iterate all adjacent vertices
	for i, v := range adjacent {
		// If last iteration, don't add a pipe character
		if i == len(adjacent)-1 {
			d.printRecursive(v, prefix+"    ")
		} else {
			// Add pipe character to show multiple items belong to same parent
			d.printRecursive(v, prefix+"   |")
		}
	}
}

// VertexCount returns the number of vertices in the digraph
func (d *Digraph) VertexCount() int {
	return d.vertexCount
}
<|endoftext|>


<commit_before>package db_test

import (
	"github.com/deiwin/luncher-api/db"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

var (
	dbClient              *db.Client
	offersCollection      db.Offers
	tagsCollection        db.Tags
	restaurantsCollection db.Restaurants
	mocks                 *Mocks
)

func TestDb(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Db Suite")
}

var _ = BeforeSuite(func(done Done) {
	defer close(done)
	mocks = createMocks()
	createClient()
	wipeDb()
	initCollections()
})

var _ = AfterSuite(func(done Done) {
	defer close(done)
	wipeDb()
	dbClient.Disconnect()
})

var _ = It("should work", func() {})

func createClient() {
	dbConfig := createTestDbConf()
	dbClient = db.NewClient(dbConfig)
	err := dbClient.Connect()
	Expect(err).NotTo(HaveOccurred())
}

func initCollections() {
	initOffersCollection()
	initTagsCollection()
	initRestaurantsCollection()
}

func initOffersCollection() {
	offersCollection = db.NewOffers(dbClient)
	err := insertOffers()
	Expect(err).NotTo(HaveOccurred())
}

func initTagsCollection() {
	tagsCollection = db.NewTags(dbClient)
	err := insertTags()
	Expect(err).NotTo(HaveOccurred())
}

func initRestaurantsCollection() {
	restaurantsCollection = db.NewRestaurants(dbClient)
	err := insertRestaurants()
	Expect(err).NotTo(HaveOccurred())
}

func createTestDbConf() (dbConfig *db.Config) {
	dbConfig = db.NewConfig()
	dbConfig.DbURL = "mongodb://localhost/test"
	dbConfig.DbName = "test"
	return
}

func insertTags() (err error) {
	return tagsCollection.Insert(mocks.tags...)
}

func insertRestaurants() (err error) {
	return restaurantsCollection.Insert(mocks.restaurants...)
}

func insertOffers() (err error) {
	return offersCollection.Insert(mocks.offers...)
}

func wipeDb() {
	err := dbClient.DropDb()
	Expect(err).NotTo(HaveOccurred())
}
<commit_msg>remove unnecessary env var reading in db test conf creation<commit_after>package db_test

import (
	"github.com/deiwin/luncher-api/db"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

var (
	dbClient              *db.Client
	offersCollection      db.Offers
	tagsCollection        db.Tags
	restaurantsCollection db.Restaurants
	mocks                 *Mocks
)

func TestDb(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Db Suite")
}

var _ = BeforeSuite(func(done Done) {
	defer close(done)
	mocks = createMocks()
	createClient()
	wipeDb()
	initCollections()
})

var _ = AfterSuite(func(done Done) {
	defer close(done)
	wipeDb()
	dbClient.Disconnect()
})

var _ = It("should work", func() {})

func createClient() {
	dbConfig := createTestDbConf()
	dbClient = db.NewClient(dbConfig)
	err := dbClient.Connect()
	Expect(err).NotTo(HaveOccurred())
}

func initCollections() {
	initOffersCollection()
	initTagsCollection()
	initRestaurantsCollection()
}

func initOffersCollection() {
	offersCollection = db.NewOffers(dbClient)
	err := insertOffers()
	Expect(err).NotTo(HaveOccurred())
}

func initTagsCollection() {
	tagsCollection = db.NewTags(dbClient)
	err := insertTags()
	Expect(err).NotTo(HaveOccurred())
}

func initRestaurantsCollection() {
	restaurantsCollection = db.NewRestaurants(dbClient)
	err := insertRestaurants()
	Expect(err).NotTo(HaveOccurred())
}

func createTestDbConf() (dbConfig *db.Config) {
	dbConfig = &db.Config{
		DbURL:  "mongodb://localhost/test",
		DbName: "test",
	}
	return
}

func insertTags() (err error) {
	return tagsCollection.Insert(mocks.tags...)
}

func insertRestaurants() (err error) {
	return restaurantsCollection.Insert(mocks.restaurants...)
}

func insertOffers() (err error) {
	return offersCollection.Insert(mocks.offers...)
}

func wipeDb() {
	err := dbClient.DropDb()
	Expect(err).NotTo(HaveOccurred())
}
<|endoftext|>


<commit_before>package main

import (
	"container/ring"
	"fmt"
	"strings"

	"github.com/bwmarrin/discordgo"
	"go.uber.org/zap"
)

// Bind represents a link between a Discord channel and a Redis queue
type Bind struct {
	DiscordChannel string `json:"channel_id"`
	InputQueue     string `json:"input_queue"`
	OutputQueue    string `json:"output_queue"`
}

// Message represents a message moving through this app
type Message struct {
	User string
	Text string
	// when in the `MessagesToDiscord` queue, this represents a Discord channel
	// and when in the `MessagesToGame` queue, it represents a queue name.
	Destination string
}

// DiscordManager holds program state and shares access to resources
type DiscordManager struct {
	app               *App
	DiscordClient     *discordgo.Session
	MessagesToDiscord chan Message // queue of messages waiting to be sent to Discord
	MessagesToGame    chan Message // queue of messages waiting to be sent to the game server
	MessageHistory    *ring.Ring   // ring-list of the last n messages processed to block duplicates
}

// NewDiscordManager sets up a Discord client and prepares it for starting
func NewDiscordManager(app *App) *DiscordManager {
	var err error

	dm := DiscordManager{
		app: app,
	}

	dm.app.config = app.config

	dm.DiscordClient, err = discordgo.New("Bot " + dm.app.config.DiscordToken)
	if err != nil {
		logger.Fatal("failed to create discord client", zap.Error(err))
	}

	dm.MessagesToDiscord = make(chan Message)
	dm.MessagesToGame = make(chan Message)
	dm.MessageHistory = ring.New(32)

	dm.Connect()

	return &dm
}

// Close cleans up and shuts down the Discord client
func (dm *DiscordManager) Close() error {
	return dm.DiscordClient.Close()
}

// Connect prepares the Discord client library and connects to the API.
func (dm *DiscordManager) Connect() {
	// Once the connection is ready, the client is prepared and the daemon multiplexer is started
	dm.DiscordClient.AddHandler(func(s *discordgo.Session, event *discordgo.Ready) {
		dm.Prepare()
		dm.Daemon()
	})

	err := dm.DiscordClient.Open()
	if err != nil {
		logger.Fatal("discord client connection error", zap.Error(err))
	}
}

// Prepare adds a handler for messages from Discord and a handler for messages from the game.
func (dm *DiscordManager) Prepare() {
	dm.AddDiscordHandler()
	dm.AddGameHandler()
}

// Daemon is a simple mutiplex select across four processes:
// - messages coming from Discord and waiting to get sent to the game
// - messages coming from the game and waiting to get sent to Discord
// - messages being consumed and sent to the game
// - messages being consumed and sent to Discord
// This is a blocking function and exits on fatal error.
func (dm *DiscordManager) Daemon() {
	for {
		select {
		case message := <-dm.MessagesToDiscord:
			logger.Debug("send to discord",
				zap.String("username", message.User),
				zap.String("message", message.Text),
				zap.String("Destination", message.Destination))

			dm.MessageHistory.Value = message.Text
			dm.MessageHistory.Next()
			Text := fmt.Sprintf("%s: %s", message.User, message.Text)

			sr, err := dm.DiscordClient.ChannelMessageSend(message.Destination, Text)
			if err != nil {
				logger.Warn("ChannelMessageSend error",
					zap.Error(err),
					zap.Any("sr", sr))
			}

		case message := <-dm.MessagesToGame:
			logger.Info("send to game",
				zap.String("username", message.User),
				zap.String("message", message.Text),
				zap.String("Destination", message.Destination))

			dm.MessageHistory.Value = message.Text
			dm.MessageHistory.Next()
			raw := fmt.Sprintf("%s:%s", message.User, message.Text)

			err := dm.app.sc.SendMessage(message.Destination, raw)
			if err != nil {
				logger.Warn("rm.SendMessage failed", zap.Error(err))
			}
		}
	}
}

// AddDiscordHandler uses the Discord client gateway event system to bind the Discord MessageCreate
// event to a function that consumes a message and places it on the message queue for being sent to
// the game.
func (dm DiscordManager) AddDiscordHandler() {
	dm.DiscordClient.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
		Destination := dm.GetOutgoingKeyFromChannelID(m.Message.ChannelID)

		if Destination != "" {
			split := strings.SplitN(m.Message.Content, ": ", 2)
			duplicate := false
			if len(split) >= 2 {
				dm.MessageHistory.Do(func(i interface{}) {
					if i == split[1] {
						duplicate = true
					}
				})
			}

			if !duplicate {
				dm.MessagesToGame <- Message{m.Message.Author.Username, m.Message.Content, Destination}
			}
		}
	})
}

// AddGameHandler uses the samp-go library to bind to a Redis list that contains chat messages sent
// from the game and places them on a message queue to be sent to the Discord channel.
func (dm DiscordManager) AddGameHandler() {
	for i := range dm.app.config.DiscordBinds {
		dm.app.sc.BindMessage(dm.app.config.DiscordBinds[i].InputQueue, func(message string) {
			split := strings.SplitN(message, ":", 2)

			if len(split) != 2 {
				logger.Warn("received from game: message malformed, no colon delimiter",
					zap.String("message", message))
				return
			}

			duplicate := false
			dm.MessageHistory.Do(func(i interface{}) {
				if i == split[1] {
					duplicate = true
				}
			})

			if !duplicate {
				dm.MessagesToDiscord <- Message{split[0], split[1], dm.app.config.DiscordBinds[i].DiscordChannel}
			}
		})
	}
}

// GetOutgoingKeyFromChannelID takes a Discord channel ID and returns a Redis queue if it is
// associated with one, otherwise returns an empty string.
func (dm DiscordManager) GetOutgoingKeyFromChannelID(channel string) string {
	for i := range dm.app.config.DiscordBinds {
		if channel == dm.app.config.DiscordBinds[i].DiscordChannel {
			return dm.app.config.DiscordBinds[i].OutputQueue
		}
	}
	return ""
}

// GetChannelFromInQueue takes a Redis queue name and returns a Discord channel ID if it is
// associated with one, otherwise returns an empty string.
func (dm DiscordManager) GetChannelFromInQueue(queue string) string {
	for i := range dm.app.config.DiscordBinds {
		if queue == dm.app.config.DiscordBinds[i].InputQueue {
			return dm.app.config.DiscordBinds[i].DiscordChannel
		}
	}
	return ""
}
<commit_msg>added redis key builder<commit_after>package main

import (
	"container/ring"
	"fmt"
	"strings"

	"github.com/bwmarrin/discordgo"
	"go.uber.org/zap"
)

// Bind represents a link between a Discord channel and a Redis queue
type Bind struct {
	DiscordChannel string `json:"channel_id"`
	InputQueue     string `json:"input_queue"`
	OutputQueue    string `json:"output_queue"`
}

// Message represents a message moving through this app
type Message struct {
	User string
	Text string
	// when in the `MessagesToDiscord` queue, this represents a Discord channel
	// and when in the `MessagesToGame` queue, it represents a queue name.
	Destination string
}

// DiscordManager holds program state and shares access to resources
type DiscordManager struct {
	app               *App
	DiscordClient     *discordgo.Session
	MessagesToDiscord chan Message // queue of messages waiting to be sent to Discord
	MessagesToGame    chan Message // queue of messages waiting to be sent to the game server
	MessageHistory    *ring.Ring   // ring-list of the last n messages processed to block duplicates
}

// NewDiscordManager sets up a Discord client and prepares it for starting
func NewDiscordManager(app *App) *DiscordManager {
	var err error

	dm := DiscordManager{
		app: app,
	}

	dm.app.config = app.config

	dm.DiscordClient, err = discordgo.New("Bot " + dm.app.config.DiscordToken)
	if err != nil {
		logger.Fatal("failed to create discord client", zap.Error(err))
	}

	dm.MessagesToDiscord = make(chan Message)
	dm.MessagesToGame = make(chan Message)
	dm.MessageHistory = ring.New(32)

	dm.Connect()

	return &dm
}

// Close cleans up and shuts down the Discord client
func (dm *DiscordManager) Close() error {
	return dm.DiscordClient.Close()
}

// Connect prepares the Discord client library and connects to the API.
func (dm *DiscordManager) Connect() {
	// Once the connection is ready, the client is prepared and the daemon multiplexer is started
	dm.DiscordClient.AddHandler(func(s *discordgo.Session, event *discordgo.Ready) {
		dm.Prepare()
		dm.Daemon()
	})

	err := dm.DiscordClient.Open()
	if err != nil {
		logger.Fatal("discord client connection error", zap.Error(err))
	}
}

// Prepare adds a handler for messages from Discord and a handler for messages from the game.
func (dm *DiscordManager) Prepare() {
	dm.AddDiscordHandler()
	dm.AddGameHandler()
}

// Daemon is a simple mutiplex select across four processes:
// - messages coming from Discord and waiting to get sent to the game
// - messages coming from the game and waiting to get sent to Discord
// - messages being consumed and sent to the game
// - messages being consumed and sent to Discord
// This is a blocking function and exits on fatal error.
func (dm *DiscordManager) Daemon() {
	for {
		select {
		case message := <-dm.MessagesToDiscord:
			logger.Debug("send to discord",
				zap.String("username", message.User),
				zap.String("message", message.Text),
				zap.String("Destination", message.Destination))

			dm.MessageHistory.Value = message.Text
			dm.MessageHistory.Next()
			Text := fmt.Sprintf("%s: %s", message.User, message.Text)

			sr, err := dm.DiscordClient.ChannelMessageSend(message.Destination, Text)
			if err != nil {
				logger.Warn("ChannelMessageSend error",
					zap.Error(err),
					zap.Any("sr", sr))
			}

		case message := <-dm.MessagesToGame:
			logger.Info("send to game",
				zap.String("username", message.User),
				zap.String("message", message.Text),
				zap.String("Destination", message.Destination))

			dm.MessageHistory.Value = message.Text
			dm.MessageHistory.Next()
			raw := fmt.Sprintf("%s:%s", message.User, message.Text)

			err := dm.app.sc.SendMessage(message.Destination, raw)
			if err != nil {
				logger.Warn("rm.SendMessage failed", zap.Error(err))
			}
		}
	}
}

// AddDiscordHandler uses the Discord client gateway event system to bind the Discord MessageCreate
// event to a function that consumes a message and places it on the message queue for being sent to
// the game.
func (dm DiscordManager) AddDiscordHandler() {
	dm.DiscordClient.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
		Destination := dm.GetOutgoingKeyFromChannelID(m.Message.ChannelID)

		if Destination != "" {
			split := strings.SplitN(m.Message.Content, ": ", 2)
			duplicate := false
			if len(split) >= 2 {
				dm.MessageHistory.Do(func(i interface{}) {
					if i == split[1] {
						duplicate = true
					}
				})
			}

			if !duplicate {
				dm.MessagesToGame <- Message{m.Message.Author.Username, m.Message.Content, Destination}
			}
		}
	})
}

// AddGameHandler uses the samp-go library to bind to a Redis list that contains chat messages sent
// from the game and places them on a message queue to be sent to the Discord channel.
func (dm DiscordManager) AddGameHandler() {
	for i := range dm.app.config.DiscordBinds {
		dm.app.sc.BindMessage(dm.GetFullRedisKey(dm.app.config.DiscordBinds[i].InputQueue), func(message string) {
			split := strings.SplitN(message, ":", 2)

			if len(split) != 2 {
				logger.Warn("received from game: message malformed, no colon delimiter",
					zap.String("message", message))
				return
			}

			duplicate := false
			dm.MessageHistory.Do(func(i interface{}) {
				if i == split[1] {
					duplicate = true
				}
			})

			if !duplicate {
				dm.MessagesToDiscord <- Message{split[0], split[1], dm.app.config.DiscordBinds[i].DiscordChannel}
			}
		})
	}
}

// GetOutgoingKeyFromChannelID takes a Discord channel ID and returns a Redis queue if it is
// associated with one, otherwise returns an empty string.
func (dm DiscordManager) GetOutgoingKeyFromChannelID(channel string) string {
	for i := range dm.app.config.DiscordBinds {
		if channel == dm.app.config.DiscordBinds[i].DiscordChannel {
			return dm.GetFullRedisKey(dm.app.config.DiscordBinds[i].OutputQueue)
		}
	}
	return ""
}

// GetChannelFromInQueue takes a Redis queue name and returns a Discord channel ID if it is
// associated with one, otherwise returns an empty string.
func (dm DiscordManager) GetChannelFromInQueue(queue string) string {
	for i := range dm.app.config.DiscordBinds {
		if queue == dm.GetFullRedisKey(dm.app.config.DiscordBinds[i].InputQueue) {
			return dm.app.config.DiscordBinds[i].DiscordChannel
		}
	}
	return ""
}

// GetFullRedisKey returns a full redis key for naming queues in the form: myserver.rediscord.queue
func (dm DiscordManager) GetFullRedisKey(name string) string {
	return dm.app.config.Domain + ".rediscord." + name
}
<|endoftext|>