<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package mungers

import (
	"time"

	"github.com/golang/glog"
	"github.com/google/go-github/github"
	"github.com/spf13/cobra"
	"k8s.io/contrib/mungegithub/features"
	mgh "k8s.io/contrib/mungegithub/github"
	c "k8s.io/contrib/mungegithub/mungers/matchers/comment"
	"k8s.io/contrib/mungegithub/mungers/mungerutil"
)

const (
	flakeNagNotifName = "FLAKE-PING"
	defaultTimePeriod = 14 * 24 * time.Hour
)

var (
	pinger = c.NewPinger(flakeNagNotifName).
		SetDescription("This flaky-test issue would love to have more attention...")
	// Only include priorities that you care about. Others won't be pinged
	timePeriods = map[string]time.Duration{
		"priority/P0": 2 * 24 * time.Hour,
		"priority/P1": 4 * 24 * time.Hour,
		"priority/P2": 7 * 24 * time.Hour,
	}
)

// NagFlakeIssues pings assignees on flaky-test issues
type NagFlakeIssues struct{}

var _ Munger = &NagFlakeIssues{}

func init() {
	n := &NagFlakeIssues{}
	RegisterMungerOrDie(n)
	RegisterStaleComments(n)
}

// Name is the name usable in --pr-mungers
func (NagFlakeIssues) Name() string { return "nag-flake-issues" }

// RequiredFeatures is a slice of 'features' that must be provided
func (NagFlakeIssues) RequiredFeatures() []string { return []string{} }

// Initialize will initialize the munger
func (NagFlakeIssues) Initialize(config *mgh.Config, features *features.Features) error {
	return nil
}

// EachLoop is called at the start of every munge loop
func (NagFlakeIssues) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (NagFlakeIssues) AddFlags(cmd *cobra.Command, config *mgh.Config) {
}

// findTimePeriod returns how often we should ping based on priority
func findTimePeriod(labels []github.Label) time.Duration {
	priorities := mgh.GetLabelsWithPrefix(labels, "priority/")
	if len(priorities) == 0 {
		return defaultTimePeriod
	}
	// If we have multiple priority labels (shouldn't happen), use the first one
	period, ok := timePeriods[priorities[0]]
	if !ok {
		return defaultTimePeriod
	}
	return period
}

// Munge is the workhorse the will actually make updates to the PR
func (NagFlakeIssues) Munge(obj *mgh.MungeObject) {
	if obj.IsPR() || !obj.HasLabel("kind/flake") {
		return
	}

	comments, err := obj.ListComments()
	if err != nil {
		glog.Error(err)
		return
	}

	// Use the pinger to notify assignees:
	// - Set time period based on configuration (at the top of this file)
	// - Mention list of assignees as an argument
	// - Start the ping timer after the last HumanActor comment

	// How often should we ping
	period := findTimePeriod(obj.Issue.Labels)

	// Who are we pinging
	who := mungerutil.GetIssueUsers(obj.Issue).Assignees.Mention().Join()

	// When does the pinger start
	startDate := c.LastComment(comments, c.HumanActor(), obj.Issue.CreatedAt)

	// Get a notification if it's time to ping.
	notif := pinger.SetTimePeriod(period).PingNotification(
		comments,
		who,
		startDate,
	)
	if notif != nil {
		obj.WriteComment(notif.String())
	}
}

// StaleComments returns a slice of stale comments
func (NagFlakeIssues) StaleComments(obj *mgh.MungeObject, comments []*github.IssueComment) []*github.IssueComment {
	// Remove all pings written before the last human actor comment
	return c.FilterComments(comments, c.And([]c.Matcher{
		c.MungerNotificationName(flakeNagNotifName),
		c.CreatedBefore(*c.LastComment(comments, c.HumanActor(), &time.Time{})),
	}))
}
<commit_msg>nag-flaky-issue: Do not ping if no assignee<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package mungers

import (
	"time"

	"github.com/golang/glog"
	"github.com/google/go-github/github"
	"github.com/spf13/cobra"
	"k8s.io/contrib/mungegithub/features"
	mgh "k8s.io/contrib/mungegithub/github"
	c "k8s.io/contrib/mungegithub/mungers/matchers/comment"
	"k8s.io/contrib/mungegithub/mungers/mungerutil"
)

const (
	flakeNagNotifName = "FLAKE-PING"
	defaultTimePeriod = 14 * 24 * time.Hour
)

var (
	pinger = c.NewPinger(flakeNagNotifName).
		SetDescription("This flaky-test issue would love to have more attention.")
	// Only include priorities that you care about. Others won't be pinged
	timePeriods = map[string]time.Duration{
		"priority/P0": 2 * 24 * time.Hour,
		"priority/P1": 4 * 24 * time.Hour,
		"priority/P2": 7 * 24 * time.Hour,
	}
)

// NagFlakeIssues pings assignees on flaky-test issues
type NagFlakeIssues struct{}

var _ Munger = &NagFlakeIssues{}

func init() {
	n := &NagFlakeIssues{}
	RegisterMungerOrDie(n)
	RegisterStaleComments(n)
}

// Name is the name usable in --pr-mungers
func (NagFlakeIssues) Name() string { return "nag-flake-issues" }

// RequiredFeatures is a slice of 'features' that must be provided
func (NagFlakeIssues) RequiredFeatures() []string { return []string{} }

// Initialize will initialize the munger
func (NagFlakeIssues) Initialize(config *mgh.Config, features *features.Features) error {
	return nil
}

// EachLoop is called at the start of every munge loop
func (NagFlakeIssues) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (NagFlakeIssues) AddFlags(cmd *cobra.Command, config *mgh.Config) {
}

// findTimePeriod returns how often we should ping based on priority
func findTimePeriod(labels []github.Label) time.Duration {
	priorities := mgh.GetLabelsWithPrefix(labels, "priority/")
	if len(priorities) == 0 {
		return defaultTimePeriod
	}
	// If we have multiple priority labels (shouldn't happen), use the first one
	period, ok := timePeriods[priorities[0]]
	if !ok {
		return defaultTimePeriod
	}
	return period
}

// Munge is the workhorse the will actually make updates to the PR
func (NagFlakeIssues) Munge(obj *mgh.MungeObject) {
	if obj.IsPR() || !obj.HasLabel("kind/flake") {
		return
	}

	comments, err := obj.ListComments()
	if err != nil {
		glog.Error(err)
		return
	}

	// Use the pinger to notify assignees:
	// - Set time period based on configuration (at the top of this file)
	// - Mention list of assignees as an argument
	// - Start the ping timer after the last HumanActor comment

	// How often should we ping
	period := findTimePeriod(obj.Issue.Labels)

	// Who are we pinging
	who := mungerutil.GetIssueUsers(obj.Issue).Assignees.Mention().Join()
	if who == "" {
		return
	}

	// When does the pinger start
	startDate := c.LastComment(comments, c.HumanActor(), obj.Issue.CreatedAt)

	// Get a notification if it's time to ping.
	notif := pinger.SetTimePeriod(period).PingNotification(
		comments,
		who,
		startDate,
	)
	if notif != nil {
		obj.WriteComment(notif.String())
	}
}

// StaleComments returns a slice of stale comments
func (NagFlakeIssues) StaleComments(obj *mgh.MungeObject, comments []*github.IssueComment) []*github.IssueComment {
	// Remove all pings written before the last human actor comment
	return c.FilterComments(comments, c.And([]c.Matcher{
		c.MungerNotificationName(flakeNagNotifName),
		c.CreatedBefore(*c.LastComment(comments, c.HumanActor(), &time.Time{})),
	}))
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes 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 dco

import (
	"fmt"
	"strings"

	"github.com/sirupsen/logrus"

	"regexp"

	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/pluginhelp"
	"k8s.io/test-infra/prow/plugins"
)

const (
	pluginName               = "dco"
	dcoContextName           = "dco"
	dcoContextMessageFailed  = "Commits in PR missing Signed-off-by"
	dcoContextMessageSuccess = "All commits have Signed-off-by"

	dcoYesLabel        = "dco-signoff: yes"
	dcoNoLabel         = "dco-signoff: no"
	dcoMsgPruneMatch   = "Thanks for your pull request. Before we can look at it, you'll need to add a 'DCO signoff' to your commits."
	dcoNotFoundMessage = `Thanks for your pull request. Before we can look at it, you'll need to add a 'DCO signoff' to your commits.

:memo: **Please follow instructions in the [contributing guide](%s) to update your commits with the DCO**

Full details of the Developer Certificate of Origin can be found at [developercertificate.org](https://developercertificate.org/).

**The list of commits missing DCO signoff**:

%s

<details>

%s
</details>
`
)

var (
	checkDCORe = regexp.MustCompile(`(?mi)^/check-dco\s*$`)
	testRe     = regexp.MustCompile(`(?mi)^signed-off-by:`)
)

func init() {
	plugins.RegisterPullRequestHandler(pluginName, handlePullRequestEvent, helpProvider)
	plugins.RegisterGenericCommentHandler(pluginName, handleCommentEvent, helpProvider)
}

func helpProvider(config *plugins.Configuration, enabledRepos []string) (*pluginhelp.PluginHelp, error) {
	// The {WhoCanUse, Usage, Examples, Config} fields are omitted because this plugin cannot be
	// manually triggered and is not configurable.
	pluginHelp := &pluginhelp.PluginHelp{
		Description: "The dco plugin checks pull request commits for 'DCO sign off' and maintains the '" + dcoContextName + "' status context, as well as the 'dco' label.",
	}
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/check-dco",
		Description: "Forces rechecking of the DCO status.",
		Featured:    true,
		WhoCanUse:   "Anyone",
		Examples:    []string{"/check-dco"},
	})
	return pluginHelp, nil
}

type gitHubClient interface {
	BotName() (string, error)
	CreateComment(owner, repo string, number int, comment string) error
	GetIssueLabels(org, repo string, number int) ([]github.Label, error)
	AddLabel(owner, repo string, number int, label string) error
	RemoveLabel(owner, repo string, number int, label string) error
	ListStatuses(org, repo, ref string) ([]github.Status, error)
	CreateStatus(owner, repo, ref string, status github.Status) error
	ListPRCommits(org, repo string, number int) ([]github.RepositoryCommit, error)
	GetPullRequest(owner, repo string, number int) (*github.PullRequest, error)
}

type commentPruner interface {
	PruneComments(shouldPrune func(github.IssueComment) bool)
}

// checkCommitMessages will perform the actual DCO check by retrieving all
// commits contained within the PR with the given number.
// *All* commits in the pull request *must* match the 'testRe' in order to pass.
func checkCommitMessages(gc gitHubClient, l *logrus.Entry, org, repo string, number int) ([]github.GitCommit, error) {
	allCommits, err := gc.ListPRCommits(org, repo, number)
	if err != nil {
		return nil, fmt.Errorf("error listing commits for pull request: %v", err)
	}
	l.Infof("Found %d commits in PR", len(allCommits))

	var commitsMissingDCO []github.GitCommit
	for _, commit := range allCommits {
		if !testRe.MatchString(*commit.Commit.Message) {
			c := *commit.Commit
			c.SHA = &(*commit.SHA)
			commitsMissingDCO = append(commitsMissingDCO, c)
		}
	}

	l.Infof("All commits in PR have DCO signoff: %t", len(commitsMissingDCO) == 0)
	return commitsMissingDCO, nil
}

// checkExistingStatus will retrieve the current status of the DCO context for
// the provided SHA.
func checkExistingStatus(gc gitHubClient, l *logrus.Entry, org, repo, sha string) (string, error) {
	statuses, err := gc.ListStatuses(org, repo, sha)
	if err != nil {
		return "", fmt.Errorf("error listing pull request statuses: %v", err)
	}

	existingStatus := ""
	for _, status := range statuses {
		if status.Context != dcoContextName {
			continue
		}
		existingStatus = status.State
		break
	}
	l.Infof("Existing DCO status context status is %q", existingStatus)
	return existingStatus, nil
}

// checkExistingLabels will check the provided PR for the dco sign off labels,
// returning bool's indicating whether the 'yes' and the 'no' label are present.
func checkExistingLabels(gc gitHubClient, l *logrus.Entry, org, repo string, number int) (hasYesLabel, hasNoLabel bool, err error) {
	labels, err := gc.GetIssueLabels(org, repo, number)
	if err != nil {
		return false, false, fmt.Errorf("error getting pull request labels: %v", err)
	}

	for _, l := range labels {
		if l.Name == dcoYesLabel {
			hasYesLabel = true
		}
		if l.Name == dcoNoLabel {
			hasNoLabel = true
		}
	}

	return hasYesLabel, hasNoLabel, nil
}

// takeAction will take appropriate action on the pull request according to its
// current state.
func takeAction(gc gitHubClient, cp commentPruner, l *logrus.Entry, org, repo string, pr github.PullRequest, commitsMissingDCO []github.GitCommit, existingStatus string, hasYesLabel, hasNoLabel, addComment bool) error {
	targetURL := fmt.Sprintf("https://github.com/%s/%s/blob/master/CONTRIBUTING.md", org, repo)
	botName, err := gc.BotName()
	if err != nil {
		return fmt.Errorf("failed to get bot name: %v", err)
	}

	signedOff := len(commitsMissingDCO) == 0

	// handle the 'all commits signed off' case by adding appropriate labels
	// TODO: clean-up old comments?
	if signedOff {
		if hasNoLabel {
			l.Infof("Removing %q label", dcoNoLabel)
			// remove 'dco-signoff: no' label
			if err := gc.RemoveLabel(org, repo, pr.Number, dcoNoLabel); err != nil {
				return fmt.Errorf("error removing label: %v", err)
			}
		}
		if !hasYesLabel {
			l.Infof("Adding %q label", dcoYesLabel)
			// add 'dco-signoff: yes' label
			if err := gc.AddLabel(org, repo, pr.Number, dcoYesLabel); err != nil {
				return fmt.Errorf("error adding label: %v", err)
			}
		}
		if existingStatus != github.StatusSuccess {
			l.Infof("Setting DCO status context to succeeded")
			if err := gc.CreateStatus(org, repo, pr.Head.SHA, github.Status{
				Context:     dcoContextName,
				State:       github.StatusSuccess,
				TargetURL:   targetURL,
				Description: dcoContextMessageSuccess,
			}); err != nil {
				return fmt.Errorf("error setting pull request status: %v", err)
			}
		}

		cp.PruneComments(shouldPrune(l, botName))
		return nil
	}

	// handle the 'not all commits signed off' case
	if !hasNoLabel {
		l.Infof("Adding %q label", dcoNoLabel)
		// add 'dco-signoff: no' label
		if err := gc.AddLabel(org, repo, pr.Number, dcoNoLabel); err != nil {
			return fmt.Errorf("error adding label: %v", err)
		}
	}
	if hasYesLabel {
		l.Infof("Removing %q label", dcoYesLabel)
		// remove 'dco-signoff: yes' label
		if err := gc.RemoveLabel(org, repo, pr.Number, dcoYesLabel); err != nil {
			return fmt.Errorf("error removing label: %v", err)
		}
	}
	if existingStatus != github.StatusFailure {
		l.Infof("Setting DCO status context to failed")
		if err := gc.CreateStatus(org, repo, pr.Head.SHA, github.Status{
			Context:     dcoContextName,
			State:       github.StatusFailure,
			TargetURL:   targetURL,
			Description: dcoContextMessageFailed,
		}); err != nil {
			return fmt.Errorf("error setting pull request status: %v", err)
		}
	}

	if addComment {
		// we only prune and recreate the comment on *synchronize* events to
		// prevent us commenting every time a label is added/deleted.
		cp.PruneComments(shouldPrune(l, botName))
		l.Infof("Commenting on PR to advise users of DCO check")
		if err := gc.CreateComment(org, repo, pr.Number, fmt.Sprintf(dcoNotFoundMessage, targetURL, markdownSHAList(org, repo, commitsMissingDCO), plugins.AboutThisBot)); err != nil {
			l.WithError(err).Warning("Could not create DCO not found comment.")
		}
	}

	return nil
}

// 1. Check commit messages in the pull request for the sign-off string
// 2. Check the existing status context value
// 3. Check the existing PR labels
// 4. If signed off, apply appropriate labels and status context.
// 5. If not signed off, apply appropriate labels and status context and add a comment.
func handle(gc gitHubClient, cp commentPruner, log *logrus.Entry, org, repo string, pr github.PullRequest, addComment bool) error {
	l := log.WithField("pr", pr.Number)

	commitsMissingDCO, err := checkCommitMessages(gc, l, org, repo, pr.Number)
	if err != nil {
		l.WithError(err).Infof("Error running DCO check against commits in PR")
		return err
	}

	existingStatus, err := checkExistingStatus(gc, l, org, repo, pr.Head.SHA)
	if err != nil {
		l.WithError(err).Infof("Error checking existing PR status")
		return err
	}

	hasYesLabel, hasNoLabel, err := checkExistingLabels(gc, l, org, repo, pr.Number)
	if err != nil {
		l.WithError(err).Infof("Error checking existing PR labels")
		return err
	}

	return takeAction(gc, cp, l, org, repo, pr, commitsMissingDCO, existingStatus, hasYesLabel, hasNoLabel, addComment)
}

func markdownSHAList(org, repo string, list []github.GitCommit) string {
	lines := make([]string, len(list))
	lineFmt := "- [%s](https://github.com/%s/%s/commits/%s) %s"
	for i, commit := range list {
		if commit.SHA == nil {
			continue
		}
		// if we somehow encounter a SHA that's less than 7 characters, we will
		// just use it as is.
		shortSHA := *commit.SHA
		if len(shortSHA) > 7 {
			shortSHA = shortSHA[:7]
		}
		message := ""
		if commit.Message != nil {
			message = *commit.Message
		}

		// get the first line of the commit
		message = strings.Split(message, "\n")[0]

		lines[i] = fmt.Sprintf(lineFmt, shortSHA, org, repo, *commit.SHA, message)
	}
	return strings.Join(lines, "\n")
}

// shouldPrune finds comments left by this plugin.
func shouldPrune(log *logrus.Entry, botName string) func(github.IssueComment) bool {
	return func(comment github.IssueComment) bool {
		if comment.User.Login != botName {
			return false
		}
		return strings.Contains(comment.Body, dcoMsgPruneMatch)
	}
}

func handlePullRequestEvent(pc plugins.PluginClient, pe github.PullRequestEvent) error {
	org := pe.Repo.Owner.Login
	repo := pe.Repo.Name

	// we only reprocess on label, unlabel, open, reopen and synchronize events
	// this will reduce our API token usage and save processing of unrelated events
	switch pe.Action {
	case github.PullRequestActionLabeled,
		github.PullRequestActionUnlabeled,
		github.PullRequestActionOpened,
		github.PullRequestActionReopened,
		github.PullRequestActionSynchronize:
	default:
		return nil
	}

	shouldComment := pe.Action == github.PullRequestActionSynchronize || pe.Action == github.PullRequestActionOpened

	return handle(pc.GitHubClient, pc.CommentPruner, pc.Logger, org, repo, pe.PullRequest, shouldComment)
}

func handleCommentEvent(pc plugins.PluginClient, ce github.GenericCommentEvent) error {
	// Only consider open PRs and new comments.
	if ce.IssueState != "open" || ce.Action != github.GenericCommentActionCreated || !ce.IsPR {
		return nil
	}
	// Only consider "/check-dco" comments.
	if !checkDCORe.MatchString(ce.Body) {
		return nil
	}

	gc := pc.GitHubClient
	org := ce.Repo.Owner.Login
	repo := ce.Repo.Name

	pr, err := gc.GetPullRequest(org, repo, ce.Number)
	if err != nil {
		return fmt.Errorf("error getting pull request for comment: %v", err)
	}

	return handle(gc, pc.CommentPruner, pc.Logger, org, repo, *pr, true)
}
<commit_msg>Split handlePullRequestEvent and handleCommentEvent for testing<commit_after>/*
Copyright 2018 The Kubernetes 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 dco

import (
	"fmt"
	"strings"

	"github.com/sirupsen/logrus"

	"regexp"

	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/pluginhelp"
	"k8s.io/test-infra/prow/plugins"
)

const (
	pluginName               = "dco"
	dcoContextName           = "dco"
	dcoContextMessageFailed  = "Commits in PR missing Signed-off-by"
	dcoContextMessageSuccess = "All commits have Signed-off-by"

	dcoYesLabel        = "dco-signoff: yes"
	dcoNoLabel         = "dco-signoff: no"
	dcoMsgPruneMatch   = "Thanks for your pull request. Before we can look at it, you'll need to add a 'DCO signoff' to your commits."
	dcoNotFoundMessage = `Thanks for your pull request. Before we can look at it, you'll need to add a 'DCO signoff' to your commits.

:memo: **Please follow instructions in the [contributing guide](%s) to update your commits with the DCO**

Full details of the Developer Certificate of Origin can be found at [developercertificate.org](https://developercertificate.org/).

**The list of commits missing DCO signoff**:

%s

<details>

%s
</details>
`
)

var (
	checkDCORe = regexp.MustCompile(`(?mi)^/check-dco\s*$`)
	testRe     = regexp.MustCompile(`(?mi)^signed-off-by:`)
)

func init() {
	plugins.RegisterPullRequestHandler(pluginName, handlePullRequestEvent, helpProvider)
	plugins.RegisterGenericCommentHandler(pluginName, handleCommentEvent, helpProvider)
}

func helpProvider(config *plugins.Configuration, enabledRepos []string) (*pluginhelp.PluginHelp, error) {
	// The {WhoCanUse, Usage, Examples, Config} fields are omitted because this plugin cannot be
	// manually triggered and is not configurable.
	pluginHelp := &pluginhelp.PluginHelp{
		Description: "The dco plugin checks pull request commits for 'DCO sign off' and maintains the '" + dcoContextName + "' status context, as well as the 'dco' label.",
	}
	pluginHelp.AddCommand(pluginhelp.Command{
		Usage:       "/check-dco",
		Description: "Forces rechecking of the DCO status.",
		Featured:    true,
		WhoCanUse:   "Anyone",
		Examples:    []string{"/check-dco"},
	})
	return pluginHelp, nil
}

type gitHubClient interface {
	BotName() (string, error)
	CreateComment(owner, repo string, number int, comment string) error
	GetIssueLabels(org, repo string, number int) ([]github.Label, error)
	AddLabel(owner, repo string, number int, label string) error
	RemoveLabel(owner, repo string, number int, label string) error
	ListStatuses(org, repo, ref string) ([]github.Status, error)
	CreateStatus(owner, repo, ref string, status github.Status) error
	ListPRCommits(org, repo string, number int) ([]github.RepositoryCommit, error)
	GetPullRequest(owner, repo string, number int) (*github.PullRequest, error)
}

type commentPruner interface {
	PruneComments(shouldPrune func(github.IssueComment) bool)
}

// checkCommitMessages will perform the actual DCO check by retrieving all
// commits contained within the PR with the given number.
// *All* commits in the pull request *must* match the 'testRe' in order to pass.
func checkCommitMessages(gc gitHubClient, l *logrus.Entry, org, repo string, number int) ([]github.GitCommit, error) {
	allCommits, err := gc.ListPRCommits(org, repo, number)
	if err != nil {
		return nil, fmt.Errorf("error listing commits for pull request: %v", err)
	}
	l.Infof("Found %d commits in PR", len(allCommits))

	var commitsMissingDCO []github.GitCommit
	for _, commit := range allCommits {
		if !testRe.MatchString(*commit.Commit.Message) {
			c := *commit.Commit
			c.SHA = &(*commit.SHA)
			commitsMissingDCO = append(commitsMissingDCO, c)
		}
	}

	l.Infof("All commits in PR have DCO signoff: %t", len(commitsMissingDCO) == 0)
	return commitsMissingDCO, nil
}

// checkExistingStatus will retrieve the current status of the DCO context for
// the provided SHA.
func checkExistingStatus(gc gitHubClient, l *logrus.Entry, org, repo, sha string) (string, error) {
	statuses, err := gc.ListStatuses(org, repo, sha)
	if err != nil {
		return "", fmt.Errorf("error listing pull request statuses: %v", err)
	}

	existingStatus := ""
	for _, status := range statuses {
		if status.Context != dcoContextName {
			continue
		}
		existingStatus = status.State
		break
	}
	l.Infof("Existing DCO status context status is %q", existingStatus)
	return existingStatus, nil
}

// checkExistingLabels will check the provided PR for the dco sign off labels,
// returning bool's indicating whether the 'yes' and the 'no' label are present.
func checkExistingLabels(gc gitHubClient, l *logrus.Entry, org, repo string, number int) (hasYesLabel, hasNoLabel bool, err error) {
	labels, err := gc.GetIssueLabels(org, repo, number)
	if err != nil {
		return false, false, fmt.Errorf("error getting pull request labels: %v", err)
	}

	for _, l := range labels {
		if l.Name == dcoYesLabel {
			hasYesLabel = true
		}
		if l.Name == dcoNoLabel {
			hasNoLabel = true
		}
	}

	return hasYesLabel, hasNoLabel, nil
}

// takeAction will take appropriate action on the pull request according to its
// current state.
func takeAction(gc gitHubClient, cp commentPruner, l *logrus.Entry, org, repo string, pr github.PullRequest, commitsMissingDCO []github.GitCommit, existingStatus string, hasYesLabel, hasNoLabel, addComment bool) error {
	targetURL := fmt.Sprintf("https://github.com/%s/%s/blob/master/CONTRIBUTING.md", org, repo)
	botName, err := gc.BotName()
	if err != nil {
		return fmt.Errorf("failed to get bot name: %v", err)
	}

	signedOff := len(commitsMissingDCO) == 0

	// handle the 'all commits signed off' case by adding appropriate labels
	// TODO: clean-up old comments?
	if signedOff {
		if hasNoLabel {
			l.Infof("Removing %q label", dcoNoLabel)
			// remove 'dco-signoff: no' label
			if err := gc.RemoveLabel(org, repo, pr.Number, dcoNoLabel); err != nil {
				return fmt.Errorf("error removing label: %v", err)
			}
		}
		if !hasYesLabel {
			l.Infof("Adding %q label", dcoYesLabel)
			// add 'dco-signoff: yes' label
			if err := gc.AddLabel(org, repo, pr.Number, dcoYesLabel); err != nil {
				return fmt.Errorf("error adding label: %v", err)
			}
		}
		if existingStatus != github.StatusSuccess {
			l.Infof("Setting DCO status context to succeeded")
			if err := gc.CreateStatus(org, repo, pr.Head.SHA, github.Status{
				Context:     dcoContextName,
				State:       github.StatusSuccess,
				TargetURL:   targetURL,
				Description: dcoContextMessageSuccess,
			}); err != nil {
				return fmt.Errorf("error setting pull request status: %v", err)
			}
		}

		cp.PruneComments(shouldPrune(l, botName))
		return nil
	}

	// handle the 'not all commits signed off' case
	if !hasNoLabel {
		l.Infof("Adding %q label", dcoNoLabel)
		// add 'dco-signoff: no' label
		if err := gc.AddLabel(org, repo, pr.Number, dcoNoLabel); err != nil {
			return fmt.Errorf("error adding label: %v", err)
		}
	}
	if hasYesLabel {
		l.Infof("Removing %q label", dcoYesLabel)
		// remove 'dco-signoff: yes' label
		if err := gc.RemoveLabel(org, repo, pr.Number, dcoYesLabel); err != nil {
			return fmt.Errorf("error removing label: %v", err)
		}
	}
	if existingStatus != github.StatusFailure {
		l.Infof("Setting DCO status context to failed")
		if err := gc.CreateStatus(org, repo, pr.Head.SHA, github.Status{
			Context:     dcoContextName,
			State:       github.StatusFailure,
			TargetURL:   targetURL,
			Description: dcoContextMessageFailed,
		}); err != nil {
			return fmt.Errorf("error setting pull request status: %v", err)
		}
	}

	if addComment {
		// we only prune and recreate the comment on *synchronize* events to
		// prevent us commenting every time a label is added/deleted.
		cp.PruneComments(shouldPrune(l, botName))
		l.Infof("Commenting on PR to advise users of DCO check")
		if err := gc.CreateComment(org, repo, pr.Number, fmt.Sprintf(dcoNotFoundMessage, targetURL, markdownSHAList(org, repo, commitsMissingDCO), plugins.AboutThisBot)); err != nil {
			l.WithError(err).Warning("Could not create DCO not found comment.")
		}
	}

	return nil
}

// 1. Check commit messages in the pull request for the sign-off string
// 2. Check the existing status context value
// 3. Check the existing PR labels
// 4. If signed off, apply appropriate labels and status context.
// 5. If not signed off, apply appropriate labels and status context and add a comment.
func handle(gc gitHubClient, cp commentPruner, log *logrus.Entry, org, repo string, pr github.PullRequest, addComment bool) error {
	l := log.WithField("pr", pr.Number)

	commitsMissingDCO, err := checkCommitMessages(gc, l, org, repo, pr.Number)
	if err != nil {
		l.WithError(err).Infof("Error running DCO check against commits in PR")
		return err
	}

	existingStatus, err := checkExistingStatus(gc, l, org, repo, pr.Head.SHA)
	if err != nil {
		l.WithError(err).Infof("Error checking existing PR status")
		return err
	}

	hasYesLabel, hasNoLabel, err := checkExistingLabels(gc, l, org, repo, pr.Number)
	if err != nil {
		l.WithError(err).Infof("Error checking existing PR labels")
		return err
	}

	return takeAction(gc, cp, l, org, repo, pr, commitsMissingDCO, existingStatus, hasYesLabel, hasNoLabel, addComment)
}

func markdownSHAList(org, repo string, list []github.GitCommit) string {
	lines := make([]string, len(list))
	lineFmt := "- [%s](https://github.com/%s/%s/commits/%s) %s"
	for i, commit := range list {
		if commit.SHA == nil {
			continue
		}
		// if we somehow encounter a SHA that's less than 7 characters, we will
		// just use it as is.
		shortSHA := *commit.SHA
		if len(shortSHA) > 7 {
			shortSHA = shortSHA[:7]
		}
		message := ""
		if commit.Message != nil {
			message = *commit.Message
		}

		// get the first line of the commit
		message = strings.Split(message, "\n")[0]

		lines[i] = fmt.Sprintf(lineFmt, shortSHA, org, repo, *commit.SHA, message)
	}
	return strings.Join(lines, "\n")
}

// shouldPrune finds comments left by this plugin.
func shouldPrune(log *logrus.Entry, botName string) func(github.IssueComment) bool {
	return func(comment github.IssueComment) bool {
		if comment.User.Login != botName {
			return false
		}
		return strings.Contains(comment.Body, dcoMsgPruneMatch)
	}
}

func handlePullRequestEvent(pc plugins.PluginClient, pe github.PullRequestEvent) error {
	return handlePullRequest(pc.GitHubClient, pc.CommentPruner, pc.Logger, pe)
}

func handlePullRequest(gc gitHubClient, cp commentPruner, log *logrus.Entry, pe github.PullRequestEvent) error {
	org := pe.Repo.Owner.Login
	repo := pe.Repo.Name

	// we only reprocess on label, unlabel, open, reopen and synchronize events
	// this will reduce our API token usage and save processing of unrelated events
	switch pe.Action {
	case github.PullRequestActionLabeled,
		github.PullRequestActionUnlabeled,
		github.PullRequestActionOpened,
		github.PullRequestActionReopened,
		github.PullRequestActionSynchronize:
	default:
		return nil
	}

	shouldComment := pe.Action == github.PullRequestActionSynchronize || pe.Action == github.PullRequestActionOpened

	return handle(gc, cp, log, org, repo, pe.PullRequest, shouldComment)
}

func handleCommentEvent(pc plugins.PluginClient, ce github.GenericCommentEvent) error {
	return handleComment(pc.GitHubClient, pc.CommentPruner, pc.Logger, ce)
}

func handleComment(gc gitHubClient, cp commentPruner, log *logrus.Entry, ce github.GenericCommentEvent) error {
	// Only consider open PRs and new comments.
	if ce.IssueState != "open" || ce.Action != github.GenericCommentActionCreated || !ce.IsPR {
		return nil
	}
	// Only consider "/check-dco" comments.
	if !checkDCORe.MatchString(ce.Body) {
		return nil
	}

	org := ce.Repo.Owner.Login
	repo := ce.Repo.Name

	pr, err := gc.GetPullRequest(org, repo, ce.Number)
	if err != nil {
		return fmt.Errorf("error getting pull request for comment: %v", err)
	}

	return handle(gc, cp, log, org, repo, *pr, true)
}
<|endoftext|>


<commit_before>package slog

import (
	stdLog "log"

	"github.com/davecgh/go-spew/spew"
)

var (
	Fatal   = stdLog.Fatal
	Fatalf  = stdLog.Fatalf
	Fatalln = stdLog.Fatalln
	Panic   = stdLog.Panic
	Panicf  = stdLog.Panicf
	Panicln = stdLog.Panicln
	Prefix  = stdLog.Prefix
	Print   = stdLog.Print
	Printf  = stdLog.Printf
	Println = stdLog.Println
)

func init() {
	stdLog.SetFlags(0)
}

func Dump(v ...interface{}) {
	spew.Dump(v...)
}

func FatalIf(err error) {
	if err != nil {
		stdLog.Fatal(err)
	}
}

func PanicIf(err error) {
	if err != nil {
		stdLog.Panic(err)
	}
}

func SetPrefix(s string) {
	stdLog.SetPrefix("[" + s + "] ")
}
<commit_msg>PanicIff<commit_after>package slog

import (
	stdLog "log"

	"github.com/davecgh/go-spew/spew"
)

var (
	Fatal   = stdLog.Fatal
	Fatalf  = stdLog.Fatalf
	Fatalln = stdLog.Fatalln
	Panic   = stdLog.Panic
	Panicf  = stdLog.Panicf
	Panicln = stdLog.Panicln
	Prefix  = stdLog.Prefix
	Print   = stdLog.Print
	Printf  = stdLog.Printf
	Println = stdLog.Println
)

func init() {
	stdLog.SetFlags(0)
}

func Dump(v ...interface{}) {
	spew.Dump(v...)
}

func FatalIf(err error) {
	if err != nil {
		stdLog.Fatal(err)
	}
}

func PanicIf(err error) {
	if err != nil {
		stdLog.Panic(err)
	}
}

func PanicIff(err error, template string, v ...interface{}) {
	if err != nil {
		stdLog.Panicf(template, v...)
	}
}

func SetPrefix(s string) {
	stdLog.SetPrefix("[" + s + "] ")
}
<|endoftext|>


<commit_before>package jman

import (
	"bytes"
	"errors"
	"io/ioutil"
	"strconv"
	"strings"
	"sync"
)

var (
	// ErrSpecificNode is for when retrieving a node does not return a specific key/value, but perhaps a map
	ErrSpecificNode = errors.New("Could not find a specific node that matched the given path")
)

// JFile represents a JSON file and contains the filename and root node
type JFile struct {
	filename string
	rootnode *Node
	rw       *sync.RWMutex
}

// NewFile will read the given filename and return a JFile struct
func NewFile(filename string) (*JFile, error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	js, err := New(data)
	if err != nil {
		return nil, err
	}
	rw := &sync.RWMutex{}
	return &JFile{filename, js, rw}, nil
}

// Add two byte slices together
func badd(a, b []byte) []byte {
	var buf bytes.Buffer
	buf.Write(a)
	buf.Write(b)
	return buf.Bytes()
}

// Recursively look up a given JSON path
func jsonpath(js *Node, JSONpath string) (*Node, string, error) {
	if JSONpath == "" {
		// Could not find end node
		return js, JSONpath, ErrSpecificNode
	}
	firstpart := JSONpath
	secondpart := ""
	if strings.Contains(JSONpath, ".") {
		fields := strings.SplitN(JSONpath, ".", 2)
		firstpart = fields[0]
		secondpart = fields[1]
	}
	if firstpart == "x" {
		return jsonpath(js, secondpart)
	} else if strings.Contains(firstpart, "[") && strings.Contains(firstpart, "]") {
		fields := strings.SplitN(firstpart, "[", 2)
		name := fields[0]
		if name != "x" {
			js = js.Get(name)
		}
		fields = strings.SplitN(fields[1], "]", 2)
		index, err := strconv.Atoi(fields[0])
		if err != nil {
			return js, JSONpath, errors.New("Invalid index: " + fields[0] + " (" + err.Error() + ")")
		}
		node, ok := js.GetIndex(index)
		if !ok {
			return js, JSONpath, errors.New("Could not find index: " + fields[0])
		}
		return jsonpath(node, secondpart)
	}
	name := firstpart
	if secondpart != "" {
		return js, JSONpath, errors.New("JSON path left unparsed: " + secondpart)
	}
	return js.Get(name), "", nil
}

// GetNode will find the JSON node that corresponds to the given JSON path
func (jf *JFile) GetNode(JSONpath string) (*Node, error) {
	foundnode, leftoverpath, err := jsonpath(jf.rootnode, JSONpath)
	if err != nil {
		return nil, err
	}
	if leftoverpath != "" {
		return nil, errors.New("JSON path left unparsed: " + leftoverpath)
	}
	if foundnode == nil {
		return nil, errors.New("Could not lookup: " + JSONpath)
	}
	return foundnode, nil
}

// GetString will find the string that corresponds to the given JSON path
func (jf *JFile) GetString(JSONpath string) (string, error) {
	node, err := jf.GetNode(JSONpath)
	if err != nil {
		return "", err
	}
	return node.String(), nil
}

// SetString will change the value of the key that the given JSON path points to
func (jf *JFile) SetString(JSONpath, value string) error {
	firstpart := ""
	lastpart := JSONpath
	if strings.Contains(JSONpath, ".") {
		pos := strings.LastIndex(JSONpath, ".")
		firstpart = JSONpath[:pos]
		lastpart = JSONpath[pos+1:]
	}

	node, _, err := jsonpath(jf.rootnode, firstpart)
	if (err != nil) && (err != ErrSpecificNode) {
		return err
	}

	_, hasNode := node.CheckGet(lastpart)
	if !hasNode {
		return errors.New("Index out of range? Could not set value.")
	}

	// It's weird that simplejson Set does not return an error value
	node.Set(lastpart, value)

	newdata, err := jf.rootnode.EncodePretty()
	if err != nil {
		return err
	}

	return jf.Write(newdata)
}

// Write writes the current JSON data to the file
func (jf *JFile) Write(data []byte) error {
	jf.rw.Lock()
	defer jf.rw.Unlock()
	// TODO: Add newline as well?
	return ioutil.WriteFile(jf.filename, data, 0666)
}

// AddJSON adds JSON data at the given JSON path
func (jf *JFile) AddJSON(JSONpath, JSONdata string) error {
	firstpart := ""
	lastpart := JSONpath
	if strings.Contains(JSONpath, ".") {
		pos := strings.LastIndex(JSONpath, ".")
		firstpart = JSONpath[:pos]
		lastpart = JSONpath[pos+1:]
	}

	node, _, err := jsonpath(jf.rootnode, firstpart)
	if (err != nil) && (err != ErrSpecificNode) {
		return err
	}

	_, hasNode := node.CheckGet(lastpart)
	if hasNode {
		return errors.New("The JSON path should not point to a single key when adding JSON data.")
	}

	listJSON, err := node.Encode()
	if err != nil {
		return err
	}

	fullJSON, err := jf.rootnode.Encode()
	if err != nil {
		return err
	}

	// TODO: Implement a more efficient way of adding data.
	newFullJSON := bytes.Replace(fullJSON, listJSON, badd(listJSON[:len(listJSON)-1], []byte(","+JSONdata+"]")), 1)

	js, err := New(newFullJSON)
	if err != nil {
		return err
	}

	// Update the root node
	jf.rootnode = js

	newFullJSON, err = js.EncodePretty()
	if err != nil {
		return err
	}

	return jf.Write(newFullJSON)
}

// GetAll returns the current JSON data
func (jf *JFile) GetAll() ([]byte, error) {
	return jf.rootnode.EncodePretty()
}

// SetString sets a value to the given JSON file at the given JSON path
func SetString(filename, JSONpath, value string) error {
	jf, err := NewFile(filename)
	if err != nil {
		return err
	}
	return jf.SetString(JSONpath, value)
}

// AddJSON adds JSON data to the given JSON file at the given JSON path
func AddJSON(filename, JSONpath, JSONdata string) error {
	jf, err := NewFile(filename)
	if err != nil {
		return err
	}
	return jf.AddJSON(JSONpath, JSONdata)
}

// GetString will find the string that corresponds to the given JSON Path,
// given a filename and a simple JSON path expression.
func GetString(filename, JSONpath string) (string, error) {
	jf, err := NewFile(filename)
	if err != nil {
		return "", err
	}
	return jf.GetString(JSONpath)
}
<commit_msg>Adding items to empty lists<commit_after>package jman

import (
	"bytes"
	"errors"
	"io/ioutil"
	"strconv"
	"strings"
	"sync"
)

var (
	// ErrSpecificNode is for when retrieving a node does not return a specific key/value, but perhaps a map
	ErrSpecificNode = errors.New("Could not find a specific node that matched the given path")
)

// JFile represents a JSON file and contains the filename and root node
type JFile struct {
	filename string
	rootnode *Node
	rw       *sync.RWMutex
}

// NewFile will read the given filename and return a JFile struct
func NewFile(filename string) (*JFile, error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	js, err := New(data)
	if err != nil {
		return nil, err
	}
	rw := &sync.RWMutex{}
	return &JFile{filename, js, rw}, nil
}

// Add two byte slices together
func badd(a, b []byte) []byte {
	var buf bytes.Buffer
	buf.Write(a)
	buf.Write(b)
	return buf.Bytes()
}

// Recursively look up a given JSON path
func jsonpath(js *Node, JSONpath string) (*Node, string, error) {
	if JSONpath == "" {
		// Could not find end node
		return js, JSONpath, ErrSpecificNode
	}
	firstpart := JSONpath
	secondpart := ""
	if strings.Contains(JSONpath, ".") {
		fields := strings.SplitN(JSONpath, ".", 2)
		firstpart = fields[0]
		secondpart = fields[1]
	}
	if firstpart == "x" {
		return jsonpath(js, secondpart)
	} else if strings.Contains(firstpart, "[") && strings.Contains(firstpart, "]") {
		fields := strings.SplitN(firstpart, "[", 2)
		name := fields[0]
		if name != "x" {
			js = js.Get(name)
		}
		fields = strings.SplitN(fields[1], "]", 2)
		index, err := strconv.Atoi(fields[0])
		if err != nil {
			return js, JSONpath, errors.New("Invalid index: " + fields[0] + " (" + err.Error() + ")")
		}
		node, ok := js.GetIndex(index)
		if !ok {
			return js, JSONpath, errors.New("Could not find index: " + fields[0])
		}
		return jsonpath(node, secondpart)
	}
	name := firstpart
	if secondpart != "" {
		return js, JSONpath, errors.New("JSON path left unparsed: " + secondpart)
	}
	return js.Get(name), "", nil
}

// GetNode will find the JSON node that corresponds to the given JSON path
func (jf *JFile) GetNode(JSONpath string) (*Node, error) {
	foundnode, leftoverpath, err := jsonpath(jf.rootnode, JSONpath)
	if err != nil {
		return nil, err
	}
	if leftoverpath != "" {
		return nil, errors.New("JSON path left unparsed: " + leftoverpath)
	}
	if foundnode == nil {
		return nil, errors.New("Could not lookup: " + JSONpath)
	}
	return foundnode, nil
}

// GetString will find the string that corresponds to the given JSON path
func (jf *JFile) GetString(JSONpath string) (string, error) {
	node, err := jf.GetNode(JSONpath)
	if err != nil {
		return "", err
	}
	return node.String(), nil
}

// SetString will change the value of the key that the given JSON path points to
func (jf *JFile) SetString(JSONpath, value string) error {
	firstpart := ""
	lastpart := JSONpath
	if strings.Contains(JSONpath, ".") {
		pos := strings.LastIndex(JSONpath, ".")
		firstpart = JSONpath[:pos]
		lastpart = JSONpath[pos+1:]
	}

	node, _, err := jsonpath(jf.rootnode, firstpart)
	if (err != nil) && (err != ErrSpecificNode) {
		return err
	}

	_, hasNode := node.CheckGet(lastpart)
	if !hasNode {
		return errors.New("Index out of range? Could not set value.")
	}

	// It's weird that simplejson Set does not return an error value
	node.Set(lastpart, value)

	newdata, err := jf.rootnode.EncodePretty()
	if err != nil {
		return err
	}

	return jf.Write(newdata)
}

// Write writes the current JSON data to the file
func (jf *JFile) Write(data []byte) error {
	jf.rw.Lock()
	defer jf.rw.Unlock()
	// TODO: Add newline as well?
	return ioutil.WriteFile(jf.filename, data, 0666)
}

// AddJSON adds JSON data at the given JSON path
func (jf *JFile) AddJSON(JSONpath, JSONdata string) error {
	firstpart := ""
	lastpart := JSONpath
	if strings.Contains(JSONpath, ".") {
		pos := strings.LastIndex(JSONpath, ".")
		firstpart = JSONpath[:pos]
		lastpart = JSONpath[pos+1:]
	}

	node, _, err := jsonpath(jf.rootnode, firstpart)
	if (err != nil) && (err != ErrSpecificNode) {
		return err
	}

	_, hasNode := node.CheckGet(lastpart)
	if hasNode {
		return errors.New("The JSON path should not point to a single key when adding JSON data.")
	}

	listJSON, err := node.Encode()
	if err != nil {
		return err
	}

	fullJSON, err := jf.rootnode.Encode()
	if err != nil {
		return err
	}

	if len(listJSON) == 0 {
		return errors.New("Can not add to a completely empty list.")
	}

	// TODO: Implement a safer and more efficient way of adding data.
	var newFullJSON []byte
	if string(fullJSON) == "[]" {
		newFullJSON = []byte("[" + JSONdata + "]")
	} else {
		newFullJSON = bytes.Replace(fullJSON, listJSON, badd(listJSON[:len(listJSON)-1], []byte(","+JSONdata+"]")), 1)
	}
	js, err := New(newFullJSON)
	if err != nil {
		return err
	}

	// Update the root node
	jf.rootnode = js

	newFullJSON, err = js.EncodePretty()
	if err != nil {
		return err
	}

	return jf.Write(newFullJSON)
}

// GetAll returns the current JSON data
func (jf *JFile) GetAll() ([]byte, error) {
	return jf.rootnode.EncodePretty()
}

// SetString sets a value to the given JSON file at the given JSON path
func SetString(filename, JSONpath, value string) error {
	jf, err := NewFile(filename)
	if err != nil {
		return err
	}
	return jf.SetString(JSONpath, value)
}

// AddJSON adds JSON data to the given JSON file at the given JSON path
func AddJSON(filename, JSONpath, JSONdata string) error {
	jf, err := NewFile(filename)
	if err != nil {
		return err
	}
	return jf.AddJSON(JSONpath, JSONdata)
}

// GetString will find the string that corresponds to the given JSON Path,
// given a filename and a simple JSON path expression.
func GetString(filename, JSONpath string) (string, error) {
	jf, err := NewFile(filename)
	if err != nil {
		return "", err
	}
	return jf.GetString(JSONpath)
}
<|endoftext|>


<commit_before>package uuid

import (
	"fmt"
	"crypto/rand"
	"encoding/binary"
	"time"
)

func uint32rand() (value uint32) {
	err := binary.Read(rand.Reader, binary.LittleEndian, &value)
	if err != nil {
		panic(err)
	}
	return
}

// Generates a time ordered UUID. Top 32 bits are a timestamp,
// bottom 96 are random.
func TimeOrderedUUID() string {
	unix := uint32(time.Now().UTC().Unix())
	rand1 := uint32rand()
	rand2 := uint32rand()
	rand3 := uint32rand()
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%04x%08x",
		unix,
		uint16(rand1>>16),
		uint16(rand1&0xffff),
		uint16(rand2>>16),
		uint16(rand2&0xffff),
		rand3)
}
<commit_msg>Simplifies the implementation of common/uuid<commit_after>package uuid

import (
	"crypto/rand"
	"fmt"
	"time"
)

// Generates a time ordered UUID. Top 32 bits are a timestamp,
// bottom 96 are random.
func TimeOrderedUUID() string {
	unix := uint32(time.Now().UTC().Unix())

	b := make([]byte, 12)
	n, err := rand.Read(b)
	if n != len(b) {
		err = fmt.Errorf("Not enough entropy available")
	}
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%04x%08x",
		unix, b[0:2], b[2:4], b[4:6], b[6:8], b[8:])
}
<|endoftext|>


<commit_before>package imgmeta

import (
	//"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

var supportedExts = regexp.MustCompile(".+(jpg|png)$")

type ImgMeta interface {
	Name() string
	Path() string
	FullPath() string
	RelPath() string
	Meta() map[string]string
}

type imgMeta struct {
	basepath string
	name     string
	path     string
	meta     map[string]string
}

func (im *imgMeta) Name() string {
	return im.name
}

func (im *imgMeta) Path() string {
	return im.path
}

func (im *imgMeta) RelPath() string {
	rel, _ := filepath.Rel(im.basepath, im.FullPath())
	return rel
}

func (im *imgMeta) FullPath() string {
	return im.path + "/" + im.Name()
}

func (im *imgMeta) Meta() map[string]string {
	return im.meta
}

type ImgFolder interface {
	Name() string
	Path() string
	Folders() map[string]ImgFolder
	Images() []ImgMeta
	Find(path string) ImgFolder
	RelPath() string
}

type imgFolder struct {
	basepath string
	name     string
	path     string
	folders  map[string]ImgFolder
	images   []ImgMeta
}

func (im *imgFolder) RelPath() string {
	rel, _ := filepath.Rel(im.basepath, im.path)
	return rel
}

func (im *imgFolder) Name() string {
	return im.name
}

func (im *imgFolder) Path() string {
	return im.path
}

func (im *imgFolder) Folders() map[string]ImgFolder {
	return im.folders
}

func (im *imgFolder) Images() []ImgMeta {
	return im.images
}

func (im *imgFolder) Find(path string) ImgFolder {
	var parent ImgFolder
	folders := strings.Split(strings.Trim(path, "/"), "/")
	parent = im
	for _, f := range folders {
		ff := parent.Folders()
		p, ok := ff[f]
		if !ok {
			return nil
		}
		parent = p
	}
	return parent
}

func scanFolder(path string, basepath string, name string) (ImgFolder, error) {
	re := &imgFolder{
		basepath: basepath,
		path:     path,
		folders:  make(map[string]ImgFolder),
		name:     name,
	}
	folder, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	files, err := folder.Readdir(-1)

	if err != nil {
		return nil, err
	}
	for _, fi := range files {
		if fi.IsDir() {
			subfolder, err := scanFolder(path+"/"+fi.Name(), basepath, fi.Name())
			if err == nil {
				re.folders[fi.Name()] = subfolder
			}
		}
		if supportedExts.MatchString(fi.Name()) {
			img := &imgMeta{name: fi.Name(), path: path, basepath: basepath}
			img.meta = getMeta(img.FullPath())
			re.images = append(re.images, img)
		}
	}

	return re, nil
}

func Scan(path string) (ImgFolder, error) {
	p, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}
	return scanFolder(p, p, "")
}
<commit_msg>Should be case insensitive<commit_after>package imgmeta

import (
	//"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

var supportedExts = regexp.MustCompile(".+(?i)(jpg|png)$")

type ImgMeta interface {
	Name() string
	Path() string
	FullPath() string
	RelPath() string
	Meta() map[string]string
}

type imgMeta struct {
	basepath string
	name     string
	path     string
	meta     map[string]string
}

func (im *imgMeta) Name() string {
	return im.name
}

func (im *imgMeta) Path() string {
	return im.path
}

func (im *imgMeta) RelPath() string {
	rel, _ := filepath.Rel(im.basepath, im.FullPath())
	return rel
}

func (im *imgMeta) FullPath() string {
	return im.path + "/" + im.Name()
}

func (im *imgMeta) Meta() map[string]string {
	return im.meta
}

type ImgFolder interface {
	Name() string
	Path() string
	Folders() map[string]ImgFolder
	Images() []ImgMeta
	Find(path string) ImgFolder
	RelPath() string
}

type imgFolder struct {
	basepath string
	name     string
	path     string
	folders  map[string]ImgFolder
	images   []ImgMeta
}

func (im *imgFolder) RelPath() string {
	rel, _ := filepath.Rel(im.basepath, im.path)
	return rel
}

func (im *imgFolder) Name() string {
	return im.name
}

func (im *imgFolder) Path() string {
	return im.path
}

func (im *imgFolder) Folders() map[string]ImgFolder {
	return im.folders
}

func (im *imgFolder) Images() []ImgMeta {
	return im.images
}

func (im *imgFolder) Find(path string) ImgFolder {
	var parent ImgFolder
	folders := strings.Split(strings.Trim(path, "/"), "/")
	parent = im
	for _, f := range folders {
		ff := parent.Folders()
		p, ok := ff[f]
		if !ok {
			return nil
		}
		parent = p
	}
	return parent
}

func scanFolder(path string, basepath string, name string) (ImgFolder, error) {
	re := &imgFolder{
		basepath: basepath,
		path:     path,
		folders:  make(map[string]ImgFolder),
		name:     name,
	}
	folder, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	files, err := folder.Readdir(-1)

	if err != nil {
		return nil, err
	}
	for _, fi := range files {
		if fi.IsDir() {
			subfolder, err := scanFolder(path+"/"+fi.Name(), basepath, fi.Name())
			if err == nil {
				re.folders[fi.Name()] = subfolder
			}
		}
		if supportedExts.MatchString(fi.Name()) {
			img := &imgMeta{name: fi.Name(), path: path, basepath: basepath}
			img.meta = getMeta(img.FullPath())
			re.images = append(re.images, img)
		}
	}

	return re, nil
}

func Scan(path string) (ImgFolder, error) {
	p, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}
	return scanFolder(p, p, "")
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"github.com/hashicorp/consul/testutil"
	"github.com/hashicorp/serf/serf"
	"github.com/mitchellh/cli"
	"strings"
	"testing"
	"errors"
)

func TestForceLeaveCommand_implements(t *testing.T) {
	var _ cli.Command = &ForceLeaveCommand{}
}

func TestForceLeaveCommandRun(t *testing.T) {
	a1 := testAgent(t)
	a2 := testAgent(t)
	defer a1.Shutdown()
	defer a2.Shutdown()

	addr := fmt.Sprintf("127.0.0.1:%d", a2.config.Ports.SerfLan)
	_, err := a1.agent.JoinLAN([]string{addr})
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Forcibly shutdown a2 so that it appears "failed" in a1
	a2.Shutdown()

	ui := new(cli.MockUi)
	c := &ForceLeaveCommand{Ui: ui}
	args := []string{
		"-rpc-addr=" + a1.addr,
		a2.config.NodeName,
	}

	code := c.Run(args)
	if code != 0 {
		t.Fatalf("bad: %d. %#v", code, ui.ErrorWriter.String())
	}

	m := a1.agent.LANMembers()
	if len(m) != 2 {
		t.Fatalf("should have 2 members: %#v", m)
	}

	testutil.WaitForResult(func() (bool, error) {
		m = a1.agent.LANMembers()
		success := m[1].Status == serf.StatusLeft
		err := errors.New(m[1].Status.String())
		return success, err
	}, func(err error) {
		t.Fatalf("member status is %v, not left", err)
	})
}

func TestForceLeaveCommandRun_noAddrs(t *testing.T) {
	ui := new(cli.MockUi)
	c := &ForceLeaveCommand{Ui: ui}
	args := []string{"-rpc-addr=foo"}

	code := c.Run(args)
	if code != 1 {
		t.Fatalf("bad: %d", code)
	}

	if !strings.Contains(ui.ErrorWriter.String(), "node name") {
		t.Fatalf("bad: %#v", ui.ErrorWriter.String())
	}
}
<commit_msg>Fix formatting in `TestForceLeaveCommandRun`<commit_after>package command

import (
	"fmt"
	"github.com/hashicorp/consul/testutil"
	"github.com/hashicorp/serf/serf"
	"github.com/mitchellh/cli"
	"strings"
	"testing"
	"errors"
)

func TestForceLeaveCommand_implements(t *testing.T) {
	var _ cli.Command = &ForceLeaveCommand{}
}

func TestForceLeaveCommandRun(t *testing.T) {
	a1 := testAgent(t)
	a2 := testAgent(t)
	defer a1.Shutdown()
	defer a2.Shutdown()

	addr := fmt.Sprintf("127.0.0.1:%d", a2.config.Ports.SerfLan)
	_, err := a1.agent.JoinLAN([]string{addr})
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Forcibly shutdown a2 so that it appears "failed" in a1
	a2.Shutdown()

	ui := new(cli.MockUi)
	c := &ForceLeaveCommand{Ui: ui}
	args := []string{
		"-rpc-addr=" + a1.addr,
		a2.config.NodeName,
	}

	code := c.Run(args)
	if code != 0 {
		t.Fatalf("bad: %d. %#v", code, ui.ErrorWriter.String())
	}

	m := a1.agent.LANMembers()
	if len(m) != 2 {
		t.Fatalf("should have 2 members: %#v", m)
	}

	testutil.WaitForResult(func() (bool, error) {
		m = a1.agent.LANMembers()
		success := m[1].Status == serf.StatusLeft
		return success, errors.New(m[1].Status.String())
	}, func(err error) {
		t.Fatalf("member status is %v, should be left", err)
	})
}

func TestForceLeaveCommandRun_noAddrs(t *testing.T) {
	ui := new(cli.MockUi)
	c := &ForceLeaveCommand{Ui: ui}
	args := []string{"-rpc-addr=foo"}

	code := c.Run(args)
	if code != 1 {
		t.Fatalf("bad: %d", code)
	}

	if !strings.Contains(ui.ErrorWriter.String(), "node name") {
		t.Fatalf("bad: %#v", ui.ErrorWriter.String())
	}
}
<|endoftext|>


<commit_before>// for each target domain in the given prioritized list:
//     for front candidate in the given prioritized list:
//        check whether we can reach the target domain by attempting domain
//        fronting through the front candidate.  Log successfully proxied pairs
//        to stdout.  Log detected errors to stderr.
//
// run with `-tags debug` to enable debug logging.  See `-help` for other
// options.

// XXX: record where we left off, to enable resuming an interrupted scan.

package main

import (
	"bufio"
	"crypto/sha1"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"sync"
	"time"

	dbg "github.com/getlantern/sitescanner/log"
	"gopkg.in/getlantern/tlsdialer.v1"
)

var headRe *regexp.Regexp = regexp.MustCompile("(?s)(.*)")

var titleRe *regexp.Regexp = regexp.MustCompile("(?s)<title>(.*)</title>")

const reportFmt = "%-24s %-24s %-12s %-12s %-12s %-20s %s\n"

var dialTimeout = 10 * time.Second

var fronts = flag.String("fronts", "fronts.txt", "Path to the file containing front candidates, one domain per line.")

var targets = flag.String("targets", "targets.txt", "Path to the file containing fronting targets, one domain per line.")

var workers = flag.Int("workers", runtime.NumCPU() * 100, "Number of concurrent worker goroutines.")

// ResponseFeatures contains the characteristics we'll use when comparing
// direct vs domain fronted Responses, to estimate the likelihood that we are
// successfully fronting to a target.
//
// We are also checking whether we get a 200 OK response from the site at all,
// but we don't include this in ResponseFeatures becouse we are not reporting
// sites for which we don't get a 200 OK through domain fronting.
type ResponseFeatures struct {
	// Hash of the whole body, or "" if the body is the empty string
	BodyHash string
	// Hash of the response body's  element, or "" if it has none
	HeadHash string
	// Hash of the <title> element, or "" if the response body has none.
	TitleHash string
}

type Pairing struct {
	// The domain we're trying to reach.
	Target string
	// A URL where we can hopefully reach Target without redirects.
	TargetURL url.URL
	// The domain we're trying to domain front through.
	Front string
	// A URL where we can hopefully reach Front without redirects.
	FrontURL url.URL
	// A string of the form ip-address:port
	FrontIPPort string

	Features ResponseFeatures
}

func main() {
	flag.Parse()
	runtime.GOMAXPROCS(runtime.NumCPU())
	start := time.Now()
	fmt.Println()
	fmt.Printf(reportFmt, "TARGET", "FRONT", "BODIES", "HEADS", "TITLES", "FRONT IP:PORT", "REDIRECTED URLS")
	fmt.Printf(reportFmt, "======", "=====", "======", "=====", "======", "=============", "===============")
	fmt.Println()
	main_()
	fmt.Println("**************** COMPLETE ******************")
	fmt.Printf("Scan took %.2fs\n", time.Since(start).Seconds())
}

func main_() {

	tasksChan := make(chan Pairing, *workers)

	workersWg := sync.WaitGroup{}

	// Size of work queue is O(n^2) on the size of the input lists, so we
	// generate it lazily.
	go feedTasks(tasksChan)

	workersWg.Add(*workers)
	for i := 0; i < *workers; i++ {
		go work(tasksChan, &workersWg)
	}

	// Task completion handling:
	//
	// The feedTasks goroutine will close tasksChan when it has processed all
	// input.  Each worker reports workersWg.Done() when it learns that
	// tasksChan is closed.  By that time, the worker has sent all its results
	// to the resultsChan.  Therefore, by the time workersWg.Wait() returns,
	// all `Pairing`s have been checked and all the successful ones have been
	// printed out.
	workersWg.Wait()
}

func work(tasksChan <-chan Pairing, workersWg *sync.WaitGroup) {
	for task := range tasksChan {
		features, err := proxiedResponseFeatures(task)
		if err != nil {
			dbg.Debugln("proxiedResponseFeatures error:", err)
			continue
		}
		bodiesMatch := matchReport(task.Features.BodyHash, features.BodyHash)
		headsMatch := matchReport(task.Features.HeadHash, features.HeadHash)
		titlesMatch := matchReport(task.Features.TitleHash, features.TitleHash)
		if bodiesMatch != "match" && headsMatch != "match" && titlesMatch != "match" {
			dbg.Debugln("Got an OK response fronting", task.Target, "through", task.Front, "but nothing matches.")
			continue
		}
		fmt.Printf(
			reportFmt,
			task.Target,
			task.Front,
			bodiesMatch,
			headsMatch,
			titlesMatch,
			task.FrontIPPort,
			task.TargetURL.String() + ", " + task.FrontURL.String(),
		)
	}
	dbg.Debugln("One worker done")
	workersWg.Done()
}

func proxiedResponseFeatures(p Pairing) (ResponseFeatures, error) {
	dbg.Debugln("canProxy", p)
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				return tlsdialer.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					p.FrontURL.Host+":443",
					// We can't use a domain front that requires a properly
					// populated SNI, so let's make those fail.
					false,
					&tls.Config{InsecureSkipVerify: true},
				)
			},
		},
		CheckRedirect: noRedirect,
	}
	req, err := http.NewRequest("GET", p.TargetURL.String(), nil)
	if err != nil {
		logErr("building GET request", err)
		return ResponseFeatures{}, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return ResponseFeatures{}, err
	}
	return responseFeatures(resp)
}

// feedTasks reads input files, enqueues testing tasks in taskChan, and closes
// taskChan when done.
func feedTasks(taskChan chan<- Pairing) {
	fronts, err := readLines(*fronts)
	if err != nil {
		panic(err)
	}
	targets, err := readLines(*targets)
	if err != nil {
		panic(err)
	}

	// Domains for which we can't get a hash, even without domain fronting.
	badDomains := make(map[string]bool)

	// No-redirect URLs
	nru := make(map[string]url.URL)

	// No-redirect response features
	nrf := make(map[string]ResponseFeatures)

	// Strings with form ip-address:port
	ipPorts := make(map[string]string)

	nilURL := url.URL{}
	nilRF := ResponseFeatures{}

	// URLnRF just provides caching and some logging on top of followRedirects.
	URLnRF := func(domain string) (url.URL, ResponseFeatures) {
		if _, ok := badDomains[domain]; ok {
			dbg.Debugln("Known bad domain", domain)
			return nilURL, nilRF
		}
		u, ok := nru[domain]
		if ok {
			dbg.Debugln("Cached domain", domain, u.String())
			return u, nrf[domain]
		} else {
			u, ipPort, rf, err := followRedirects(domain)
			if err != nil {
				logErr("can't reach "+domain+" *without* a proxy", err)
				badDomains[domain] = true
				return nilURL, nilRF
			} else {
				dbg.Debugln("Got new url", u.String(), "features", rf)
				nru[domain] = u
				nrf[domain] = rf
				ipPorts[domain] = ipPort
				return u, rf
			}
		}
	}

	for _, target := range targets {
		dbg.Debugln("Trying", target)
		targetURL, features := URLnRF(target)
		if targetURL == nilURL {
			dbg.Debugln("Got null URL for", target)
			// This works for Baidu at least.  Worth trying, should there be
			// more sites for which www. works, the naked domain doesn't, and
			// which won't automatically redirect.
			targetURL, features = URLnRF("www." + target)
			if targetURL == nilURL {
				continue
			}
		}
		for _, front := range fronts {
			dbg.Debugln("Enqueuing pairing:", front, target)
			frontURL, _ := URLnRF(front)
			if frontURL == nilURL {
				// Some CDNs like akamai or Fastly don't like to have their /s
				// requested, but will still domain front.
				frontURL = url.URL{Scheme: "http", Host: front}
			}
			taskChan <- Pairing{
				Target:    target,
				TargetURL: targetURL,
				Front:     front,
				FrontURL:  frontURL,
				FrontIPPort: ipPorts[front],
				Features:  features,
			}
		}
	}
	dbg.Debugln("Enqueued all input")
	close(taskChan)
}

// followRedirects makes a request for the given domain without domain fronting,
// returning a URL that will hopefully not redirect, and the relevant features
// of the response.
func followRedirects(domain string) (u url.URL, ipPort string, rf ResponseFeatures, err error) {
	retriesLeft := 10
	// use http scheme to avoid getting double-TLSed
	u = url.URL{Scheme: "http", Host: domain}
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				dbg.Debugln("** trying to get", addr)
				conn, err = tls.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					domain+":443",
					&tls.Config{InsecureSkipVerify: true},
				)
				ipPort = conn.RemoteAddr().String()
				return conn, err
			},
		},
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			retriesLeft--
			if retriesLeft == 0 {
				return errors.New("Too many redirects")
			}
			// use http scheme to avoid getting double-TLSed
			req.URL.Scheme = "http"
			u = *req.URL
			dbg.Debugln("(re)setting URL to", u.String())
			return nil
		},
	}
	// use http scheme to avoid getting double-TLSed
	req, err := http.NewRequest("GET", "http://"+domain, nil)
	if err != nil {
		logErr("building GET request", err)
		return url.URL{}, ipPort, rf, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return url.URL{}, ipPort, rf, err
	}
	if resp.StatusCode != 200 {
		return url.URL{}, ipPort, rf, errors.New("Non-200 response:" + resp.Status)
	}
	rf, err = responseFeatures(resp)
	if err != nil {
		logErr("getting response features", err)
		return url.URL{}, ipPort, rf, err
	}
	return
}

func responseFeatures(resp *http.Response) (ret ResponseFeatures, err error) {
	body, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		return ret, err
	}
	// In what follows, we take advantage of the fact that all ret fields start
	// at their zero value (in particular, "" for strings).
	if body == nil {
		return ret, nil
	}
	ret.BodyHash = hash(body)
	if headMatches := headRe.FindSubmatch(body); headMatches != nil {
		ret.HeadHash = hash(headMatches[1])
	}
	if titleMatches := titleRe.FindSubmatch(body); titleMatches != nil {
		dbg.Debugln("Title:", string(titleMatches[1]))
		ret.TitleHash = hash(titleMatches[1])
	}
	return
}

func hash(s []byte) string {
	h := sha1.New()
	h.Write(s)
	bs := h.Sum(nil)
	return fmt.Sprintf("%x", bs)
}

func matchReport(expected, actual string) string {
	if actual == expected {
		if actual == "" {
			return "both empty"
		} else {
			return "match"
		}
	} else {
		if actual == "" {
			return "actual empty"
		} else {
			return "differ"
		}
	}
}

func readLines(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(file)
	ret := make([]string, 0)
	for scanner.Scan() {
		ret = append(ret, scanner.Text())
	}
	return ret, scanner.Err()
}

func noRedirect(req *http.Request, via []*http.Request) error {
	return errors.New("Don't redirect!: " + req.URL.String())
}

func logErr(msg string, err error) {
	if err != nil {
		log.Printf("ERROR %s: %v\n", msg, err)
	}
}
<commit_msg>only use ipPort if successfully connected<commit_after>// for each target domain in the given prioritized list:
//     for front candidate in the given prioritized list:
//        check whether we can reach the target domain by attempting domain
//        fronting through the front candidate.  Log successfully proxied pairs
//        to stdout.  Log detected errors to stderr.
//
// run with `-tags debug` to enable debug logging.  See `-help` for other
// options.

// XXX: record where we left off, to enable resuming an interrupted scan.

package main

import (
	"bufio"
	"crypto/sha1"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"sync"
	"time"

	dbg "github.com/getlantern/sitescanner/log"
	"gopkg.in/getlantern/tlsdialer.v1"
)

var headRe *regexp.Regexp = regexp.MustCompile("(?s)(.*)")

var titleRe *regexp.Regexp = regexp.MustCompile("(?s)<title>(.*)</title>")

const reportFmt = "%-24s %-24s %-12s %-12s %-12s %-20s %s\n"

var dialTimeout = 10 * time.Second

var fronts = flag.String("fronts", "fronts.txt", "Path to the file containing front candidates, one domain per line.")

var targets = flag.String("targets", "targets.txt", "Path to the file containing fronting targets, one domain per line.")

var workers = flag.Int("workers", runtime.NumCPU() * 100, "Number of concurrent worker goroutines.")

// ResponseFeatures contains the characteristics we'll use when comparing
// direct vs domain fronted Responses, to estimate the likelihood that we are
// successfully fronting to a target.
//
// We are also checking whether we get a 200 OK response from the site at all,
// but we don't include this in ResponseFeatures becouse we are not reporting
// sites for which we don't get a 200 OK through domain fronting.
type ResponseFeatures struct {
	// Hash of the whole body, or "" if the body is the empty string
	BodyHash string
	// Hash of the response body's  element, or "" if it has none
	HeadHash string
	// Hash of the <title> element, or "" if the response body has none.
	TitleHash string
}

type Pairing struct {
	// The domain we're trying to reach.
	Target string
	// A URL where we can hopefully reach Target without redirects.
	TargetURL url.URL
	// The domain we're trying to domain front through.
	Front string
	// A URL where we can hopefully reach Front without redirects.
	FrontURL url.URL
	// A string of the form ip-address:port
	FrontIPPort string

	Features ResponseFeatures
}

func main() {
	flag.Parse()
	runtime.GOMAXPROCS(runtime.NumCPU())
	start := time.Now()
	fmt.Println()
	fmt.Printf(reportFmt, "TARGET", "FRONT", "BODIES", "HEADS", "TITLES", "FRONT IP:PORT", "REDIRECTED URLS")
	fmt.Printf(reportFmt, "======", "=====", "======", "=====", "======", "=============", "===============")
	fmt.Println()
	main_()
	fmt.Println("**************** COMPLETE ******************")
	fmt.Printf("Scan took %.2fs\n", time.Since(start).Seconds())
}

func main_() {

	tasksChan := make(chan Pairing, *workers)

	workersWg := sync.WaitGroup{}

	// Size of work queue is O(n^2) on the size of the input lists, so we
	// generate it lazily.
	go feedTasks(tasksChan)

	workersWg.Add(*workers)
	for i := 0; i < *workers; i++ {
		go work(tasksChan, &workersWg)
	}

	// Task completion handling:
	//
	// The feedTasks goroutine will close tasksChan when it has processed all
	// input.  Each worker reports workersWg.Done() when it learns that
	// tasksChan is closed.  By that time, the worker has sent all its results
	// to the resultsChan.  Therefore, by the time workersWg.Wait() returns,
	// all `Pairing`s have been checked and all the successful ones have been
	// printed out.
	workersWg.Wait()
}

func work(tasksChan <-chan Pairing, workersWg *sync.WaitGroup) {
	for task := range tasksChan {
		features, err := proxiedResponseFeatures(task)
		if err != nil {
			dbg.Debugln("proxiedResponseFeatures error:", err)
			continue
		}
		bodiesMatch := matchReport(task.Features.BodyHash, features.BodyHash)
		headsMatch := matchReport(task.Features.HeadHash, features.HeadHash)
		titlesMatch := matchReport(task.Features.TitleHash, features.TitleHash)
		if bodiesMatch != "match" && headsMatch != "match" && titlesMatch != "match" {
			dbg.Debugln("Got an OK response fronting", task.Target, "through", task.Front, "but nothing matches.")
			continue
		}
		fmt.Printf(
			reportFmt,
			task.Target,
			task.Front,
			bodiesMatch,
			headsMatch,
			titlesMatch,
			task.FrontIPPort,
			task.TargetURL.String() + ", " + task.FrontURL.String(),
		)
	}
	dbg.Debugln("One worker done")
	workersWg.Done()
}

func proxiedResponseFeatures(p Pairing) (ResponseFeatures, error) {
	dbg.Debugln("canProxy", p)
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				return tlsdialer.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					p.FrontURL.Host+":443",
					// We can't use a domain front that requires a properly
					// populated SNI, so let's make those fail.
					false,
					&tls.Config{InsecureSkipVerify: true},
				)
			},
		},
		CheckRedirect: noRedirect,
	}
	req, err := http.NewRequest("GET", p.TargetURL.String(), nil)
	if err != nil {
		logErr("building GET request", err)
		return ResponseFeatures{}, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return ResponseFeatures{}, err
	}
	return responseFeatures(resp)
}

// feedTasks reads input files, enqueues testing tasks in taskChan, and closes
// taskChan when done.
func feedTasks(taskChan chan<- Pairing) {
	fronts, err := readLines(*fronts)
	if err != nil {
		panic(err)
	}
	targets, err := readLines(*targets)
	if err != nil {
		panic(err)
	}

	// Domains for which we can't get a hash, even without domain fronting.
	badDomains := make(map[string]bool)

	// No-redirect URLs
	nru := make(map[string]url.URL)

	// No-redirect response features
	nrf := make(map[string]ResponseFeatures)

	// Strings with form ip-address:port
	ipPorts := make(map[string]string)

	nilURL := url.URL{}
	nilRF := ResponseFeatures{}

	// URLnRF just provides caching and some logging on top of followRedirects.
	URLnRF := func(domain string) (url.URL, ResponseFeatures) {
		if _, ok := badDomains[domain]; ok {
			dbg.Debugln("Known bad domain", domain)
			return nilURL, nilRF
		}
		u, ok := nru[domain]
		if ok {
			dbg.Debugln("Cached domain", domain, u.String())
			return u, nrf[domain]
		} else {
			u, ipPort, rf, err := followRedirects(domain)
			if err != nil {
				logErr("can't reach "+domain+" *without* a proxy", err)
				badDomains[domain] = true
				return nilURL, nilRF
			} else {
				dbg.Debugln("Got new url", u.String(), "features", rf)
				nru[domain] = u
				nrf[domain] = rf
				ipPorts[domain] = ipPort
				return u, rf
			}
		}
	}

	for _, target := range targets {
		dbg.Debugln("Trying", target)
		targetURL, features := URLnRF(target)
		if targetURL == nilURL {
			dbg.Debugln("Got null URL for", target)
			// This works for Baidu at least.  Worth trying, should there be
			// more sites for which www. works, the naked domain doesn't, and
			// which won't automatically redirect.
			targetURL, features = URLnRF("www." + target)
			if targetURL == nilURL {
				continue
			}
		}
		for _, front := range fronts {
			dbg.Debugln("Enqueuing pairing:", front, target)
			frontURL, _ := URLnRF(front)
			if frontURL == nilURL {
				// Some CDNs like akamai or Fastly don't like to have their /s
				// requested, but will still domain front.
				frontURL = url.URL{Scheme: "http", Host: front}
			}
			taskChan <- Pairing{
				Target:    target,
				TargetURL: targetURL,
				Front:     front,
				FrontURL:  frontURL,
				FrontIPPort: ipPorts[front],
				Features:  features,
			}
		}
	}
	dbg.Debugln("Enqueued all input")
	close(taskChan)
}

// followRedirects makes a request for the given domain without domain fronting,
// returning a URL that will hopefully not redirect, and the relevant features
// of the response.
func followRedirects(domain string) (u url.URL, ipPort string, rf ResponseFeatures, err error) {
	retriesLeft := 10
	// use http scheme to avoid getting double-TLSed
	u = url.URL{Scheme: "http", Host: domain}
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				dbg.Debugln("** trying to get", addr)
				conn, err = tls.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					domain+":443",
					&tls.Config{InsecureSkipVerify: true},
				)
				if err == nil {
					ipPort = conn.RemoteAddr().String()
				}
				return conn, err
			},
		},
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			retriesLeft--
			if retriesLeft == 0 {
				return errors.New("Too many redirects")
			}
			// use http scheme to avoid getting double-TLSed
			req.URL.Scheme = "http"
			u = *req.URL
			dbg.Debugln("(re)setting URL to", u.String())
			return nil
		},
	}
	// use http scheme to avoid getting double-TLSed
	req, err := http.NewRequest("GET", "http://"+domain, nil)
	if err != nil {
		logErr("building GET request", err)
		return url.URL{}, ipPort, rf, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return url.URL{}, ipPort, rf, err
	}
	if resp.StatusCode != 200 {
		return url.URL{}, ipPort, rf, errors.New("Non-200 response:" + resp.Status)
	}
	rf, err = responseFeatures(resp)
	if err != nil {
		logErr("getting response features", err)
		return url.URL{}, ipPort, rf, err
	}
	return
}

func responseFeatures(resp *http.Response) (ret ResponseFeatures, err error) {
	body, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		return ret, err
	}
	// In what follows, we take advantage of the fact that all ret fields start
	// at their zero value (in particular, "" for strings).
	if body == nil {
		return ret, nil
	}
	ret.BodyHash = hash(body)
	if headMatches := headRe.FindSubmatch(body); headMatches != nil {
		ret.HeadHash = hash(headMatches[1])
	}
	if titleMatches := titleRe.FindSubmatch(body); titleMatches != nil {
		dbg.Debugln("Title:", string(titleMatches[1]))
		ret.TitleHash = hash(titleMatches[1])
	}
	return
}

func hash(s []byte) string {
	h := sha1.New()
	h.Write(s)
	bs := h.Sum(nil)
	return fmt.Sprintf("%x", bs)
}

func matchReport(expected, actual string) string {
	if actual == expected {
		if actual == "" {
			return "both empty"
		} else {
			return "match"
		}
	} else {
		if actual == "" {
			return "actual empty"
		} else {
			return "differ"
		}
	}
}

func readLines(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(file)
	ret := make([]string, 0)
	for scanner.Scan() {
		ret = append(ret, scanner.Text())
	}
	return ret, scanner.Err()
}

func noRedirect(req *http.Request, via []*http.Request) error {
	return errors.New("Don't redirect!: " + req.URL.String())
}

func logErr(msg string, err error) {
	if err != nil {
		log.Printf("ERROR %s: %v\n", msg, err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fatih/structs"
	"github.com/gorilla/mux"
	"github.com/malice-plugins/go-plugin-utils/database/elasticsearch"
	"github.com/malice-plugins/go-plugin-utils/utils"
	"github.com/parnurzeal/gorequest"
	"github.com/urfave/cli"
)

var (
	// Version stores the plugin's version
	Version string
	// BuildTime stores the plugin's build time
	BuildTime string

	path string
)

const (
	name     = "avg"
	category = "av"
)

type pluginResults struct {
	ID   string      `json:"id" structs:"id,omitempty"`
	Data ResultsData `json:"avast" structs:"avg"`
}

// AVG json object
type AVG struct {
	Results ResultsData `json:"avg"`
}

// ResultsData json object
type ResultsData struct {
	Infected bool   `json:"infected" structs:"infected"`
	Result   string `json:"result" structs:"result"`
	Engine   string `json:"engine" structs:"engine"`
	Database string `json:"database" structs:"database"`
	Updated  string `json:"updated" structs:"updated"`
	MarkDown string `json:"markdown,omitempty" structs:"markdown,omitempty"`
}

func assert(err error) {
	if err != nil {
		// AVG exits with error status 5 if it finds a virus
		if err.Error() != "exit status 5" {
			log.WithFields(log.Fields{
				"plugin":   name,
				"category": category,
				"path":     path,
			}).Fatal(err)
		}
	}
}

// AvScan performs antivirus scan
func AvScan(timeout int) AVG {

	// Give avgd 10 seconds to finish
	avgdCtx, avgdCancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	defer avgdCancel()
	// AVG needs to have the daemon started first
	_, err := utils.RunCommand(avgdCtx, "/etc/init.d/avgd", "start")
	assert(err)

	var results ResultsData

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	defer cancel()

	output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
	log.WithFields(log.Fields{
		"plugin":   name,
		"category": category,
		"path":     path,
	}).Debug("AVG output: ", output)
	assert(err)

	results, err = ParseAVGOutput(output, nil, path)

	if err != nil {
		// If fails try a second time
		output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
		log.WithFields(log.Fields{
			"plugin":   name,
			"category": category,
			"path":     path,
		}).Debug("AVG output: ", output)
		assert(err)

		results, err = ParseAVGOutput(output, nil, path)
		assert(err)
	}

	return AVG{
		Results: results,
	}
}

// ParseAVGOutput convert avg output into ResultsData struct
func ParseAVGOutput(avgout string, err error, path string) (ResultsData, error) {

	if err != nil {
		return ResultsData{}, err
	}

	log.Debug("AVG Output: ", avgout)

	avg := ResultsData{
		Infected: false,
		Engine:   getAvgVersion(),
	}
	colonSeparated := []string{}

	lines := strings.Split(avgout, "\n")
	// Extract Virus string and extract colon separated lines into an slice
	for _, line := range lines {
		if len(line) != 0 {
			if strings.Contains(line, ":") {
				colonSeparated = append(colonSeparated, line)
			}
			if strings.Contains(line, path) {
				pathVirusString := strings.Split(line, "  ")
				avg.Result = strings.TrimSpace(pathVirusString[1])
			}
		}
	}
	// fmt.Println(lines)

	// Extract AVG Details from scan output
	if len(colonSeparated) != 0 {
		for _, line := range colonSeparated {
			if len(line) != 0 {
				keyvalue := strings.Split(line, ":")
				if len(keyvalue) != 0 {
					switch {
					case strings.Contains(line, "Virus database version"):
						avg.Database = strings.TrimSpace(keyvalue[1])
					case strings.Contains(line, "Virus database release date"):
						date := strings.TrimSpace(strings.TrimPrefix(line, "Virus database release date:"))
						avg.Updated = parseUpdatedDate(date)
					case strings.Contains(line, "Infections found"):
						if strings.Contains(keyvalue[1], "1") {
							avg.Infected = true
						}
					}
				}
			}
		}
	} else {
		log.Error("[ERROR] colonSeparated was empty: ", colonSeparated)
		log.Errorf("[ERROR] AVG output was: \n%s", avgout)
		// fmt.Println("[ERROR] colonSeparated was empty: ", colonSeparated)
		// fmt.Printf("[ERROR] AVG output was: \n%s", avgout)
		return ResultsData{}, errors.New("Unable to parse AVG output")
	}

	return avg, nil
}

// Get Anti-Virus scanner version
func getAvgVersion() string {
	versionOut, err := utils.RunCommand(nil, "/usr/bin/avgscan", "-v")
	assert(err)

	log.Debug("AVG Version: ", versionOut)

	lines := strings.Split(versionOut, "\n")
	for _, line := range lines {
		if len(line) != 0 {
			keyvalue := strings.Split(line, ":")
			if len(keyvalue) != 0 {
				if strings.Contains(keyvalue[0], "Anti-Virus scanner version") {
					return strings.TrimSpace(keyvalue[1])
				}
			}
		}
	}
	return ""
}

func parseUpdatedDate(date string) string {
	layout := "Mon, 02 Jan 2006 15:04:05 +0000"
	t, _ := time.Parse(layout, date)
	return fmt.Sprintf("%d%02d%02d", t.Year(), t.Month(), t.Day())
}

func getUpdatedDate() string {
	if _, err := os.Stat("/opt/malice/UPDATED"); os.IsNotExist(err) {
		return BuildTime
	}
	updated, err := ioutil.ReadFile("/opt/malice/UPDATED")
	assert(err)
	return string(updated)
}

func updateAV(ctx context.Context) error {
	fmt.Println("Updating AVG...")
	// AVG needs to have the daemon started first
	exec.Command("/etc/init.d/avgd", "start").Output()

	output, err := utils.RunCommand(nil, "avgupdate")
	log.WithFields(log.Fields{
		"plugin":   name,
		"category": category,
		"path":     path,
	}).Debug("AVG update: ", output)
	assert(err)

	// Update UPDATED file
	t := time.Now().Format("20060102")
	err := ioutil.WriteFile("/opt/malice/UPDATED", []byte(t), 0644)
	return err
}

func generateMarkDownTable(a AVG) string {
	var tplOut bytes.Buffer

	t := template.Must(template.New("avg").Parse(tpl))

	err := t.Execute(&tplOut, a)
	if err != nil {
		log.Println("executing template:", err)
	}

	return tplOut.String()
}

func printStatus(resp gorequest.Response, body string, errs []error) {
	fmt.Println(body)
}

func webService() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/scan", webAvScan).Methods("POST")
	log.Info("web service listening on port :3993")
	log.Fatal(http.ListenAndServe(":3993", router))
}

func webAvScan(w http.ResponseWriter, r *http.Request) {

	r.ParseMultipartForm(32 << 20)
	file, header, err := r.FormFile("malware")
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, "Please supply a valid file to scan.")
		log.Error(err)
	}
	defer file.Close()

	log.Debug("Uploaded fileName: ", header.Filename)

	tmpfile, err := ioutil.TempFile("/malware", "web_")
	if err != nil {
		log.Fatal(err)
	}
	defer os.Remove(tmpfile.Name()) // clean up

	data, err := ioutil.ReadAll(file)
	assert(err)

	if _, err = tmpfile.Write(data); err != nil {
		log.Fatal(err)
	}
	if err = tmpfile.Close(); err != nil {
		log.Fatal(err)
	}

	// Do AV scan
	path = tmpfile.Name()
	avg := AvScan(120)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(avg); err != nil {
		log.Fatal(err)
	}
}

func main() {

	var elastic string

	cli.AppHelpTemplate = utils.AppHelpTemplate
	app := cli.NewApp()

	app.Name = "avg"
	app.Author = "blacktop"
	app.Email = "https://github.com/blacktop"
	app.Version = Version + ", BuildTime: " + BuildTime
	app.Compiled, _ = time.Parse("20060102", BuildTime)
	app.Usage = "Malice AVG AntiVirus Plugin"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "verbose output",
		},
		cli.StringFlag{
			Name:        "elasitcsearch",
			Value:       "",
			Usage:       "elasitcsearch address for Malice to store results",
			EnvVar:      "MALICE_ELASTICSEARCH",
			Destination: &elastic,
		},
		cli.BoolFlag{
			Name:  "table, t",
			Usage: "output as Markdown table",
		},
		cli.BoolFlag{
			Name:   "callback, c",
			Usage:  "POST results back to Malice webhook",
			EnvVar: "MALICE_ENDPOINT",
		},
		cli.BoolFlag{
			Name:   "proxy, x",
			Usage:  "proxy settings for Malice webhook endpoint",
			EnvVar: "MALICE_PROXY",
		},
		cli.IntFlag{
			Name:   "timeout",
			Value:  120,
			Usage:  "malice plugin timeout (in seconds)",
			EnvVar: "MALICE_TIMEOUT",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "update",
			Usage: "Update virus definitions",
			Action: func(c *cli.Context) error {
				return updateAV(nil)
			},
		},
		{
			Name:  "web",
			Usage: "Create a AVG scan web service",
			Action: func(c *cli.Context) error {
				// ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.Int("timeout"))*time.Second)
				// defer cancel()

				webService()

				return nil
			},
		},
	}
	app.Action = func(c *cli.Context) error {

		var err error

		if c.Bool("verbose") {
			log.SetLevel(log.DebugLevel)
		}

		if c.Args().Present() {
			path, err = filepath.Abs(c.Args().First())
			assert(err)

			if _, err = os.Stat(path); os.IsNotExist(err) {
				assert(err)
			}

			avg := AvScan(c.Int("timeout"))
			avg.Results.MarkDown = generateMarkDownTable(avg)

			// upsert into Database
			elasticsearch.InitElasticSearch(elastic)
			elasticsearch.WritePluginResultsToDatabase(elasticsearch.PluginResults{
				ID:       utils.Getopt("MALICE_SCANID", utils.GetSHA256(path)),
				Name:     name,
				Category: category,
				Data:     structs.Map(avg.Results),
			})

			if c.Bool("table") {
				fmt.Println(avg.Results.MarkDown)
			} else {
				avg.Results.MarkDown = ""
				avgJSON, err := json.Marshal(avg)
				assert(err)
				if c.Bool("callback") {
					request := gorequest.New()
					if c.Bool("proxy") {
						request = gorequest.New().Proxy(os.Getenv("MALICE_PROXY"))
					}
					request.Post(os.Getenv("MALICE_ENDPOINT")).
						Set("X-Malice-ID", utils.Getopt("MALICE_SCANID", utils.GetSHA256(path))).
						Send(string(avgJSON)).
						End(printStatus)

					return nil
				}
				fmt.Println(string(avgJSON))
			}
		} else {
			log.Fatal(fmt.Errorf("Please supply a file to scan with malice/avg"))
		}
		return nil
	}

	err := app.Run(os.Args)
	assert(err)
}
<commit_msg>fix scan.go typo<commit_after>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fatih/structs"
	"github.com/gorilla/mux"
	"github.com/malice-plugins/go-plugin-utils/database/elasticsearch"
	"github.com/malice-plugins/go-plugin-utils/utils"
	"github.com/parnurzeal/gorequest"
	"github.com/urfave/cli"
)

var (
	// Version stores the plugin's version
	Version string
	// BuildTime stores the plugin's build time
	BuildTime string

	path string
)

const (
	name     = "avg"
	category = "av"
)

type pluginResults struct {
	ID   string      `json:"id" structs:"id,omitempty"`
	Data ResultsData `json:"avast" structs:"avg"`
}

// AVG json object
type AVG struct {
	Results ResultsData `json:"avg"`
}

// ResultsData json object
type ResultsData struct {
	Infected bool   `json:"infected" structs:"infected"`
	Result   string `json:"result" structs:"result"`
	Engine   string `json:"engine" structs:"engine"`
	Database string `json:"database" structs:"database"`
	Updated  string `json:"updated" structs:"updated"`
	MarkDown string `json:"markdown,omitempty" structs:"markdown,omitempty"`
}

func assert(err error) {
	if err != nil {
		// AVG exits with error status 5 if it finds a virus
		if err.Error() != "exit status 5" {
			log.WithFields(log.Fields{
				"plugin":   name,
				"category": category,
				"path":     path,
			}).Fatal(err)
		}
	}
}

// AvScan performs antivirus scan
func AvScan(timeout int) AVG {

	// Give avgd 10 seconds to finish
	avgdCtx, avgdCancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	defer avgdCancel()
	// AVG needs to have the daemon started first
	_, err := utils.RunCommand(avgdCtx, "/etc/init.d/avgd", "start")
	assert(err)

	var results ResultsData

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	defer cancel()

	output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
	log.WithFields(log.Fields{
		"plugin":   name,
		"category": category,
		"path":     path,
	}).Debug("AVG output: ", output)
	assert(err)

	results, err = ParseAVGOutput(output, nil, path)

	if err != nil {
		// If fails try a second time
		output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
		log.WithFields(log.Fields{
			"plugin":   name,
			"category": category,
			"path":     path,
		}).Debug("AVG output: ", output)
		assert(err)

		results, err = ParseAVGOutput(output, nil, path)
		assert(err)
	}

	return AVG{
		Results: results,
	}
}

// ParseAVGOutput convert avg output into ResultsData struct
func ParseAVGOutput(avgout string, err error, path string) (ResultsData, error) {

	if err != nil {
		return ResultsData{}, err
	}

	log.Debug("AVG Output: ", avgout)

	avg := ResultsData{
		Infected: false,
		Engine:   getAvgVersion(),
	}
	colonSeparated := []string{}

	lines := strings.Split(avgout, "\n")
	// Extract Virus string and extract colon separated lines into an slice
	for _, line := range lines {
		if len(line) != 0 {
			if strings.Contains(line, ":") {
				colonSeparated = append(colonSeparated, line)
			}
			if strings.Contains(line, path) {
				pathVirusString := strings.Split(line, "  ")
				avg.Result = strings.TrimSpace(pathVirusString[1])
			}
		}
	}
	// fmt.Println(lines)

	// Extract AVG Details from scan output
	if len(colonSeparated) != 0 {
		for _, line := range colonSeparated {
			if len(line) != 0 {
				keyvalue := strings.Split(line, ":")
				if len(keyvalue) != 0 {
					switch {
					case strings.Contains(line, "Virus database version"):
						avg.Database = strings.TrimSpace(keyvalue[1])
					case strings.Contains(line, "Virus database release date"):
						date := strings.TrimSpace(strings.TrimPrefix(line, "Virus database release date:"))
						avg.Updated = parseUpdatedDate(date)
					case strings.Contains(line, "Infections found"):
						if strings.Contains(keyvalue[1], "1") {
							avg.Infected = true
						}
					}
				}
			}
		}
	} else {
		log.Error("[ERROR] colonSeparated was empty: ", colonSeparated)
		log.Errorf("[ERROR] AVG output was: \n%s", avgout)
		// fmt.Println("[ERROR] colonSeparated was empty: ", colonSeparated)
		// fmt.Printf("[ERROR] AVG output was: \n%s", avgout)
		return ResultsData{}, errors.New("Unable to parse AVG output")
	}

	return avg, nil
}

// Get Anti-Virus scanner version
func getAvgVersion() string {
	versionOut, err := utils.RunCommand(nil, "/usr/bin/avgscan", "-v")
	assert(err)

	log.Debug("AVG Version: ", versionOut)

	lines := strings.Split(versionOut, "\n")
	for _, line := range lines {
		if len(line) != 0 {
			keyvalue := strings.Split(line, ":")
			if len(keyvalue) != 0 {
				if strings.Contains(keyvalue[0], "Anti-Virus scanner version") {
					return strings.TrimSpace(keyvalue[1])
				}
			}
		}
	}
	return ""
}

func parseUpdatedDate(date string) string {
	layout := "Mon, 02 Jan 2006 15:04:05 +0000"
	t, _ := time.Parse(layout, date)
	return fmt.Sprintf("%d%02d%02d", t.Year(), t.Month(), t.Day())
}

func getUpdatedDate() string {
	if _, err := os.Stat("/opt/malice/UPDATED"); os.IsNotExist(err) {
		return BuildTime
	}
	updated, err := ioutil.ReadFile("/opt/malice/UPDATED")
	assert(err)
	return string(updated)
}

func updateAV(ctx context.Context) error {
	fmt.Println("Updating AVG...")
	// AVG needs to have the daemon started first
	exec.Command("/etc/init.d/avgd", "start").Output()

	output, err := utils.RunCommand(nil, "avgupdate")
	log.WithFields(log.Fields{
		"plugin":   name,
		"category": category,
		"path":     path,
	}).Debug("AVG update: ", output)
	assert(err)

	// Update UPDATED file
	t := time.Now().Format("20060102")
	err = ioutil.WriteFile("/opt/malice/UPDATED", []byte(t), 0644)
	return err
}

func generateMarkDownTable(a AVG) string {
	var tplOut bytes.Buffer

	t := template.Must(template.New("avg").Parse(tpl))

	err := t.Execute(&tplOut, a)
	if err != nil {
		log.Println("executing template:", err)
	}

	return tplOut.String()
}

func printStatus(resp gorequest.Response, body string, errs []error) {
	fmt.Println(body)
}

func webService() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/scan", webAvScan).Methods("POST")
	log.Info("web service listening on port :3993")
	log.Fatal(http.ListenAndServe(":3993", router))
}

func webAvScan(w http.ResponseWriter, r *http.Request) {

	r.ParseMultipartForm(32 << 20)
	file, header, err := r.FormFile("malware")
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, "Please supply a valid file to scan.")
		log.Error(err)
	}
	defer file.Close()

	log.Debug("Uploaded fileName: ", header.Filename)

	tmpfile, err := ioutil.TempFile("/malware", "web_")
	if err != nil {
		log.Fatal(err)
	}
	defer os.Remove(tmpfile.Name()) // clean up

	data, err := ioutil.ReadAll(file)
	assert(err)

	if _, err = tmpfile.Write(data); err != nil {
		log.Fatal(err)
	}
	if err = tmpfile.Close(); err != nil {
		log.Fatal(err)
	}

	// Do AV scan
	path = tmpfile.Name()
	avg := AvScan(120)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(avg); err != nil {
		log.Fatal(err)
	}
}

func main() {

	var elastic string

	cli.AppHelpTemplate = utils.AppHelpTemplate
	app := cli.NewApp()

	app.Name = "avg"
	app.Author = "blacktop"
	app.Email = "https://github.com/blacktop"
	app.Version = Version + ", BuildTime: " + BuildTime
	app.Compiled, _ = time.Parse("20060102", BuildTime)
	app.Usage = "Malice AVG AntiVirus Plugin"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "verbose output",
		},
		cli.StringFlag{
			Name:        "elasitcsearch",
			Value:       "",
			Usage:       "elasitcsearch address for Malice to store results",
			EnvVar:      "MALICE_ELASTICSEARCH",
			Destination: &elastic,
		},
		cli.BoolFlag{
			Name:  "table, t",
			Usage: "output as Markdown table",
		},
		cli.BoolFlag{
			Name:   "callback, c",
			Usage:  "POST results back to Malice webhook",
			EnvVar: "MALICE_ENDPOINT",
		},
		cli.BoolFlag{
			Name:   "proxy, x",
			Usage:  "proxy settings for Malice webhook endpoint",
			EnvVar: "MALICE_PROXY",
		},
		cli.IntFlag{
			Name:   "timeout",
			Value:  120,
			Usage:  "malice plugin timeout (in seconds)",
			EnvVar: "MALICE_TIMEOUT",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "update",
			Usage: "Update virus definitions",
			Action: func(c *cli.Context) error {
				return updateAV(nil)
			},
		},
		{
			Name:  "web",
			Usage: "Create a AVG scan web service",
			Action: func(c *cli.Context) error {
				// ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.Int("timeout"))*time.Second)
				// defer cancel()

				webService()

				return nil
			},
		},
	}
	app.Action = func(c *cli.Context) error {

		var err error

		if c.Bool("verbose") {
			log.SetLevel(log.DebugLevel)
		}

		if c.Args().Present() {
			path, err = filepath.Abs(c.Args().First())
			assert(err)

			if _, err = os.Stat(path); os.IsNotExist(err) {
				assert(err)
			}

			avg := AvScan(c.Int("timeout"))
			avg.Results.MarkDown = generateMarkDownTable(avg)

			// upsert into Database
			elasticsearch.InitElasticSearch(elastic)
			elasticsearch.WritePluginResultsToDatabase(elasticsearch.PluginResults{
				ID:       utils.Getopt("MALICE_SCANID", utils.GetSHA256(path)),
				Name:     name,
				Category: category,
				Data:     structs.Map(avg.Results),
			})

			if c.Bool("table") {
				fmt.Println(avg.Results.MarkDown)
			} else {
				avg.Results.MarkDown = ""
				avgJSON, err := json.Marshal(avg)
				assert(err)
				if c.Bool("callback") {
					request := gorequest.New()
					if c.Bool("proxy") {
						request = gorequest.New().Proxy(os.Getenv("MALICE_PROXY"))
					}
					request.Post(os.Getenv("MALICE_ENDPOINT")).
						Set("X-Malice-ID", utils.Getopt("MALICE_SCANID", utils.GetSHA256(path))).
						Send(string(avgJSON)).
						End(printStatus)

					return nil
				}
				fmt.Println(string(avgJSON))
			}
		} else {
			log.Fatal(fmt.Errorf("Please supply a file to scan with malice/avg"))
		}
		return nil
	}

	err := app.Run(os.Args)
	assert(err)
}
<|endoftext|>