<commit_before>package utils

import (
	"fmt"
	"io/ioutil"
	//"log"
	"net/http"
	"net/url"
	"regexp"
	"strings"
)

// JobqueueHTTPClient provides some utility methods to communicate with the Task Manager runner via HTTP
type JobqueueHTTPClient struct {
	URL url.URL
}

func (client *JobqueueHTTPClient) setDefaultURL() {
	u, _ := url.Parse("http://localhost:8000/")
	client.URL = *u
}

//TODO: support JSON content type

// Open an HTTP connection to control the task manager runner
func (client *JobqueueHTTPClient) Open(url string) (string, error) {
	re1, err := regexp.Compile(`([hftps]+://)?([\w\.]+)?(:\d+)?(/[^\?]+)?(?:\?(.*))?`)
	result := re1.FindAllStringSubmatch(url, -1)[0]
	if nil != err {
		return "", err
	}
	client.setDefaultURL()

	if "" != result[1] {
		client.URL.Scheme = result[1]
	}
	if "" != result[2] {
		client.URL.Host = result[2] + result[3]
	} else if "" != result[3] {
		client.URL.Host = "localhost" + result[3]
	}
	if "" != result[4] {
		client.URL.Path = result[4]
	}
	if "" != result[5] {
		client.URL.RawQuery = result[5]
	}
	return "Connected to " + client.URL.String(), err
}

// get the full URL from the path
func (client *JobqueueHTTPClient) getAddress(path string) string {
	u := client.URL
	u.Path = path
	return u.String()
}

// List the task managers listening at this address
func (client *JobqueueHTTPClient) List() (string, error) {
	resp, err := http.Get(client.getAddress("/list"))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// ListWorkers gets the status of each worker process for a given task
func (client *JobqueueHTTPClient) ListWorkers(name string) (string, error) {
	//log.Println("Listing workers for task", name)
	resp, err := http.Get(client.getAddress("/tasks/" + name + "/list"))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// ListAsList returns a list of task names (as a string slice)
func (client *JobqueueHTTPClient) ListAsList() []string {
	res, err := client.List()
	if err != nil {
		return []string{}
	}
	return strings.Split(res, "\n")
}

// Set an option on a certain task
func (client *JobqueueHTTPClient) Set(name string, param string, value string) (string, error) {
	//log.Println("Stopping task", name)
	path := fmt.Sprintf("/tasks/%s/set/%s/%s", name, param, value)
	resp, err := http.PostForm(client.getAddress(path), nil)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// Start a stopped task
func (client *JobqueueHTTPClient) Start(name string) (string, error) {
	//log.Println("Starting task", name)
	resp, err := http.PostForm(client.getAddress("/tasks/"+name+"/start"), nil)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// Stop a running task, or all of them
func (client *JobqueueHTTPClient) Stop(name string) (string, error) {
	//log.Println("Stopping task", name)
	//resp, err := http.PostForm(client.getAddress("/tasks/"+name+"/stop"), nil)
	req, err := http.NewRequest("DELETE", client.getAddress("/tasks/"+name), nil)
	if err != nil {
		return "", err
	}
	httpClient := &http.Client{}
	resp, err := httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// Status gets some information about the status of a task (or all of them)
func (client *JobqueueHTTPClient) Status(name string) (string, error) {
	//log.Println("Status")
	resp, err := http.Get(client.getAddress("/tasks/" + name))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}
<commit_msg>Fix: Don't capture :// in the httpclient regex so that specifying a protocol works<commit_after>package utils

import (
	"fmt"
	"io/ioutil"
	//"log"
	"net/http"
	"net/url"
	"regexp"
	"strings"
)

// JobqueueHTTPClient provides some utility methods to communicate with the Task Manager runner via HTTP
type JobqueueHTTPClient struct {
	URL url.URL
}

func (client *JobqueueHTTPClient) setDefaultURL() {
	u, _ := url.Parse("http://localhost:8000/")
	client.URL = *u
}

//TODO: support JSON content type

// Open an HTTP connection to control the task manager runner
func (client *JobqueueHTTPClient) Open(url string) (string, error) {
	re1, err := regexp.Compile(`([hftps]+)?://([\w\.]+)?(:\d+)?(/[^\?]+)?(?:\?(.*))?`)
	result := re1.FindAllStringSubmatch(url, -1)[0]
	if nil != err {
		return "", err
	}
	client.setDefaultURL()

	if "" != result[1] {
		client.URL.Scheme = result[1]
	}
	if "" != result[2] {
		client.URL.Host = result[2] + result[3]
	} else if "" != result[3] {
		client.URL.Host = "localhost" + result[3]
	}
	if "" != result[4] {
		client.URL.Path = result[4]
	}
	if "" != result[5] {
		client.URL.RawQuery = result[5]
	}

	return "Connected to " + client.URL.String(), err
}

// get the full URL from the path
func (client *JobqueueHTTPClient) getAddress(path string) string {
	u := client.URL
	u.Path = path
	return u.String()
}

// List the task managers listening at this address
func (client *JobqueueHTTPClient) List() (string, error) {
	resp, err := http.Get(client.getAddress("/list"))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// ListWorkers gets the status of each worker process for a given task
func (client *JobqueueHTTPClient) ListWorkers(name string) (string, error) {
	//log.Println("Listing workers for task", name)
	resp, err := http.Get(client.getAddress("/tasks/" + name + "/list"))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// ListAsList returns a list of task names (as a string slice)
func (client *JobqueueHTTPClient) ListAsList() []string {
	res, err := client.List()
	if err != nil {
		return []string{}
	}
	return strings.Split(res, "\n")
}

// Set an option on a certain task
func (client *JobqueueHTTPClient) Set(name string, param string, value string) (string, error) {
	//log.Println("Stopping task", name)
	path := fmt.Sprintf("/tasks/%s/set/%s/%s", name, param, value)
	resp, err := http.PostForm(client.getAddress(path), nil)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// Start a stopped task
func (client *JobqueueHTTPClient) Start(name string) (string, error) {
	//log.Println("Starting task", name)
	resp, err := http.PostForm(client.getAddress("/tasks/"+name+"/start"), nil)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// Stop a running task, or all of them
func (client *JobqueueHTTPClient) Stop(name string) (string, error) {
	//log.Println("Stopping task", name)
	//resp, err := http.PostForm(client.getAddress("/tasks/"+name+"/stop"), nil)
	req, err := http.NewRequest("DELETE", client.getAddress("/tasks/"+name), nil)
	if err != nil {
		return "", err
	}
	httpClient := &http.Client{}
	resp, err := httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}

// Status gets some information about the status of a task (or all of them)
func (client *JobqueueHTTPClient) Status(name string) (string, error) {
	//log.Println("Status")
	resp, err := http.Get(client.getAddress("/tasks/" + name))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	bytes, err := ioutil.ReadAll(resp.Body)
	return string(bytes), err
}
<|endoftext|>


<commit_before>package itchio

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"reflect"
	"strings"
)

type Client struct {
	Key        string
	HTTPClient *http.Client
	BaseURL    string
}

type Response struct {
	Errors []string
}

type User struct {
	ID       int64
	Username string
	CoverUrl string `json:"cover_url"`
}

type Game struct {
	ID  int64
	Url string
}

type Upload struct {
	ID       int64
	Filename string
	Size     int64

	OSX     bool `json:"p_osx"`
	Linux   bool `json:"p_linux"`
	Windows bool `json:"p_windows"`
	Android bool `json:"p_android"`
}

func ClientWithKey(key string) *Client {
	return &Client{
		Key:        key,
		HTTPClient: &http.Client{},
		BaseURL:    "https://itch.io/api/1",
	}
}

type MyGamesResponse struct {
	Response

	Games []Game
}

func (c *Client) MyGames() (r MyGamesResponse, err error) {
	path := c.MakePath("my-games")

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type GameUploadsResponse struct {
	Response

	Uploads []Upload `json:"uploads"`
}

func (c *Client) GameUploads(gameID int64) (r GameUploadsResponse, err error) {
	path := c.MakePath("game/%d/uploads", gameID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type UploadDownloadResponse struct {
	Response

	Url string
}

func (c *Client) UploadDownload(uploadID int64) (r UploadDownloadResponse, err error) {
	path := c.MakePath("upload/%d/download", uploadID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return r, err
}

type NewBuildResponse struct {
	Response

	Build struct {
		ID          int64 `json:"id"`
		UploadID    int64 `json:"upload_id"`
		ParentBuild struct {
			ID int64 `json:"id"`
		} `json:"parent_build"`
	}
}

func (c *Client) CreateBuild(target string, channel string) (r NewBuildResponse, err error) {
	path := c.MakePath("wharf/builds")

	form := url.Values{}
	form.Add("target", target)
	form.Add("channel", channel)

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type BuildFileType string

const (
	BuildFileType_PATCH     BuildFileType = "patch"
	BuildFileType_ARCHIVE                 = "archive"
	BuildFileType_SIGNATURE               = "signature"
)

type BuildFile struct {
	ID   int64
	Type BuildFileType
	Size int64
}

type ListBuildFilesResponse struct {
	Response

	Files []BuildFile
}

func (c *Client) ListBuildFiles(buildID int64) (r ListBuildFilesResponse, err error) {
	path := c.MakePath("wharf/builds/%d/files")

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type NewBuildFileResponse struct {
	Response

	File struct {
		ID           int64
		UploadURL    string            `json:"upload_url"`
		UploadParams map[string]string `json:"upload_params"`
	}
}

func (c *Client) CreateBuildFile(buildID int64, fileType BuildFileType) (r NewBuildFileResponse, err error) {
	path := c.MakePath("wharf/builds/%d/files", buildID)

	form := url.Values{}
	form.Add("type", string(fileType))

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type FinalizeBuildFileResponse struct {
	Response
}

func (c *Client) FinalizeBuildFile(buildID int64, fileID int64, size int64) (r FinalizeBuildFileResponse, err error) {
	path := c.MakePath("wharf/builds/%d/files/%d", buildID, fileID)

	form := url.Values{}
	form.Add("size", fmt.Sprintf("%d", size))

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

func (c *Client) DownloadBuildFile(buildID int64, fileID int64) (reader io.ReadCloser, err error) {
	path := c.MakePath("wharf/builds/%d/files/%d/download", buildID, fileID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	if resp.StatusCode/100 != 3 {
		err = fmt.Errorf("expected redirection")
		return
	}

	dlPath := resp.Header.Get("Location")

	req, err := http.NewRequest("GET", dlPath, nil)
	if err != nil {
		return
	}

	reader = req.Body
	return
}

type BuildEventType string

const (
	BuildEvent_JOB_STARTED   BuildEventType = "job_started"
	BuildEvent_JOB_FAILED                   = "job_failed"
	BuildEvent_JOB_COMPLETED                = "job_completed"
)

type BuildEventData map[string]interface{}

type NewBuildEventResponse struct {
	Response
}

func (c *Client) CreateBuildEvent(buildID int64, eventType BuildEventType, message string, data BuildEventData) (r NewBuildEventResponse, err error) {
	path := c.MakePath("wharf/builds/%d/events", buildID)

	form := url.Values{}
	form.Add("type", string(eventType))
	form.Add("message", message)

	jsonData, err := json.Marshal(data)
	if err != nil {
		return
	}
	form.Add("data", string(jsonData))

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type BuildEvent struct {
	Type    BuildEventType
	Message string
	Data    BuildEventData
}

type ListBuildEventsResponse struct {
	Response

	Events []BuildEvent
}

func (c *Client) ListBuildEvents(buildID int64) (r ListBuildEventsResponse, err error) {
	path := c.MakePath("wharf/builds/%d/events", buildID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

// Helpers

func (c *Client) Do(req *http.Request) (*http.Response, error) {
	if strings.HasPrefix(c.Key, "jwt:") {
		req.Header.Add("Authorization", strings.Split(c.Key, ":")[1])
	}
	return http.DefaultClient.Do(req)
}

func (c *Client) MakePath(format string, a ...interface{}) string {
	base := strings.Trim(c.BaseURL, "/")
	subPath := strings.Trim(fmt.Sprintf(format, a...), "/")

	var key string
	if strings.HasPrefix(c.Key, "jwt:") {
		key = "jwt"
	} else {
		key = c.Key
	}
	return fmt.Sprintf("%s/%s/%s", base, key, subPath)
}

func ParseAPIResponse(dst interface{}, res *http.Response) error {
	bodyReader := res.Body
	defer bodyReader.Close()

	if res.StatusCode/100 != 2 {
		return fmt.Errorf("Server returned %s for %s", res.Status, res.Request.URL.String())
	}

	err := json.NewDecoder(bodyReader).Decode(dst)
	if err != nil {
		return err
	}

	errs := reflect.Indirect(reflect.ValueOf(dst)).FieldByName("Errors")
	if errs.Len() > 0 {
		// TODO: handle other errors too
		return errors.New(errs.Index(0).String())
	}

	return nil
}
<commit_msg>:bug: Fix ListBuildFiles endpoint<commit_after>package itchio

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"reflect"
	"strings"
)

type Client struct {
	Key        string
	HTTPClient *http.Client
	BaseURL    string
}

type Response struct {
	Errors []string
}

type User struct {
	ID       int64
	Username string
	CoverUrl string `json:"cover_url"`
}

type Game struct {
	ID  int64
	Url string
}

type Upload struct {
	ID       int64
	Filename string
	Size     int64

	OSX     bool `json:"p_osx"`
	Linux   bool `json:"p_linux"`
	Windows bool `json:"p_windows"`
	Android bool `json:"p_android"`
}

func ClientWithKey(key string) *Client {
	return &Client{
		Key:        key,
		HTTPClient: &http.Client{},
		BaseURL:    "https://itch.io/api/1",
	}
}

type MyGamesResponse struct {
	Response

	Games []Game
}

func (c *Client) MyGames() (r MyGamesResponse, err error) {
	path := c.MakePath("my-games")

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type GameUploadsResponse struct {
	Response

	Uploads []Upload `json:"uploads"`
}

func (c *Client) GameUploads(gameID int64) (r GameUploadsResponse, err error) {
	path := c.MakePath("game/%d/uploads", gameID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type UploadDownloadResponse struct {
	Response

	Url string
}

func (c *Client) UploadDownload(uploadID int64) (r UploadDownloadResponse, err error) {
	path := c.MakePath("upload/%d/download", uploadID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return r, err
}

type NewBuildResponse struct {
	Response

	Build struct {
		ID          int64 `json:"id"`
		UploadID    int64 `json:"upload_id"`
		ParentBuild struct {
			ID int64 `json:"id"`
		} `json:"parent_build"`
	}
}

func (c *Client) CreateBuild(target string, channel string) (r NewBuildResponse, err error) {
	path := c.MakePath("wharf/builds")

	form := url.Values{}
	form.Add("target", target)
	form.Add("channel", channel)

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type BuildFileType string

const (
	BuildFileType_PATCH     BuildFileType = "patch"
	BuildFileType_ARCHIVE                 = "archive"
	BuildFileType_SIGNATURE               = "signature"
)

type BuildFile struct {
	ID   int64
	Type BuildFileType
	Size int64
}

type ListBuildFilesResponse struct {
	Response

	Files []BuildFile
}

func (c *Client) ListBuildFiles(buildID int64) (r ListBuildFilesResponse, err error) {
	path := c.MakePath("wharf/builds/%d/files", buildID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type NewBuildFileResponse struct {
	Response

	File struct {
		ID           int64
		UploadURL    string            `json:"upload_url"`
		UploadParams map[string]string `json:"upload_params"`
	}
}

func (c *Client) CreateBuildFile(buildID int64, fileType BuildFileType) (r NewBuildFileResponse, err error) {
	path := c.MakePath("wharf/builds/%d/files", buildID)

	form := url.Values{}
	form.Add("type", string(fileType))

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type FinalizeBuildFileResponse struct {
	Response
}

func (c *Client) FinalizeBuildFile(buildID int64, fileID int64, size int64) (r FinalizeBuildFileResponse, err error) {
	path := c.MakePath("wharf/builds/%d/files/%d", buildID, fileID)

	form := url.Values{}
	form.Add("size", fmt.Sprintf("%d", size))

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

func (c *Client) DownloadBuildFile(buildID int64, fileID int64) (reader io.ReadCloser, err error) {
	path := c.MakePath("wharf/builds/%d/files/%d/download", buildID, fileID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	if resp.StatusCode/100 != 3 {
		err = fmt.Errorf("expected redirection")
		return
	}

	dlPath := resp.Header.Get("Location")

	req, err := http.NewRequest("GET", dlPath, nil)
	if err != nil {
		return
	}

	reader = req.Body
	return
}

type BuildEventType string

const (
	BuildEvent_JOB_STARTED   BuildEventType = "job_started"
	BuildEvent_JOB_FAILED                   = "job_failed"
	BuildEvent_JOB_COMPLETED                = "job_completed"
)

type BuildEventData map[string]interface{}

type NewBuildEventResponse struct {
	Response
}

func (c *Client) CreateBuildEvent(buildID int64, eventType BuildEventType, message string, data BuildEventData) (r NewBuildEventResponse, err error) {
	path := c.MakePath("wharf/builds/%d/events", buildID)

	form := url.Values{}
	form.Add("type", string(eventType))
	form.Add("message", message)

	jsonData, err := json.Marshal(data)
	if err != nil {
		return
	}
	form.Add("data", string(jsonData))

	resp, err := http.PostForm(path, form)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

type BuildEvent struct {
	Type    BuildEventType
	Message string
	Data    BuildEventData
}

type ListBuildEventsResponse struct {
	Response

	Events []BuildEvent
}

func (c *Client) ListBuildEvents(buildID int64) (r ListBuildEventsResponse, err error) {
	path := c.MakePath("wharf/builds/%d/events", buildID)

	resp, err := http.Get(path)
	if err != nil {
		return
	}

	err = ParseAPIResponse(&r, resp)
	return
}

// Helpers

func (c *Client) Do(req *http.Request) (*http.Response, error) {
	if strings.HasPrefix(c.Key, "jwt:") {
		req.Header.Add("Authorization", strings.Split(c.Key, ":")[1])
	}
	return http.DefaultClient.Do(req)
}

func (c *Client) MakePath(format string, a ...interface{}) string {
	base := strings.Trim(c.BaseURL, "/")
	subPath := strings.Trim(fmt.Sprintf(format, a...), "/")

	var key string
	if strings.HasPrefix(c.Key, "jwt:") {
		key = "jwt"
	} else {
		key = c.Key
	}
	return fmt.Sprintf("%s/%s/%s", base, key, subPath)
}

func ParseAPIResponse(dst interface{}, res *http.Response) error {
	bodyReader := res.Body
	defer bodyReader.Close()

	if res.StatusCode/100 != 2 {
		return fmt.Errorf("Server returned %s for %s", res.Status, res.Request.URL.String())
	}

	err := json.NewDecoder(bodyReader).Decode(dst)
	if err != nil {
		return err
	}

	errs := reflect.Indirect(reflect.ValueOf(dst)).FieldByName("Errors")
	if errs.Len() > 0 {
		// TODO: handle other errors too
		return errors.New(errs.Index(0).String())
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 e2e

import (
	"log"
	"strings"
	"testing"
	"time"

	"github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/apis/sparkoperator.k8s.io/v1alpha1"
	appFramework "github.com/GoogleCloudPlatform/spark-on-k8s-operator/test/e2e/framework"
	"github.com/stretchr/testify/assert"

	"k8s.io/api/core/v1"
)

func getJobStatus(t *testing.T) v1alpha1.ApplicationStateType {
	app, err := appFramework.GetSparkApplication(framework.SparkApplicationClient, "default", "spark-pi")
	assert.Equal(t, nil, err)
	return app.Status.AppState.State
}

func TestSubmitSparkPiYaml(t *testing.T) {
	t.Parallel()

	// Wait for test job to finish. Time out after 90 seconds.
	timeout := 90

	sa, err := appFramework.MakeSparkApplicationFromYaml("../../examples/spark-pi.yaml")
	assert.Equal(t, nil, err)
	err = appFramework.CreateSparkApplication(framework.SparkApplicationClient, "default", sa)
	assert.Equal(t, nil, err)

	status := getJobStatus(t)

	timePassed := 0
	// Update job status every 5 seconds until job is done or timeout threshold is reached.
	for status != "COMPLETED" && timePassed <= timeout {
		log.Print("Waiting for the Spark job to finish...")
		time.Sleep(5 * time.Second)
		timePassed += 5
		status = getJobStatus(t)
	}
	if timePassed > timeout {
		log.Fatalf("Time out waiting for Spark job to finish!")
	}

	app, _ := appFramework.GetSparkApplication(framework.SparkApplicationClient, "default", "spark-pi")
	podName := app.Status.DriverInfo.PodName
	rawLogs, err := framework.KubeClient.CoreV1().Pods("default").GetLogs(podName, &v1.PodLogOptions{}).Do().Raw()
	assert.Equal(t, nil, err)
	assert.NotEqual(t, -1, strings.Index(string(rawLogs), "Pi is roughly 3"))

	err = appFramework.DeleteSparkApplication(framework.SparkApplicationClient, "default", "spark-pi")
	assert.Equal(t, nil, err)
}
<commit_msg>Use wait.Poll<commit_after>/*
Copyright 2018 Google LLC

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 e2e

import (
	"k8s.io/apimachinery/pkg/util/wait"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"

	"k8s.io/api/core/v1"

	"github.com/GoogleCloudPlatform/spark-on-k8s-operator/pkg/apis/sparkoperator.k8s.io/v1alpha1"
	appFramework "github.com/GoogleCloudPlatform/spark-on-k8s-operator/test/e2e/framework"
)

func getJobStatus(t *testing.T) v1alpha1.ApplicationStateType {
	app, err := appFramework.GetSparkApplication(framework.SparkApplicationClient, "default", "spark-pi")
	assert.Equal(t, nil, err)
	return app.Status.AppState.State
}

func TestSubmitSparkPiYaml(t *testing.T) {
	t.Parallel()

	// Wait for test job to finish. Time out after 90 seconds.
	timeout := 100 * time.Second
	interval := 5 * time.Second

	sa, err := appFramework.MakeSparkApplicationFromYaml("../../examples/spark-pi.yaml")
	assert.Equal(t, nil, err)
	err = appFramework.CreateSparkApplication(framework.SparkApplicationClient, "default", sa)
	assert.Equal(t, nil, err)

	status := getJobStatus(t)

	wait.Poll(interval, timeout, func() (done bool, err error) {
		if status == "COMPLETED" {
			return true, nil
		} else {
			status = getJobStatus(t)
			return false, nil
		}
	})

	app, _ := appFramework.GetSparkApplication(framework.SparkApplicationClient, "default", "spark-pi")
	podName := app.Status.DriverInfo.PodName
	rawLogs, err := framework.KubeClient.CoreV1().Pods("default").GetLogs(podName, &v1.PodLogOptions{}).Do().Raw()
	assert.Equal(t, nil, err)
	assert.NotEqual(t, -1, strings.Index(string(rawLogs), "Pi is roughly 3"))

	err = appFramework.DeleteSparkApplication(framework.SparkApplicationClient, "default", "spark-pi")
	assert.Equal(t, nil, err)
}
<|endoftext|>


<commit_before>package twitter

import (
  "fmt"
  "github.com/garyburd/redigo/redis"
  "github.com/gin-gonic/gin"
  "github.com/kurrik/oauth1a"
  "github.com/kurrik/twittergo"
  "net/http"
  "net/url"
  "os"
  "time"
)

type Twitter struct{
}

var (
  TWITTER_COSUMER_KEY     = os.Getenv("TWITTER_CONSUMER_KEY")
  TWITTER_CONSUMER_SECRET = os.Getenv("TWITTER_CONSUMER_SECRET")
  REDISTOGO, _ = url.Parse(os.Getenv("REDIS_URL"))
   conn, _ = redis.Dial("tcp", REDISTOGO.Host)
)

func (tw Twitter) LoadCredentials() (client *twittergo.Client, err error) {
  config := &oauth1a.ClientConfig{
    ConsumerKey:    TWITTER_COSUMER_KEY,
    ConsumerSecret: TWITTER_CONSUMER_SECRET,
  }
  user := oauth1a.NewAuthorizedConfig("", "")
  client = twittergo.NewClient(config, user)
  return
}

func (tw Twitter) Memoize(resp map[string]string, tweetId uint64,
    key string, value string) map[string]string {

    tweetRedisKey	:= fmt.Sprintf("tweet:%v", tweetId)
    cache, err := redis.String(conn.Do("HGET", tweetRedisKey, key))

    if err != nil && cache == "" {
      conn.Do("HSET", tweetRedisKey, key, value)
      cache = value
    }

  resp[key] = cache
  return resp
}

func (tw Twitter) TweetsFromResults(
    c *gin.Context,
    results *twittergo.SearchResults,
    t chan map[string]string) {

  tweets := []uint64{}

  for _, tweet := range results.Statuses() {

    user 		:= tweet.User()
    entities 	:= tweet["entities"].(map[string]interface{})
    urls 		:= entities["urls"].([]interface{})

    if len(urls) > 0 {
      source 	:= urls[0].(map[string]interface{})
      resp 	:= make(map[string]string)

      //url	:= media[0].(map[string]interface{})
      tweetId 		:= tweet.Id()
      tweetRedisKey 	:= fmt.Sprintf("tweet:%v", tweetId)
      reply, _ 	:= redis.Values(conn.Do("KEYS", tweetRedisKey))
      //fmt.Printf("values: %v", reply)
      if len(reply) == 0 {
        tweets = append(tweets, tweet.Id())
      }

      media, ok 	:= entities["media"]
      //fmt.Printf("media: %v", ok)
      if ok {
        content := media.([]interface{})
        first := content[0].(map[string]interface{})
        if first["type"] == "photo" {
          resp = tw.Memoize(resp, tweetId, "image_source",
            fmt.Sprintf("%v", first["media_url"]))

          sizes := first["sizes"].(map[string]interface{})
          small := sizes["small"].(map[string]interface{})
          h := small["h"].(float64)
          w := small["w"].(float64)
          resp = tw.Memoize(resp, tweetId, "width",
            fmt.Sprintf("%d", w))
          resp = tw.Memoize(resp, tweetId, "height",
            fmt.Sprintf("%d", h))
        }
      }else {
        resp = tw.Memoize(resp, tweetId, "image_source",
          fmt.Sprintf("%v", source["expanded_url"]))
      }
      resp = tw.Memoize(resp, tweetId, "tweet_id",
        fmt.Sprintf("%v", tweetId))
      resp = tw.Memoize(resp, tweetId, "tweet",
        fmt.Sprintf("%v", tweet.Text()))
      resp = tw.Memoize(resp, tweetId, "user",
        fmt.Sprintf("%v (@%v)", user.Name(), user.ScreenName()))
      resp = tw.Memoize(resp, tweetId, "created",
        fmt.Sprintf("%v", tweet.CreatedAt().Format(time.RFC1123)))

      resp["source"] = "search"
      fmt.Println("+ return parsed tweet")
      t <- resp
    }
  }
  for i:=len(tweets)-1; i >= 0; i-- {
    conn.Do("LPUSH", "tweets", tweets[i])
  }
}

func (tw Twitter) RetrieveSinceId() (string, bool){
  reply, _ := redis.String( conn.Do("LINDEX", "tweets", 0) )
  if reply == "" {
    return reply, false
  } else {
    return reply, true
  }
}

func (tw Twitter) TwitterImages() url.Values {
  query := url.Values{}
  query.Set("q", "#beautifulsky filter:images")
  query.Set("result_type", "mixed")
  query.Set("count", "1000")

  if id, present := tw.RetrieveSinceId(); present {
    fmt.Println("+ sinceId(%v), using query %v", id, query)
    query.Set("since_id", id)
  }else{
    fmt.Println("+ no sinceId, using query %v", query)
  }
  return query
}

func (tw Twitter) Instagram() url.Values {
  query := url.Values{}
  query.Set("q", "#beautifulsky")
  query.Set("result_type", "mixed")
  query.Set("count", "1000")

  if id, present := tw.RetrieveSinceId(); present {
    fmt.Println("+ sinceId(%v), using query %v", id, query)
    query.Set("since_id", id)
  }else{
    fmt.Println("+ no sinceId, using query %v", query)
  }
  return query
}

func (tw Twitter) SearchTweets(
  query url.Values,
  k chan *twittergo.SearchResults,
  searchDone chan bool){

  var (
    err     error
    results *twittergo.SearchResults
    req     *http.Request
    resp    *twittergo.APIResponse
    client *twittergo.Client
  )
  client, _ = tw.LoadCredentials()

  url := fmt.Sprintf("/1.1/search/tweets.json?%v", query.Encode())

  req, err = http.NewRequest("GET", url, nil)
  if err != nil {
    fmt.Printf("Could not parse request: %v\n", err)
  }

  fmt.Printf("+ Getting tweets: %v\n", req)
  resp, err = client.SendRequest(req)
  if err != nil {
    fmt.Printf("Could not send request: %v\n", err)
  }

  results = &twittergo.SearchResults{}
  err = resp.Parse(results)
  if err != nil {
    fmt.Printf("Problem parsing response: %v\n", err)
  }

  k <- results
  searchDone <- true
}

func (tw Twitter) TweetsFromCache(t chan map[string]string) {
  l, _ := redis.Int(conn.Do("LLEN", "tweets"))
  for i:=0; i< l && i< 50; i++ {
    tweetId, _ := redis.String(conn.Do("LINDEX", "tweets", i))
    var resp map[string]string = make(map[string]string)
    for _, key := range []string{"tweet", "image_source", "image_url","user", "created", "tweet_id", "width", "height"} {
      reply,_ := redis.String( conn.Do("HGET", "tweet:"+tweetId, key) )
      resp[key] = reply
    }
    fmt.Printf("+ return cache #%v\n", i+1)
    t <- resp
  }
}

func (tw Twitter) PrintRateLimit(r chan *twittergo.APIResponse) {
  fmt.Printf("start PrintRateLimit \n")
  resp := <-r
  if resp.HasRateLimit() {
    fmt.Printf("Rate limit:           %v\n", resp.RateLimit())
    fmt.Printf("Rate limit remaining: %v\n", resp.RateLimitRemaining())
    fmt.Printf("Rate limit reset:     %v\n", resp.RateLimitReset())
  } else {
    fmt.Printf("Could not parse rate limit from response.\n")
  }
  fmt.Printf("end PrintRateLimit \n")
}
<commit_msg>use redisurl<commit_after>package twitter

import (
  "fmt"
  "github.com/garyburd/redigo/redis"
  "github.com/gin-gonic/gin"
  "github.com/kurrik/oauth1a"
  "github.com/kurrik/twittergo"
  "github.com/soveran/redisurl"
  "net/http"
  "net/url"
  "os"
  "time"
)

type Twitter struct{
}

var (
  TWITTER_COSUMER_KEY     = os.Getenv("TWITTER_CONSUMER_KEY")
  TWITTER_CONSUMER_SECRET = os.Getenv("TWITTER_CONSUMER_SECRET")
  conn, _ = redisurl.Connect()
)

func (tw Twitter) LoadCredentials() (client *twittergo.Client, err error) {
  config := &oauth1a.ClientConfig{
    ConsumerKey:    TWITTER_COSUMER_KEY,
    ConsumerSecret: TWITTER_CONSUMER_SECRET,
  }
  user := oauth1a.NewAuthorizedConfig("", "")
  client = twittergo.NewClient(config, user)
  return
}

func (tw Twitter) Memoize(resp map[string]string, tweetId uint64,
    key string, value string) map[string]string {

    tweetRedisKey	:= fmt.Sprintf("tweet:%v", tweetId)
    cache, err := redis.String(conn.Do("HGET", tweetRedisKey, key))

    if err != nil && cache == "" {
      conn.Do("HSET", tweetRedisKey, key, value)
      cache = value
    }

  resp[key] = cache
  return resp
}

func (tw Twitter) TweetsFromResults(
    c *gin.Context,
    results *twittergo.SearchResults,
    t chan map[string]string) {

  tweets := []uint64{}

  for _, tweet := range results.Statuses() {

    user 		:= tweet.User()
    entities 	:= tweet["entities"].(map[string]interface{})
    urls 		:= entities["urls"].([]interface{})

    if len(urls) > 0 {
      source 	:= urls[0].(map[string]interface{})
      resp 	:= make(map[string]string)

      //url	:= media[0].(map[string]interface{})
      tweetId 		:= tweet.Id()
      tweetRedisKey 	:= fmt.Sprintf("tweet:%v", tweetId)
      reply, _ 	:= redis.Values(conn.Do("KEYS", tweetRedisKey))
      //fmt.Printf("values: %v", reply)
      if len(reply) == 0 {
        tweets = append(tweets, tweet.Id())
      }

      media, ok 	:= entities["media"]
      //fmt.Printf("media: %v", ok)
      if ok {
        content := media.([]interface{})
        first := content[0].(map[string]interface{})
        if first["type"] == "photo" {
          resp = tw.Memoize(resp, tweetId, "image_source",
            fmt.Sprintf("%v", first["media_url"]))

          sizes := first["sizes"].(map[string]interface{})
          small := sizes["small"].(map[string]interface{})
          h := small["h"].(float64)
          w := small["w"].(float64)
          resp = tw.Memoize(resp, tweetId, "width",
            fmt.Sprintf("%d", w))
          resp = tw.Memoize(resp, tweetId, "height",
            fmt.Sprintf("%d", h))
        }
      }else {
        resp = tw.Memoize(resp, tweetId, "image_source",
          fmt.Sprintf("%v", source["expanded_url"]))
      }
      resp = tw.Memoize(resp, tweetId, "tweet_id",
        fmt.Sprintf("%v", tweetId))
      resp = tw.Memoize(resp, tweetId, "tweet",
        fmt.Sprintf("%v", tweet.Text()))
      resp = tw.Memoize(resp, tweetId, "user",
        fmt.Sprintf("%v (@%v)", user.Name(), user.ScreenName()))
      resp = tw.Memoize(resp, tweetId, "created",
        fmt.Sprintf("%v", tweet.CreatedAt().Format(time.RFC1123)))

      resp["source"] = "search"
      fmt.Println("+ return parsed tweet")
      t <- resp
    }
  }
  for i:=len(tweets)-1; i >= 0; i-- {
    conn.Do("LPUSH", "tweets", tweets[i])
  }
}

func (tw Twitter) RetrieveSinceId() (string, bool){
  reply, _ := redis.String( conn.Do("LINDEX", "tweets", 0) )
  if reply == "" {
    return reply, false
  } else {
    return reply, true
  }
}

func (tw Twitter) TwitterImages() url.Values {
  query := url.Values{}
  query.Set("q", "#beautifulsky filter:images")
  query.Set("result_type", "mixed")
  query.Set("count", "1000")

  if id, present := tw.RetrieveSinceId(); present {
    fmt.Println("+ sinceId(%v), using query %v", id, query)
    query.Set("since_id", id)
  }else{
    fmt.Println("+ no sinceId, using query %v", query)
  }
  return query
}

func (tw Twitter) Instagram() url.Values {
  query := url.Values{}
  query.Set("q", "#beautifulsky")
  query.Set("result_type", "mixed")
  query.Set("count", "1000")

  if id, present := tw.RetrieveSinceId(); present {
    fmt.Println("+ sinceId(%v), using query %v", id, query)
    query.Set("since_id", id)
  }else{
    fmt.Println("+ no sinceId, using query %v", query)
  }
  return query
}

func (tw Twitter) SearchTweets(
  query url.Values,
  k chan *twittergo.SearchResults,
  searchDone chan bool){

  var (
    err     error
    results *twittergo.SearchResults
    req     *http.Request
    resp    *twittergo.APIResponse
    client *twittergo.Client
  )
  client, _ = tw.LoadCredentials()

  url := fmt.Sprintf("/1.1/search/tweets.json?%v", query.Encode())

  req, err = http.NewRequest("GET", url, nil)
  if err != nil {
    fmt.Printf("Could not parse request: %v\n", err)
  }

  fmt.Printf("+ Getting tweets: %v\n", req)
  resp, err = client.SendRequest(req)
  if err != nil {
    fmt.Printf("Could not send request: %v\n", err)
  }

  results = &twittergo.SearchResults{}
  err = resp.Parse(results)
  if err != nil {
    fmt.Printf("Problem parsing response: %v\n", err)
  }

  k <- results
  searchDone <- true
}

func (tw Twitter) TweetsFromCache(t chan map[string]string) {
  l, _ := redis.Int(conn.Do("LLEN", "tweets"))
  for i:=0; i< l && i< 50; i++ {
    tweetId, _ := redis.String(conn.Do("LINDEX", "tweets", i))
    var resp map[string]string = make(map[string]string)
    for _, key := range []string{"tweet", "image_source", "image_url","user", "created", "tweet_id", "width", "height"} {
      reply,_ := redis.String( conn.Do("HGET", "tweet:"+tweetId, key) )
      resp[key] = reply
    }
    fmt.Printf("+ return cache #%v\n", i+1)
    t <- resp
  }
}

func (tw Twitter) PrintRateLimit(r chan *twittergo.APIResponse) {
  fmt.Printf("start PrintRateLimit \n")
  resp := <-r
  if resp.HasRateLimit() {
    fmt.Printf("Rate limit:           %v\n", resp.RateLimit())
    fmt.Printf("Rate limit remaining: %v\n", resp.RateLimitRemaining())
    fmt.Printf("Rate limit reset:     %v\n", resp.RateLimitReset())
  } else {
    fmt.Printf("Could not parse rate limit from response.\n")
  }
  fmt.Printf("end PrintRateLimit \n")
}
<|endoftext|>


<commit_before>package riffle

import (
	"github.com/exis-io/core"
	"github.com/exis-io/core/goRiffle"
)

var fabric string = core.ProudctionFabric

type App struct {
	coreApp core.App
}

type Domain struct {
	app        *App
	coreDomain core.Domain
}

func (a *App) Init() {
	a.coreApp = core.NewApp()
}

func (a *App) NewDomain(name string) Domain {
	return Domain{
		app:        a,
		coreDomain: a.coreApp.NewDomain(name),
	}
}

// Blocks on callbacks from the core. TODO: trigger a close meta callback when connection is lost
func (a *App) Receive() string {
	return core.MantleMarshall(a.coreApp.CallbackListen())
}

func (d *Domain) Join(cb uint, eb uint) {
	if c, err := goRiffle.Open(fabric); err != nil {
		d.app.coreApp.CallbackSend(eb, err.Error())
	} else {
		c.App = d.app.coreApp
		if err := d.coreDomain.Join(c); err != nil {
			d.app.coreApp.CallbackSend(eb, err.Error())
		} else {
			d.app.coreApp.CallbackSend(cb)
		}
	}
}

func (d Domain) Subscribe(cb uint, endpoint string) {
	go func() {
		d.coreDomain.Subscribe(endpoint, cb, make([]interface{}, 0))
	}()
}

func (d Domain) Register(cb uint, endpoint string) {
	go func() {
		d.coreDomain.Register(endpoint, cb, make([]interface{}, 0))
	}()
}

// Args are string encoded json
func (d Domain) Publish(cb uint, endpoint string, args string) {
	go func() {
		d.coreDomain.Publish(endpoint, cb, core.MantleUnmarshal(args))
	}()
}

func (d Domain) Call(cb uint, endpoint string, args string) {
	go func() {
		d.coreDomain.Call(endpoint, cb, core.MantleUnmarshal(args))
	}()
}

func (a *App) Yield(request uint, args string) {
	go func() {
		a.coreApp.Yield(request, core.MantleUnmarshal(args))
	}()
}

func (d Domain) Unsubscribe(endpoint string) {

}

func (d Domain) Unregister(endpoint string) {

}

func (d Domain) Leave() {

}

func SetLogLevelApp()       { core.LogLevel = core.LogLevelApp }
func SetLogLevelErr()       { core.LogLevel = core.LogLevelErr }
func SetLogLevelWarn()      { core.LogLevel = core.LogLevelWarn }
func SetLogLevelInfo()      { core.LogLevel = core.LogLevelInfo }
func SetLogLevelDebug()     { core.LogLevel = core.LogLevelDebug }
func SetLoggingLevel(l int) { core.LogLevel = l }

func SetDevFabric()              { fabric = core.DevFabric }
func SetSandboxFabric()          { fabric = core.SandboxFabric }
func SetProductionFabric()       { fabric = core.ProudctionFabric }
func SetLocalFabric()            { fabric = core.LocalFabric }
func SetCustomFabric(url string) { fabric = url }

func Application(s string) { core.Application("%s", s) }
func Debug(s string)       { core.Debug("%s", s) }
func Info(s string)        { core.Info("%s", s) }
func Warn(s string)        { core.Warn("%s", s) }
func Error(s string)       { core.Error("%s", s) }
<commit_msg>returning values<commit_after>package riffle

import (
	"github.com/exis-io/core"
	"github.com/exis-io/core/goRiffle"
)

var fabric string = core.ProudctionFabric

type App struct {
	coreApp core.App
}

type Domain struct {
	app        *App
	coreDomain core.Domain
}

func (a *App) Init() {
	a.coreApp = core.NewApp()
}

func (a *App) NewDomain(name string) Domain {
	return Domain{
		app:        a,
		coreDomain: a.coreApp.NewDomain(name),
	}
}

// Blocks on callbacks from the core. TODO: trigger a close meta callback when connection is lost
func (a *App) Receive() string {
	return core.MantleMarshall(a.coreApp.CallbackListen())
}

func (d *Domain) Join(cb uint, eb uint) {
	if c, err := goRiffle.Open(fabric); err != nil {
		d.app.coreApp.CallbackSend(eb, err.Error())
	} else {
		c.App = d.app.coreApp

		if err := d.coreDomain.Join(c); err != nil {
			d.app.coreApp.CallbackSend(eb, err.Error())
		} else {
			d.app.coreApp.CallbackSend(cb)
		}
	}
}

func (d Domain) Subscribe(cb uint, endpoint string) {
	go func() {
		d.coreDomain.Subscribe(endpoint, cb, make([]interface{}, 0))
	}()
}

func (d Domain) Register(cb uint, endpoint string) {
	go func() {
		d.coreDomain.Register(endpoint, cb, make([]interface{}, 0))
	}()
}

// Args are string encoded json
func (d Domain) Publish(cb uint, endpoint string, args string) {
	go func() {
		d.coreDomain.Publish(endpoint, cb, core.MantleUnmarshal(args))
	}()
}

func (d Domain) Call(cb uint, endpoint string, args string) {
	go func() {
		d.coreDomain.Call(endpoint, cb, core.MantleUnmarshal(args))
	}()
}

func (a *App) Yield(request uint, args string) {
	go func() {
		a.coreApp.Yield(request, core.MantleUnmarshal(args))
	}()
}

func (d Domain) Unsubscribe(endpoint string) {

}

func (d Domain) Unregister(endpoint string) {

}

func (d Domain) Leave() {

}

func SetLogLevelApp()       { core.LogLevel = core.LogLevelApp }
func SetLogLevelErr()       { core.LogLevel = core.LogLevelErr }
func SetLogLevelWarn()      { core.LogLevel = core.LogLevelWarn }
func SetLogLevelInfo()      { core.LogLevel = core.LogLevelInfo }
func SetLogLevelDebug()     { core.LogLevel = core.LogLevelDebug }
func SetLoggingLevel(l int) { core.LogLevel = l }

func SetDevFabric()              { fabric = core.DevFabric }
func SetSandboxFabric()          { fabric = core.SandboxFabric }
func SetProductionFabric()       { fabric = core.ProudctionFabric }
func SetLocalFabric()            { fabric = core.LocalFabric }
func SetCustomFabric(url string) { fabric = url }

func Application(s string) { core.Application("%s", s) }
func Debug(s string)       { core.Debug("%s", s) }
func Info(s string)        { core.Info("%s", s) }
func Warn(s string)        { core.Warn("%s", s) }
func Error(s string)       { core.Error("%s", s) }
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"path"
	"reflect"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/cblomart/vsphere-graphite/backend"
	"github.com/cblomart/vsphere-graphite/config"
	"github.com/cblomart/vsphere-graphite/vsphere"

	"github.com/takama/daemon"

	"code.cloudfoundry.org/bytefmt"

	"github.com/vmware/govmomi/vim25/types"
)

const (
	// name of the service
	name        = "vsphere-graphite"
	description = "send vsphere stats to graphite"
)

var dependencies = []string{}

var stdlog, errlog *log.Logger

// Service has embedded daemon
type Service struct {
	daemon.Daemon
}

// EntityQuery : Informations to query about an entity
type EntityQuery struct {
	Name    string
	Entity  types.ManagedObjectReference
	Metrics []int
}

func queryVCenter(vcenter vsphere.VCenter, config config.Configuration, channel *chan backend.Point) {
	vcenter.Query(config.Interval, config.Domain, channel)
}

// Manage by daemon commands or run the daemon
func (service *Service) Manage() (string, error) {

	usage := "Usage: vsphere-graphite install | remove | start | stop | status"

	// if received any kind of command, do it
	if len(os.Args) > 1 {
		command := os.Args[1]
		switch command {
		case "install":
			return service.Install()
		case "remove":
			return service.Remove()
		case "start":
			return service.Start()
		case "stop":
			return service.Stop()
		case "status":
			return service.Status()
		default:
			return usage, nil
		}
	}

	stdlog.Println("Starting daemon:", path.Base(os.Args[0]))

	// read the configuration
	file, err := os.Open("/etc/" + path.Base(os.Args[0]) + ".json")
	if err != nil {
		return "Could not open configuration file", err
	}
	jsondec := json.NewDecoder(file)
	config := config.Configuration{}
	err = jsondec.Decode(&config)
	if err != nil {
		return "Could not decode configuration file", err
	}

	if config.FlushSize == 0 {
		config.FlushSize = 1000
	}

	if config.Profiling {
		f, err := ioutil.TempFile("/tmp", "vsphere-graphite-cpu.profile")
		stdlog.Println("Will write cpu profiling to: ", f.Name())
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
		defer pprof.StopCPUProfile()
	}

	//force backend values to environement varialbles if present
	s := reflect.ValueOf(&config.Backend).Elem()
	numfields := s.NumField()
	for i := 0; i < numfields; i++ {
		f := s.Field(i)
		if f.CanSet() {
			//exported field
			envname := strings.ToUpper(s.Type().Name() + "_" + s.Type().Field(i).Name)
			envval := os.Getenv(envname)
			if len(envval) > 0 {
				//environment variable set with name
				switch ftype := f.Type().Name(); ftype {
				case "string":
					f.SetString(envval)
				case "int":
					val, err := strconv.ParseInt(envval, 10, 64)
					if err == nil {
						f.SetInt(val)
					}
				}
			}
		}
	}

	for _, vcenter := range config.VCenters {
		vcenter.Init(config.Metrics, stdlog, errlog)
	}

	err = config.Backend.Init(stdlog, errlog)
	if err != nil {
		return "Could not initialize backend", err
	}
	defer config.Backend.Disconnect()

	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

	// Set up a channel to receive the metrics
	metrics := make(chan backend.Point, config.FlushSize)

	// Set up a ticker to collect metrics at givent interval
	ticker := time.NewTicker(time.Second * time.Duration(config.Interval))
	defer ticker.Stop()

	// Set up a ticker to garbadge collect
	memtimer := time.NewTicker(time.Second * time.Duration(config.Interval))
	defer memtimer.Stop()

	// Start retriveing and sending metrics
	stdlog.Println("Retrieving metrics")
	for _, vcenter := range config.VCenters {
		go queryVCenter(*vcenter, config, &metrics)
	}

	// Memory statisctics
	var memstats runtime.MemStats

	// buffer for points to send
	pointbuffer := make([]backend.Point, config.FlushSize)
	bufferindex := 0

	for {
		select {
		case value := <-metrics:
			pointbuffer[bufferindex] = value
			bufferindex++
			if bufferindex == len(pointbuffer) {
				config.Backend.SendMetrics(pointbuffer)
				stdlog.Printf("Sent %d logs to backend", len(pointbuffer))
				for i := 0; i < len(pointbuffer); i++ {
					pointbuffer[i] = backend.Point{}
				}
				bufferindex = 0
			}
		case <-ticker.C:
			stdlog.Println("Retrieving metrics")
			for _, vcenter := range config.VCenters {
				go queryVCenter(*vcenter, config, &metrics)
			}
		case <-memtimer.C:
			runtime.GC()
			runtime.ReadMemStats(&memstats)
			stdlog.Println("Memory usage :", bytefmt.ByteSize(memstats.Sys))
		case killSignal := <-interrupt:
			stdlog.Println("Got signal:", killSignal)
			if bufferindex > 0 {
				config.Backend.SendMetrics(pointbuffer[:bufferindex])
				stdlog.Printf("Sent %d logs to backend", bufferindex)
			}
			if killSignal == os.Interrupt {
				return "Daemon was interrupted by system signal", nil
			}
			return "Daemon was killed", nil
		}
	}
}

func init() {
	stdlog = log.New(os.Stdout, "", log.Ldate|log.Ltime)
	errlog = log.New(os.Stderr, "", log.Ldate|log.Ltime)
}

func main() {
	srv, err := daemon.New(name, description, dependencies...)
	if err != nil {
		errlog.Println("Error: ", err)
		os.Exit(1)
	}
	service := &Service{srv}
	status, err := service.Manage()
	if err != nil {
		errlog.Println(status, "Error: ", err)
		os.Exit(1)
	}
	fmt.Println(status)
}
<commit_msg>Garbadge collect after sent<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"path"
	"reflect"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/cblomart/vsphere-graphite/backend"
	"github.com/cblomart/vsphere-graphite/config"
	"github.com/cblomart/vsphere-graphite/vsphere"

	"github.com/takama/daemon"

	"code.cloudfoundry.org/bytefmt"

	"github.com/vmware/govmomi/vim25/types"
)

const (
	// name of the service
	name        = "vsphere-graphite"
	description = "send vsphere stats to graphite"
)

var dependencies = []string{}

var stdlog, errlog *log.Logger

// Service has embedded daemon
type Service struct {
	daemon.Daemon
}

// EntityQuery : Informations to query about an entity
type EntityQuery struct {
	Name    string
	Entity  types.ManagedObjectReference
	Metrics []int
}

func queryVCenter(vcenter vsphere.VCenter, config config.Configuration, channel *chan backend.Point) {
	vcenter.Query(config.Interval, config.Domain, channel)
}

// Manage by daemon commands or run the daemon
func (service *Service) Manage() (string, error) {

	usage := "Usage: vsphere-graphite install | remove | start | stop | status"

	// if received any kind of command, do it
	if len(os.Args) > 1 {
		command := os.Args[1]
		switch command {
		case "install":
			return service.Install()
		case "remove":
			return service.Remove()
		case "start":
			return service.Start()
		case "stop":
			return service.Stop()
		case "status":
			return service.Status()
		default:
			return usage, nil
		}
	}

	stdlog.Println("Starting daemon:", path.Base(os.Args[0]))

	// read the configuration
	file, err := os.Open("/etc/" + path.Base(os.Args[0]) + ".json")
	if err != nil {
		return "Could not open configuration file", err
	}
	jsondec := json.NewDecoder(file)
	config := config.Configuration{}
	err = jsondec.Decode(&config)
	if err != nil {
		return "Could not decode configuration file", err
	}

	if config.FlushSize == 0 {
		config.FlushSize = 1000
	}

	if config.Profiling {
		f, err := ioutil.TempFile("/tmp", "vsphere-graphite-cpu.profile")
		stdlog.Println("Will write cpu profiling to: ", f.Name())
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
		defer pprof.StopCPUProfile()
	}

	//force backend values to environement varialbles if present
	s := reflect.ValueOf(&config.Backend).Elem()
	numfields := s.NumField()
	for i := 0; i < numfields; i++ {
		f := s.Field(i)
		if f.CanSet() {
			//exported field
			envname := strings.ToUpper(s.Type().Name() + "_" + s.Type().Field(i).Name)
			envval := os.Getenv(envname)
			if len(envval) > 0 {
				//environment variable set with name
				switch ftype := f.Type().Name(); ftype {
				case "string":
					f.SetString(envval)
				case "int":
					val, err := strconv.ParseInt(envval, 10, 64)
					if err == nil {
						f.SetInt(val)
					}
				}
			}
		}
	}

	for _, vcenter := range config.VCenters {
		vcenter.Init(config.Metrics, stdlog, errlog)
	}

	err = config.Backend.Init(stdlog, errlog)
	if err != nil {
		return "Could not initialize backend", err
	}
	defer config.Backend.Disconnect()

	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

	// Set up a channel to receive the metrics
	metrics := make(chan backend.Point, config.FlushSize)

	// Set up a ticker to collect metrics at givent interval
	ticker := time.NewTicker(time.Second * time.Duration(config.Interval))
	defer ticker.Stop()

	// Start retriveing and sending metrics
	stdlog.Println("Retrieving metrics")
	for _, vcenter := range config.VCenters {
		go queryVCenter(*vcenter, config, &metrics)
	}

	// Memory statisctics
	var memstats runtime.MemStats
	// timer to execute memory collection
	memtimer := time.NewTimer(time.Second * time.Duration(10))

	// buffer for points to send
	pointbuffer := make([]backend.Point, config.FlushSize)
	bufferindex := 0

	for {
		select {
		case value := <-metrics:
			pointbuffer[bufferindex] = value
			bufferindex++
			if bufferindex == len(pointbuffer) {
				config.Backend.SendMetrics(pointbuffer)
				stdlog.Printf("Sent %d logs to backend", len(pointbuffer))
				for i := 0; i < len(pointbuffer); i++ {
					pointbuffer[i] = backend.Point{}
				}
				bufferindex = 0
				// reset timer as data has been set
				if !memtimer.Stop() {
					select {
					case <-memtimer.C:
					default:
					}
				}
				memtimer.Reset(time.Second * time.Duration(10))
			}
		case <-ticker.C:
			stdlog.Println("Retrieving metrics")
			for _, vcenter := range config.VCenters {
				go queryVCenter(*vcenter, config, &metrics)
			}
		case <-memtimer.C:
			runtime.GC()
			runtime.ReadMemStats(&memstats)
			stdlog.Println("Memory usage :", bytefmt.ByteSize(memstats.Sys))
		case killSignal := <-interrupt:
			stdlog.Println("Got signal:", killSignal)
			if bufferindex > 0 {
				config.Backend.SendMetrics(pointbuffer[:bufferindex])
				stdlog.Printf("Sent %d logs to backend", bufferindex)
			}
			if killSignal == os.Interrupt {
				return "Daemon was interrupted by system signal", nil
			}
			return "Daemon was killed", nil
		}
	}
}

func init() {
	stdlog = log.New(os.Stdout, "", log.Ldate|log.Ltime)
	errlog = log.New(os.Stderr, "", log.Ldate|log.Ltime)
}

func main() {
	srv, err := daemon.New(name, description, dependencies...)
	if err != nil {
		errlog.Println("Error: ", err)
		os.Exit(1)
	}
	service := &Service{srv}
	status, err := service.Manage()
	if err != nil {
		errlog.Println(status, "Error: ", err)
		os.Exit(1)
	}
	fmt.Println(status)
}
<|endoftext|>


<commit_before>package boxutil

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/matryer/is"
)

func init() {
	// quicker for testing
	readyCheckInterval = 100 * time.Millisecond
}

func TestStatusChan(t *testing.T) {
	is := is.New(t)

	i := &testBox{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	status := StatusChan(ctx, i)
	is.Equal(<-status, "starting...")
	i.setReady()
	is.Equal(<-status, "ready")
	i.setError()
	is.Equal(<-status, "unavailable")
	i.clearError()
	is.Equal(<-status, "ready")

}

func TestWaitForReady(t *testing.T) {
	is := is.New(t)
	i := &testBox{}
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(300*time.Millisecond, cancel)
	go func() {
		time.Sleep(200 * time.Millisecond)
		i.setReady()
	}()
	err := WaitForReady(ctx, i)
	is.NoErr(err)
}

func TestWaitForReadyTimeout(t *testing.T) {
	is := is.New(t)
	i := &testBox{}
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(100*time.Millisecond, cancel)
	go func() {
		time.Sleep(200 * time.Millisecond)
		i.setReady()
	}()
	err := WaitForReady(ctx, i)
	is.Equal(err, ErrCanceled)
}

type testBox struct {
	lock  sync.Mutex
	ready bool
	err   error
}

func (i *testBox) setReady() {
	i.lock.Lock()
	defer i.lock.Unlock()
	i.ready = true
}

func (i *testBox) setError() {
	i.lock.Lock()
	defer i.lock.Unlock()
	i.err = errors.New("cannot reach server")
}

func (i *testBox) clearError() {
	i.lock.Lock()
	defer i.lock.Unlock()
	i.err = nil
}

func (i *testBox) Info() (*Info, error) {
	i.lock.Lock()
	defer i.lock.Unlock()
	if i.err != nil {
		return nil, i.err
	}
	if i.ready {
		return &Info{Status: "ready"}, nil
	}
	return &Info{Status: "starting..."}, nil
}
<commit_msg>better test<commit_after>package boxutil

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/matryer/is"
)

func init() {
	// quicker for testing
	readyCheckInterval = 100 * time.Millisecond
}

func TestStatusChan(t *testing.T) {
	is := is.New(t)

	i := &testBox{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	status := StatusChan(ctx, i)
	is.Equal(<-status, "starting...")
	i.setReady()
	is.Equal(<-status, "ready")
	i.setError()
	is.Equal(<-status, "unavailable")
	i.clearError()
	is.Equal(<-status, "ready")

}

func TestWaitForReady(t *testing.T) {
	is := is.New(t)
	i := &testBox{}
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(300*time.Millisecond, cancel)
	go func() {
		i.setReady()
	}()
	err := WaitForReady(ctx, i)
	is.NoErr(err)
}

func TestWaitForReadyTimeout(t *testing.T) {
	is := is.New(t)
	i := &testBox{}
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(100*time.Millisecond, cancel)
	go func() {
		time.Sleep(200 * time.Millisecond)
		i.setReady()
	}()
	err := WaitForReady(ctx, i)
	is.Equal(err, ErrCanceled)
}

type testBox struct {
	lock  sync.Mutex
	ready bool
	err   error
}

func (i *testBox) setReady() {
	i.lock.Lock()
	defer i.lock.Unlock()
	i.ready = true
}

func (i *testBox) setError() {
	i.lock.Lock()
	defer i.lock.Unlock()
	i.err = errors.New("cannot reach server")
}

func (i *testBox) clearError() {
	i.lock.Lock()
	defer i.lock.Unlock()
	i.err = nil
}

func (i *testBox) Info() (*Info, error) {
	i.lock.Lock()
	defer i.lock.Unlock()
	if i.err != nil {
		return nil, i.err
	}
	if i.ready {
		return &Info{Status: "ready"}, nil
	}
	return &Info{Status: "starting..."}, nil
}
<|endoftext|>


<commit_before><commit_msg>Using var block.<commit_after><|endoftext|>


<commit_before>/*
   Copyright (C) 2017 Enrico Bianchi (enrico.bianchi@gmail.com)
   Project       Syncimap
   Description   A rsync like IMAP syncronization tool
   License       Apache License 2.0 (see LICENSE for details)
*/

package imap

import (
	"errors"
	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/client"
	"net/url"
	"strings"
)

type Connection struct {
	Client *client.Client
}

func (conn *Connection) Dial(uri *url.URL) error {
	var err error

	port := uri.Port()
	switch uri.Scheme {
	case "imap":
		if port == "" {
			port = "143"
		}
		conn.Client, err = client.Dial(uri.Host + ":" + port)
	case "imaps":
		if port == "" {
			port = "993"
		}
		conn.Client, err = client.DialTLS(uri.Host+":"+port, nil)
	default:
		return errors.New("Scheme not supported: " + uri.Scheme)
	}

	caps, err := conn.Client.Capability()
	if err != nil {
		return err
	}

	if caps["STARTTLS"] {
		conn.Client.StartTLS(nil)
	}

	user := uri.User.Username()
	password, _ := uri.User.Password()

	err = conn.Client.Login(user, password)
	if err != nil {
		return err
	}

	return err
}

func (conn *Connection) Folders(folder string) []*imap.MailboxInfo {
	var folders []*imap.MailboxInfo
	mailboxes := make(chan *imap.MailboxInfo)

	done := make(chan error, 1)
	go func() {
		done <- conn.Client.List("", folder, mailboxes)
	}()

	for mailbox := range mailboxes {
		folders = append(folders, mailbox)
	}

	return folders
}

func (conn *Connection) FolderExist(folder string) bool {
	conn.Client.Close()

	_, err := conn.Client.Select(folder, false)
	if err != nil {
		return false
	}

	return true
}

func (conn *Connection) FolderCreate(folder *imap.MailboxInfo) error {

	parent := "INBOX"
	name := folder.Name
	if strings.Contains(name, folder.Delimiter) {
		elements := strings.Split(folder.Name, folder.Delimiter)
		parent = parent + folder.Delimiter + strings.Join(elements[:len(elements)-1], folder.Delimiter)
		name = elements[len(elements)-1]
	}

	conn.Client.Select(parent, false)
	err := conn.Client.Create(folder.Name)

	return err
}
<commit_msg>Checked error<commit_after>/*
   Copyright (C) 2017 Enrico Bianchi (enrico.bianchi@gmail.com)
   Project       Syncimap
   Description   A rsync like IMAP syncronization tool
   License       Apache License 2.0 (see LICENSE for details)
*/

package imap

import (
	"errors"
	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/client"
	"net/url"
	"strings"
)

type Connection struct {
	Client *client.Client
}

func (conn *Connection) Dial(uri *url.URL) error {
	var err error

	port := uri.Port()
	switch uri.Scheme {
	case "imap":
		if port == "" {
			port = "143"
		}
		conn.Client, err = client.Dial(uri.Host + ":" + port)
	case "imaps":
		if port == "" {
			port = "993"
		}
		conn.Client, err = client.DialTLS(uri.Host+":"+port, nil)
	default:
		return errors.New("Scheme not supported: " + uri.Scheme)
	}

	if err != nil {
		return err
	}

	caps, err := conn.Client.Capability()
	if err != nil {
		return err
	}

	if caps["STARTTLS"] {
		conn.Client.StartTLS(nil)
	}

	user := uri.User.Username()
	password, _ := uri.User.Password()

	err = conn.Client.Login(user, password)
	if err != nil {
		return err
	}

	return err
}

func (conn *Connection) Folders(folder string) []*imap.MailboxInfo {
	var folders []*imap.MailboxInfo
	mailboxes := make(chan *imap.MailboxInfo)

	done := make(chan error, 1)
	go func() {
		done <- conn.Client.List("", folder, mailboxes)
	}()

	for mailbox := range mailboxes {
		folders = append(folders, mailbox)
	}

	return folders
}

func (conn *Connection) FolderExist(folder string) bool {
	conn.Client.Close()

	_, err := conn.Client.Select(folder, false)
	if err != nil {
		return false
	}

	return true
}

func (conn *Connection) FolderCreate(folder *imap.MailboxInfo) error {

	parent := "INBOX"
	name := folder.Name
	if strings.Contains(name, folder.Delimiter) {
		elements := strings.Split(folder.Name, folder.Delimiter)
		parent = parent + folder.Delimiter + strings.Join(elements[:len(elements)-1], folder.Delimiter)
		name = elements[len(elements)-1]
	}

	conn.Client.Select(parent, false)
	err := conn.Client.Create(folder.Name)

	return err
}
<|endoftext|>


<commit_before>package script

import "testing"

func TestPrintingConditionalOperations(t *testing.T) {
	a := Operation{
		opType:  OpConditional,
		condOp1: "Var[1]",
		condOp2: "Var[2]",
		condOp:  "==",
		condDst: 0x2345,
	}
	if a.String() != "unless (Var[1] == Var[2]) goto 2345" {
		t.Errorf("conditional operation `%v` is not properly formatted", a.String())
	}
}

func TestPrintingAssignmentOperations(t *testing.T) {
	a := Operation{
		opType:    OpAssignment,
		assignDst: "Var[123]",
		assignVal: "999",
	}
	if a.String() != "Var[123] = 999" {
		t.Errorf("operation `%v` is not properly formatted", a.String())
	}
}

func TestPrintingCallOperations(t *testing.T) {
	a := Operation{
		opType:     OpCall,
		callMethod: "myMethod",
		callParams: []string{
			"123", "456",
		},
	}

	if a.String() != "myMethod(123, 456)" {
		t.Errorf("conditional operation `%v` is not properly formatted", a.String())
	}

	b := Operation{
		opType:     OpCall,
		callMethod: "myMethod",
		callMap: map[string]string{
			"x": "1234",
			"y": "6789",
		},
	}

	if b.String() != "myMethod(x=1234, y=6789)" &&
		b.String() != "myMethod(y=6789, x=1234)" {
		t.Errorf("conditional operation `%v` is not properly formatted", b.String())
	}
}
<commit_msg>Cleaner test<commit_after>package script

import "testing"

func TestPrintingConditionalOperations(t *testing.T) {
	a := Operation{
		opType:  OpConditional,
		condOp1: "Var[1]",
		condOp2: "Var[2]",
		condOp:  "==",
		condDst: 0x2345,
	}
	if a.String() != "unless (Var[1] == Var[2]) goto 2345" {
		t.Errorf("conditional operation `%v` is not properly formatted", a.String())
	}
}

func TestPrintingAssignmentOperations(t *testing.T) {
	a := Operation{
		opType:    OpAssignment,
		assignDst: "Var[123]",
		assignVal: "999",
	}
	if a.String() != "Var[123] = 999" {
		t.Errorf("operation `%v` is not properly formatted", a.String())
	}
}

func TestPrintingCallOperations(t *testing.T) {
	a := Operation{
		opType:     OpCall,
		callMethod: "myMethod",
		callParams: []string{
			"123", "456",
		},
	}

	if a.String() != "myMethod(123, 456)" {
		t.Errorf("conditional operation `%v` is not properly formatted", a.String())
	}

	b := Operation{
		opType:     OpCall,
		callMethod: "myMethod",
		callMap: map[string]string{
			"x": "1234",
			"y": "6789",
		},
	}

	opPrint := b.String()
	if opPrint != "myMethod(x=1234, y=6789)" &&
		opPrint != "myMethod(y=6789, x=1234)" {
		t.Errorf("conditional operation `%v` is not properly formatted", opPrint)
	}
}
<|endoftext|>


<commit_before>package sendmail

import (
	"bds/lib/mailstore"
	"bds/lib/smtp"
	"errors"
	log "github.com/Sirupsen/logrus"
	"net"
	"strings"
	"sync"
)

var ErrNoLocalMailDelivery = errors.New("no local mail store for user")

// mail bounce handler
// paramters are (recipiant email address, from email address, the filepath of the message, network related error or nil for regular bounce)
type Bouncer func(string, string, string, error)

type connection struct {
	addr   string
	cl     *smtp.Client
	access sync.RWMutex
}

// safely visit a conection with a function
func (c *connection) visit(visitor func(*smtp.Client) error) (err error) {
	c.access.Lock()
	err = visitor(c.cl)
	c.access.Unlock()
	return
}

func (c *connection) Quit() {
	c.cl.Quit()
}

// mail delivery
type Mailer struct {
	// get mail storage for local user
	Local mailstore.MailRouter
	// a dial function to obtain outbound smtp client
	Dial Dialer
	// number of times to try to deliver mail
	// 0 for unlimited
	Retries int
	// called if mail bounces or fails to deliver
	Bounce Bouncer
	// domain resolver function
	Resolve Resolver
	// delivery success hook, called with (recipiant email address, from email address)
	Success func(string, string)
	// for pipelining
	conns map[string]*connection
	// mutex for conns
	cmtx sync.RWMutex
}

// create a new pooled mailer
func NewMailer() *Mailer {
	return &Mailer{
		conns: make(map[string]*connection),
	}
}

// call a visitor for each connection in connection poll
func (s *Mailer) foreach(visitor func(*smtp.Client) error) {
	var conns []*connection
	s.cmtx.Lock()
	for _, conn := range s.conns {
		conns = append(conns, conn)
	}
	s.cmtx.Unlock()
	for _, conn := range conns {
		// in paralell
		go func(c *connection, v func(*smtp.Client) error) {
			err := c.visit(v)
			if err != nil {
				// remove connection on error
				s.delConn(c.addr)
			}
		}(conn, visitor)
	}
}

// delete connection from pool
func (s *Mailer) delConn(addr string) {
	s.cmtx.Lock()
	// safe delete
	c, ok := s.conns[addr]
	if ok {
		c.Quit()
		delete(s.conns, addr)
	}
	s.cmtx.Unlock()
}

// visit a pooled connection in a safe way
func (s *Mailer) visitConn(n, addr, localname string, dialer Dialer, visitor func(*smtp.Client) error) (err error) {
	var c *connection
	c, err = s.getConn(n, addr, localname, dialer)
	if err == nil {
		err = c.visit(visitor)
		if err != nil {
			log.Errorf("failed to visit connection: %s", err.Error())
			s.delConn(addr)
		}
	}
	return
}

// get connection from pool, create if not there
func (s *Mailer) getConn(n, addr, localname string, dial Dialer) (sc *connection, err error) {
	s.cmtx.Lock()
	defer s.cmtx.Unlock()
	var ok bool
	sc, ok = s.conns[addr]
	if !ok {
		var c net.Conn
		c, err = dial(n, addr)
		if err == nil {
			// new connection
			var cl *smtp.Client
			cl, err = smtp.NewClient(c, addr)
			if err != nil {
				log.Errorf("failed to dial: %s", err.Error())
				return
			}
			err = cl.Hello(localname)
			if err != nil {
				log.Errorf("failed to helo: %s", err.Error())
				cl.Quit()
				return
			}
			sc = &connection{
				cl:   cl,
				addr: addr,
			}
			// success
			s.conns[addr] = sc
		}
	}
	return
}

// try delivering mail
// returns a DeliveryJob that can be cancelled
func (s *Mailer) Deliver(recip, from string, msg mailstore.Message) (d DeliverJob) {
	log.Infof("Delivering %s to %s from %s", msg.Filepath(), recip, from)
	dialer := s.Dial
	if dialer == nil {
		dialer = net.Dial
	}

	bounce := s.Bounce

	resolver := s.Resolve

	if resolver == nil {
		resolver = func(name string) (a net.Addr, err error) {
			log.Debugf("mx lookup for %s", name)
			var mx []*net.MX
			mx, err = net.LookupMX(name)
			if err != nil && mx != nil {
				for _, m := range mx {
					var ips []net.IP
					log.Debugf("resolve mx record %s", m.Host)
					ips, err = net.LookupIP(m.Host)
					if err == nil {
						for _, ip := range ips {
							a, err = net.ResolveIPAddr("ip", ip.String())
							if err == nil {
								log.Debugf("resolved %s to %s", name, a)
								return
							}
						}
					}
				}
			}
			return
		}
	}
	var st mailstore.Store
	if s.Local != nil {
		st, _ = s.Local.FindStoreFor(recip)
	}

	if st == nil {
		d = &RemoteDeliverJob{
			unlimited: s.Retries == 0,
			cancel:    false,
			retries:   s.Retries,
			visit: func(f func(*smtp.Client) error) error {
				parts := strings.Split(recip, "@")
				if len(parts) == 2 {
					r_addr := parts[1]
					a, err := resolver(r_addr)
					if err == nil {
						err = s.visitConn(a.Network(), r_addr, r_addr, dialer, f)
					} else {
						log.Warnf("failed to resolve %s: %s", r_addr, err.Error())
					}
					return err
				} else {
					log.Warnf("bad email address %s", recip)
				}
				return nil
			},
			bounce:    bounce,
			recip:     recip,
			from:      from,
			fpath:     msg.Filepath(),
			result:    make(chan bool),
			delivered: s.Success,
		}
	} else {
		d = &LocalDeliverJob{
			st:     st,
			result: make(chan bool),
			fpath:  msg.Filepath(),
		}
	}
	return
}

// gracefully quit all polled connections and close down
func (m *Mailer) Quit() {
	log.Info("shutting down pooled mailer")
	m.cmtx.Lock()
	// collect connections
	var conns []*connection
	for _, c := range m.conns {
		conns = append(conns, c)
	}
	// close all connections
	for _, c := range conns {
		c.visit(func(cl *smtp.Client) error {
			return cl.Quit()
		})
	}
	// collect all names
	var names []string
	for n, _ := range m.conns {
		names = append(names, n)
	}
	// delete all connections
	for _, n := range names {
		delete(m.conns, n)
	}
	m.cmtx.Unlock()
}
<commit_msg>use resolved address<commit_after>package sendmail

import (
	"bds/lib/mailstore"
	"bds/lib/smtp"
	"errors"
	log "github.com/Sirupsen/logrus"
	"net"
	"strings"
	"sync"
)

var ErrNoLocalMailDelivery = errors.New("no local mail store for user")

// mail bounce handler
// paramters are (recipiant email address, from email address, the filepath of the message, network related error or nil for regular bounce)
type Bouncer func(string, string, string, error)

type connection struct {
	addr   string
	cl     *smtp.Client
	access sync.RWMutex
}

// safely visit a conection with a function
func (c *connection) visit(visitor func(*smtp.Client) error) (err error) {
	c.access.Lock()
	err = visitor(c.cl)
	c.access.Unlock()
	return
}

func (c *connection) Quit() {
	c.cl.Quit()
}

// mail delivery
type Mailer struct {
	// get mail storage for local user
	Local mailstore.MailRouter
	// a dial function to obtain outbound smtp client
	Dial Dialer
	// number of times to try to deliver mail
	// 0 for unlimited
	Retries int
	// called if mail bounces or fails to deliver
	Bounce Bouncer
	// domain resolver function
	Resolve Resolver
	// delivery success hook, called with (recipiant email address, from email address)
	Success func(string, string)
	// for pipelining
	conns map[string]*connection
	// mutex for conns
	cmtx sync.RWMutex
}

// create a new pooled mailer
func NewMailer() *Mailer {
	return &Mailer{
		conns: make(map[string]*connection),
	}
}

// call a visitor for each connection in connection poll
func (s *Mailer) foreach(visitor func(*smtp.Client) error) {
	var conns []*connection
	s.cmtx.Lock()
	for _, conn := range s.conns {
		conns = append(conns, conn)
	}
	s.cmtx.Unlock()
	for _, conn := range conns {
		// in paralell
		go func(c *connection, v func(*smtp.Client) error) {
			err := c.visit(v)
			if err != nil {
				// remove connection on error
				s.delConn(c.addr)
			}
		}(conn, visitor)
	}
}

// delete connection from pool
func (s *Mailer) delConn(addr string) {
	s.cmtx.Lock()
	// safe delete
	c, ok := s.conns[addr]
	if ok {
		c.Quit()
		delete(s.conns, addr)
	}
	s.cmtx.Unlock()
}

// visit a pooled connection in a safe way
func (s *Mailer) visitConn(n, addr, localname string, dialer Dialer, visitor func(*smtp.Client) error) (err error) {
	var c *connection
	c, err = s.getConn(n, addr, localname, dialer)
	if err == nil {
		err = c.visit(visitor)
		if err != nil {
			log.Errorf("failed to visit connection: %s", err.Error())
			s.delConn(addr)
		}
	}
	return
}

// get connection from pool, create if not there
func (s *Mailer) getConn(n, addr, localname string, dial Dialer) (sc *connection, err error) {
	s.cmtx.Lock()
	defer s.cmtx.Unlock()
	var ok bool
	sc, ok = s.conns[addr]
	if !ok {
		var c net.Conn
		c, err = dial(n, addr)
		if err == nil {
			// new connection
			var cl *smtp.Client
			cl, err = smtp.NewClient(c, addr)
			if err != nil {
				log.Errorf("failed to dial: %s", err.Error())
				return
			}
			err = cl.Hello(localname)
			if err != nil {
				log.Errorf("failed to helo: %s", err.Error())
				cl.Quit()
				return
			}
			sc = &connection{
				cl:   cl,
				addr: addr,
			}
			// success
			s.conns[addr] = sc
		}
	}
	return
}

// try delivering mail
// returns a DeliveryJob that can be cancelled
func (s *Mailer) Deliver(recip, from string, msg mailstore.Message) (d DeliverJob) {
	log.Infof("Delivering %s to %s from %s", msg.Filepath(), recip, from)
	dialer := s.Dial
	if dialer == nil {
		dialer = net.Dial
	}

	bounce := s.Bounce

	resolver := s.Resolve

	if resolver == nil {
		resolver = func(name string) (a net.Addr, err error) {
			log.Debugf("mx lookup for %s", name)
			var mx []*net.MX
			mx, err = net.LookupMX(name)
			if err != nil && mx != nil {
				for _, m := range mx {
					var ips []net.IP
					log.Debugf("resolve mx record %s", m.Host)
					ips, err = net.LookupIP(m.Host)
					if err == nil {
						for _, ip := range ips {
							a, err = net.ResolveIPAddr("ip", ip.String())
							if err == nil {
								log.Debugf("resolved %s to %s", name, a)
								return
							}
						}
					}
				}
			}
			return
		}
	}
	var st mailstore.Store
	if s.Local != nil {
		st, _ = s.Local.FindStoreFor(recip)
	}

	if st == nil {
		d = &RemoteDeliverJob{
			unlimited: s.Retries == 0,
			cancel:    false,
			retries:   s.Retries,
			visit: func(f func(*smtp.Client) error) error {
				parts := strings.Split(recip, "@")
				if len(parts) == 2 {
					r_addr := parts[1]
					a, err := resolver(r_addr)
					if err == nil {
						err = s.visitConn(a.Network(), a.String(), r_addr, dialer, f)
					} else {
						log.Warnf("failed to resolve %s: %s", r_addr, err.Error())
					}
					return err
				} else {
					log.Warnf("bad email address %s", recip)
				}
				return nil
			},
			bounce:    bounce,
			recip:     recip,
			from:      from,
			fpath:     msg.Filepath(),
			result:    make(chan bool),
			delivered: s.Success,
		}
	} else {
		d = &LocalDeliverJob{
			st:     st,
			result: make(chan bool),
			fpath:  msg.Filepath(),
		}
	}
	return
}

// gracefully quit all polled connections and close down
func (m *Mailer) Quit() {
	log.Info("shutting down pooled mailer")
	m.cmtx.Lock()
	// collect connections
	var conns []*connection
	for _, c := range m.conns {
		conns = append(conns, c)
	}
	// close all connections
	for _, c := range conns {
		c.visit(func(cl *smtp.Client) error {
			return cl.Quit()
		})
	}
	// collect all names
	var names []string
	for n, _ := range m.conns {
		names = append(names, n)
	}
	// delete all connections
	for _, n := range names {
		delete(m.conns, n)
	}
	m.cmtx.Unlock()
}
<|endoftext|>


<commit_before><commit_msg>lib/raft: Remove double unlock<commit_after><|endoftext|>


<commit_before>package main

import (
	// "encoding/hex"
	"fmt"
	"os/exec"
	"bytes"
	// "time"
	// "math/big"
	// "github.com/PointCoin/btcutil"
	// "github.com/PointCoin/btcwire"
	// "github.com/PointCoin/btcrpcclient"
	// "github.com/PointCoin/btcjson"
	"strconv"
	"strings"
	"encoding/json"
	// "regexp"
	// "math/rand"
	"log"
)

func main() {
	print := fmt.Println
	s := "\"Hello\""
	fmt.Println(s)

	o,_ := strconv.Unquote(s)
	fmt.Println(o)

	// command = pointctl getrawtransaction d2011b19dea6e98ec8bf78bd224856e76b6a9c460bbb347e49adb3dcf457e548
	txid := "d2011b19dea6e98ec8bf78bd224856e76b6a9c460bbb347e49adb3dcf457e548"
	txdetails := getTransactionDetails(txid)
	txdetailsbytes := []byte(txdetails)

	var f interface{}
	_ = json.Unmarshal(txdetailsbytes, &f)
	m := f.(map[string]interface{})

	// for k, v := range m {
	//     switch vv := v.(type) {
	// 	    case string:
	// 	        fmt.Println(k, "is string", vv)
	// 	    case int:
	// 	        fmt.Println(k, "is int", vv)
	// 	    case []interface{}:
	// 	        fmt.Println(k, "is an array:")
	// 	        for i, u := range vv {
	// 	            fmt.Println(i, u)
	// 	        }
	// 	    default:
	// 	        fmt.Println(k, "is of a type I don't know how to handle")
	// 	    }
	// }

	txidreturned  := m["txid"]

	vinList := make([]vin,0)
	vinJsonList := m["vin"]

	switch vv := vinJsonList.(type) {
	case []interface{}:
		for i, u := range vv {
		            fmt.Println(i, u)
		        }
	default:
		print("nope didn't work")
	}

	_ = vinList
	print("\n\ngot txid", txidreturned)

}

type vin struct {
    txid string
    vout  int
}

func getTransactionDetails(txhash string) (string){
	// command = pointctl getrawtransaction d2011b19dea6e98ec8bf78bd224856e76b6a9c460bbb347e49adb3dcf457e548
	cmd := exec.Command("pointctl", "getrawtransaction", txhash)

	var out bytes.Buffer
    cmd.Stdout = &out
    err := cmd.Run()
    if err != nil {
    	log.Fatal(err)
    }
    // fmt.Printf("result: %s\n", out)
    // fmt.Println(out.String())

	cmd2 := exec.Command("xargs", "pointctl", "decoderawtransaction")
	cmd2.Stdin = strings.NewReader(out.String())
	var out2 bytes.Buffer
    cmd2.Stdout = &out2
    err2 := cmd2.Run()
    if err2 != nil {
    	log.Fatal(err2)
    }

    return out2.String()
    // fmt.Println(out2.String())
}

<commit_msg>did my vins work?<commit_after>package main

import (
	// "encoding/hex"
	"fmt"
	"os/exec"
	"bytes"
	// "time"
	// "math/big"
	// "github.com/PointCoin/btcutil"
	// "github.com/PointCoin/btcwire"
	// "github.com/PointCoin/btcrpcclient"
	// "github.com/PointCoin/btcjson"
	"strconv"
	"strings"
	"encoding/json"
	// "regexp"
	// "math/rand"
	"log"
)

func main() {
	print := fmt.Println
	s := "\"Hello\""
	fmt.Println(s)

	o,_ := strconv.Unquote(s)
	fmt.Println(o)

	// command = pointctl getrawtransaction d2011b19dea6e98ec8bf78bd224856e76b6a9c460bbb347e49adb3dcf457e548
	txid := "d2011b19dea6e98ec8bf78bd224856e76b6a9c460bbb347e49adb3dcf457e548"
	txdetails := getTransactionDetails(txid)
	txdetailsbytes := []byte(txdetails)

	var f interface{}
	_ = json.Unmarshal(txdetailsbytes, &f)
	m := f.(map[string]interface{})

	// for k, v := range m {
	//     switch vv := v.(type) {
	// 	    case string:
	// 	        fmt.Println(k, "is string", vv)
	// 	    case int:
	// 	        fmt.Println(k, "is int", vv)
	// 	    case []interface{}:
	// 	        fmt.Println(k, "is an array:")
	// 	        for i, u := range vv {
	// 	            fmt.Println(i, u)
	// 	        }
	// 	    default:
	// 	        fmt.Println(k, "is of a type I don't know how to handle")
	// 	    }
	// }

	txidreturned  := m["txid"]

	vinList := make([]vin,0)
	vinJsonList := m["vin"]

	switch vv := vinJsonList.(type) {
	case []interface{}:
		for _, u := range vv {
					j := u.(map[string]interface{})
					vinTxid := j["txid"]
					vinVout := j["vout"]
					newVin := vin{txid: vinTxid.(string), vout: vinVout.(int)}
					vinList = append(vinList, newVin)
		            // fmt.Println(i, u)
		        }
		// print("yes matches")
	default:
		print("nope didn't work")
	}

	for x := range vinList {
		print(x)
	}

	_ = vinList
	print("\n\ngot txid", txidreturned)

}

type vin struct {
    txid string
    vout  int
}

func getTransactionDetails(txhash string) (string){
	// command = pointctl getrawtransaction d2011b19dea6e98ec8bf78bd224856e76b6a9c460bbb347e49adb3dcf457e548
	cmd := exec.Command("pointctl", "getrawtransaction", txhash)

	var out bytes.Buffer
    cmd.Stdout = &out
    err := cmd.Run()
    if err != nil {
    	log.Fatal(err)
    }
    // fmt.Printf("result: %s\n", out)
    // fmt.Println(out.String())

	cmd2 := exec.Command("xargs", "pointctl", "decoderawtransaction")
	cmd2.Stdin = strings.NewReader(out.String())
	var out2 bytes.Buffer
    cmd2.Stdout = &out2
    err2 := cmd2.Run()
    if err2 != nil {
    	log.Fatal(err2)
    }

    return out2.String()
    // fmt.Println(out2.String())
}

<|endoftext|>


<commit_before>// Arpapp
// dRbiG
// See LICENSE.txt

package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"time"
)

type ArpEntry struct {
	online bool
	stamp  time.Time
	name   string
}

const (
	ARPREGEX  = ".*? \\((.*?)\\) "
	NAMEREGEX = "^(.*?)\\s+A\\s+(\\d+\\.\\d+\\.\\d+\\.\\d+)"
	LOGSIZE   = 64
	HTMLSTART = `<title>arpapp</title>

arpapp

`
	HTMLEND = `</pre></body></html>`
)

var (
	host     string
	port     int
	interval time.Duration
	decay    int
	names    map[string]string
	arplog   map[string]ArpEntry
	arpregex *regexp.Regexp
	decayrex *regexp.Regexp
)

func die(msg string, code int) {
	log.Fatalln(msg)
	os.Exit(code)
}

func setup() {
	var err error
	var intervalstr string
	var decayrexstr string
	var namedb string
	var namerex *regexp.Regexp

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] PORT\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.StringVar(&host, "h", "127.0.0.1", "HTTP server bind HOST")
	flag.StringVar(&intervalstr, "i", "10m", "Scan INTERVAL")
	flag.IntVar(&decay, "d", 72, "Forget after DECAY hours")
	flag.StringVar(&decayrexstr, "r", "", "Decay only IPs matching REGEXP")
	flag.StringVar(&namedb, "n", "", "Hostnames db file path")
	flag.Parse()

	if len(flag.Args()) < 1 {
		die("You have to specify a port!", 2)
	}

	interval, err = time.ParseDuration(intervalstr)
	if err != nil {
		die("Couldn't parse interval!", 2)
	}

	port, err = strconv.Atoi(flag.Args()[0])
	if err != nil || port <= 0 {
		die("Port has to be a positive integer!", 2)
	}

	if (len(decayrexstr) > 0) && (decay > 0) {
		decayrex, err = regexp.Compile(decayrexstr)
		if err != nil {
			die("Failed to compile regexp '"+decayrexstr+"'!", 2)
		}
	} else {
		decayrex = nil
	}

	if len(namedb) > 0 {
		namerex = regexp.MustCompile(NAMEREGEX)
		names = make(map[string]string, LOGSIZE)
		handle, err := os.Open(namedb)
		if err != nil {
			die("Failed to open hostnames file '"+namedb+"'!", 2)
		}
		defer handle.Close()
		scn := bufio.NewScanner(handle)
		for scn.Scan() {
			if match := namerex.FindStringSubmatch(scn.Text()); match != nil {
				names[match[2]] = match[1]
			}
		}
	}

	arpregex = regexp.MustCompile(ARPREGEX)
	arplog = make(map[string]ArpEntry, LOGSIZE)
}

func scan() {
	data, err := exec.Command("arp", "-a").Output()
	if err != nil {
		log.Println("ERROR: running 'arp -a'")
		return
	}

	matches := arpregex.FindAllStringSubmatch(string(data), -1)
	seen := make(map[string]bool, len(matches))
	for i := 0; i < len(matches); i++ {
		ip := matches[i][1]
		if exec.Command("ping", "-c1", "-W1", ip).Run() == nil {
			seen[ip] = true
		} else {
			continue
		}
	}
	for ip, _ := range seen {
		if _, present := arplog[ip]; present {
			if !arplog[ip].online {
				arplog[ip] = ArpEntry{online: true, stamp: time.Now(), name: names[ip]}
			}
		} else {
			arplog[ip] = ArpEntry{online: true, stamp: time.Now(), name: names[ip]}
		}
	}
	for ip, entry := range arplog {
		if _, present := seen[ip]; !present {
			if entry.online {
				arplog[ip] = ArpEntry{online: false, stamp: time.Now(), name: names[ip]}
			}
		}
		if (decayrex != nil) && !entry.online {
			if decayrex.MatchString(ip) {
				duration := time.Since(entry.stamp)
				if duration.Hours() >= float64(decay) {
					delete(arplog, ip)
				}
			}
		}
	}
}

func render(out http.ResponseWriter, req *http.Request) {
	keys := make([]string, len(arplog))
	i := 0
	for ip, _ := range arplog {
		keys[i] = ip
		i++
	}
	sort.Strings(keys)

	fmt.Fprintf(out, HTMLSTART)
	for _, ip := range keys {
		entry := arplog[ip]
		duration := time.Since(entry.stamp)

		fmt.Fprintf(out, "%-16s %-15s since %s (%s)\n",
			entry.name, ip, entry.stamp.Format("2006-01-02 15:04:05 MST"), duration.String())
	}
	fmt.Fprintf(out, "\ninterval: %s", interval)
	fmt.Fprintf(out, HTMLEND)
}

func main() {
	setup()

	http.HandleFunc("/", render)
	go http.ListenAndServe(fmt.Sprintf("%s:%d", host, port), nil)

	for {
		scan()
		time.Sleep(interval)
	}
}
<commit_msg>arpapp.go: explicitly print online/offline status<commit_after>// Arpapp
// dRbiG
// See LICENSE.txt

package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"time"
)

type ArpEntry struct {
	online bool
	stamp  time.Time
	name   string
}

const (
	ARPREGEX  = ".*? \\((.*?)\\) "
	NAMEREGEX = "^(.*?)\\s+A\\s+(\\d+\\.\\d+\\.\\d+\\.\\d+)"
	LOGSIZE   = 64
	HTMLSTART = `<title>arpapp</title>

arpapp

`
	HTMLEND = `</pre></body></html>`
)

var (
	host     string
	port     int
	interval time.Duration
	decay    int
	names    map[string]string
	arplog   map[string]ArpEntry
	arpregex *regexp.Regexp
	decayrex *regexp.Regexp
)

func die(msg string, code int) {
	log.Fatalln(msg)
	os.Exit(code)
}

func setup() {
	var err error
	var intervalstr string
	var decayrexstr string
	var namedb string
	var namerex *regexp.Regexp

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] PORT\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.StringVar(&host, "h", "127.0.0.1", "HTTP server bind HOST")
	flag.StringVar(&intervalstr, "i", "10m", "Scan INTERVAL")
	flag.IntVar(&decay, "d", 72, "Forget after DECAY hours")
	flag.StringVar(&decayrexstr, "r", "", "Decay only IPs matching REGEXP")
	flag.StringVar(&namedb, "n", "", "Hostnames db file path")
	flag.Parse()

	if len(flag.Args()) < 1 {
		die("You have to specify a port!", 2)
	}

	interval, err = time.ParseDuration(intervalstr)
	if err != nil {
		die("Couldn't parse interval!", 2)
	}

	port, err = strconv.Atoi(flag.Args()[0])
	if err != nil || port <= 0 {
		die("Port has to be a positive integer!", 2)
	}

	if (len(decayrexstr) > 0) && (decay > 0) {
		decayrex, err = regexp.Compile(decayrexstr)
		if err != nil {
			die("Failed to compile regexp '"+decayrexstr+"'!", 2)
		}
	} else {
		decayrex = nil
	}

	if len(namedb) > 0 {
		namerex = regexp.MustCompile(NAMEREGEX)
		names = make(map[string]string, LOGSIZE)
		handle, err := os.Open(namedb)
		if err != nil {
			die("Failed to open hostnames file '"+namedb+"'!", 2)
		}
		defer handle.Close()
		scn := bufio.NewScanner(handle)
		for scn.Scan() {
			if match := namerex.FindStringSubmatch(scn.Text()); match != nil {
				names[match[2]] = match[1]
			}
		}
	}

	arpregex = regexp.MustCompile(ARPREGEX)
	arplog = make(map[string]ArpEntry, LOGSIZE)
}

func scan() {
	data, err := exec.Command("arp", "-a").Output()
	if err != nil {
		log.Println("ERROR: running 'arp -a'")
		return
	}

	matches := arpregex.FindAllStringSubmatch(string(data), -1)
	seen := make(map[string]bool, len(matches))
	for i := 0; i < len(matches); i++ {
		ip := matches[i][1]
		if exec.Command("ping", "-c1", "-W1", ip).Run() == nil {
			seen[ip] = true
		} else {
			continue
		}
	}
	for ip, _ := range seen {
		if _, present := arplog[ip]; present {
			if !arplog[ip].online {
				arplog[ip] = ArpEntry{online: true, stamp: time.Now(), name: names[ip]}
			}
		} else {
			arplog[ip] = ArpEntry{online: true, stamp: time.Now(), name: names[ip]}
		}
	}
	for ip, entry := range arplog {
		if _, present := seen[ip]; !present {
			if entry.online {
				arplog[ip] = ArpEntry{online: false, stamp: time.Now(), name: names[ip]}
			}
		}
		if (decayrex != nil) && !entry.online {
			if decayrex.MatchString(ip) {
				duration := time.Since(entry.stamp)
				if duration.Hours() >= float64(decay) {
					delete(arplog, ip)
				}
			}
		}
	}
}

func render(out http.ResponseWriter, req *http.Request) {
	keys := make([]string, len(arplog))
	i := 0
	for ip, _ := range arplog {
		keys[i] = ip
		i++
	}
	sort.Strings(keys)

	fmt.Fprintf(out, HTMLSTART)
	for _, ip := range keys {
		entry := arplog[ip]
		duration := time.Since(entry.stamp)
		var state string

		fmt.Fprintf(out, "%-16s %-15s %s since %s (%s)\n",
			entry.name, ip, state, entry.stamp.Format("2006-01-02 15:04:05 MST"), duration.String())
	}
	fmt.Fprintf(out, "\ninterval: %s", interval)
	fmt.Fprintf(out, HTMLEND)
}

func main() {
	setup()

	http.HandleFunc("/", render)
	go http.ListenAndServe(fmt.Sprintf("%s:%d", host, port), nil)

	for {
		scan()
		time.Sleep(interval)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 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 remote

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"strings"
	"sync"

	"k8s.io/klog/v2"
)

var sshOptions = flag.String("ssh-options", "", "Commandline options passed to ssh.")
var sshEnv = flag.String("ssh-env", "", "Use predefined ssh options for environment.  Options: gce")
var sshKey = flag.String("ssh-key", "", "Path to ssh private key.")
var sshUser = flag.String("ssh-user", "", "Use predefined user for ssh.")

var sshOptionsMap map[string]string
var sshDefaultKeyMap map[string]string

func init() {
	usr, err := user.Current()
	if err != nil {
		klog.Fatal(err)
	}
	sshOptionsMap = map[string]string{
		"gce": "-o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -o CheckHostIP=no -o StrictHostKeyChecking=no -o ServerAliveInterval=30 -o LogLevel=ERROR",
	}
	defaultGceKey := os.Getenv("GCE_SSH_PRIVATE_KEY_FILE")
	if defaultGceKey == "" {
		defaultGceKey = fmt.Sprintf("%s/.ssh/google_compute_engine", usr.HomeDir)
	}
	sshDefaultKeyMap = map[string]string{
		"gce": defaultGceKey,
	}
}

var hostnameIPOverrides = struct {
	sync.RWMutex
	m map[string]string
}{m: make(map[string]string)}

// AddHostnameIP adds <hostname,ip> pair into hostnameIPOverrides map.
func AddHostnameIP(hostname, ip string) {
	hostnameIPOverrides.Lock()
	defer hostnameIPOverrides.Unlock()
	hostnameIPOverrides.m[hostname] = ip
}

// GetHostnameOrIP converts hostname into ip and apply user if necessary.
func GetHostnameOrIP(hostname string) string {
	hostnameIPOverrides.RLock()
	defer hostnameIPOverrides.RUnlock()
	host := hostname
	if ip, found := hostnameIPOverrides.m[hostname]; found {
		host = ip
	}

	if *sshUser == "" {
		*sshUser = os.Getenv("KUBE_SSH_USER")
	}

	if *sshUser != "" {
		host = fmt.Sprintf("%s@%s", *sshUser, host)
	}
	return host
}

// getSSHCommand handles proper quoting so that multiple commands are executed in the same shell over ssh
func getSSHCommand(sep string, args ...string) string {
	return fmt.Sprintf("'%s'", strings.Join(args, sep))
}

// SSH executes ssh command with runSSHCommand as root. The `sudo` makes sure that all commands
// are executed by root, so that there won't be permission mismatch between different commands.
func SSH(host string, cmd ...string) (string, error) {
	return runSSHCommand("ssh", append([]string{GetHostnameOrIP(host), "--", "sudo"}, cmd...)...)
}

// SSHNoSudo executes ssh command with runSSHCommand as normal user. Sometimes we need this,
// for example creating a directory that we'll copy files there with scp.
func SSHNoSudo(host string, cmd ...string) (string, error) {
	return runSSHCommand("ssh", append([]string{GetHostnameOrIP(host), "--"}, cmd...)...)
}

// runSSHCommand executes the ssh or scp command, adding the flag provided --ssh-options
func runSSHCommand(cmd string, args ...string) (string, error) {
	if key := getPrivateSSHKey(); len(key) != 0 {
		if _, err := os.Stat(key); err != nil {
			klog.Errorf("private SSH key (%s) not found. Check if the SSH key is configured properly:, err: %v", key, err)
			return "", fmt.Errorf("private SSH key (%s) does not exist", key)
		}

		args = append([]string{"-i", key}, args...)
	}
	if env, found := sshOptionsMap[*sshEnv]; found {
		args = append(strings.Split(env, " "), args...)
	}
	if *sshOptions != "" {
		args = append(strings.Split(*sshOptions, " "), args...)
	}
	klog.Infof("Running the command %s, with args: %v", cmd, args)
	output, err := exec.Command(cmd, args...).CombinedOutput()
	if err != nil {
		klog.Errorf("failed to run SSH command: out: %s, err: %v", output, err)
		return string(output), fmt.Errorf("command [%s %s] failed with error: %v", cmd, strings.Join(args, " "), err)
	}
	return string(output), nil
}

// getPrivateSSHKey returns the path to ssh private key
func getPrivateSSHKey() string {
	if *sshKey != "" {
		return *sshKey
	}

	if key, found := sshDefaultKeyMap[*sshEnv]; found {
		return key
	}

	return ""
}
<commit_msg>Check key existence in the getPrivateSSHKey function<commit_after>/*
Copyright 2016 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 remote

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"strings"
	"sync"

	"k8s.io/klog/v2"
)

var sshOptions = flag.String("ssh-options", "", "Commandline options passed to ssh.")
var sshEnv = flag.String("ssh-env", "", "Use predefined ssh options for environment.  Options: gce")
var sshKey = flag.String("ssh-key", "", "Path to ssh private key.")
var sshUser = flag.String("ssh-user", "", "Use predefined user for ssh.")

var sshOptionsMap map[string]string
var sshDefaultKeyMap map[string]string

func init() {
	usr, err := user.Current()
	if err != nil {
		klog.Fatal(err)
	}
	sshOptionsMap = map[string]string{
		"gce": "-o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -o CheckHostIP=no -o StrictHostKeyChecking=no -o ServerAliveInterval=30 -o LogLevel=ERROR",
	}
	defaultGceKey := os.Getenv("GCE_SSH_PRIVATE_KEY_FILE")
	if defaultGceKey == "" {
		defaultGceKey = fmt.Sprintf("%s/.ssh/google_compute_engine", usr.HomeDir)
	}
	sshDefaultKeyMap = map[string]string{
		"gce": defaultGceKey,
	}
}

var hostnameIPOverrides = struct {
	sync.RWMutex
	m map[string]string
}{m: make(map[string]string)}

// AddHostnameIP adds <hostname,ip> pair into hostnameIPOverrides map.
func AddHostnameIP(hostname, ip string) {
	hostnameIPOverrides.Lock()
	defer hostnameIPOverrides.Unlock()
	hostnameIPOverrides.m[hostname] = ip
}

// GetHostnameOrIP converts hostname into ip and apply user if necessary.
func GetHostnameOrIP(hostname string) string {
	hostnameIPOverrides.RLock()
	defer hostnameIPOverrides.RUnlock()
	host := hostname
	if ip, found := hostnameIPOverrides.m[hostname]; found {
		host = ip
	}

	if *sshUser == "" {
		*sshUser = os.Getenv("KUBE_SSH_USER")
	}

	if *sshUser != "" {
		host = fmt.Sprintf("%s@%s", *sshUser, host)
	}
	return host
}

// getSSHCommand handles proper quoting so that multiple commands are executed in the same shell over ssh
func getSSHCommand(sep string, args ...string) string {
	return fmt.Sprintf("'%s'", strings.Join(args, sep))
}

// SSH executes ssh command with runSSHCommand as root. The `sudo` makes sure that all commands
// are executed by root, so that there won't be permission mismatch between different commands.
func SSH(host string, cmd ...string) (string, error) {
	return runSSHCommand("ssh", append([]string{GetHostnameOrIP(host), "--", "sudo"}, cmd...)...)
}

// SSHNoSudo executes ssh command with runSSHCommand as normal user. Sometimes we need this,
// for example creating a directory that we'll copy files there with scp.
func SSHNoSudo(host string, cmd ...string) (string, error) {
	return runSSHCommand("ssh", append([]string{GetHostnameOrIP(host), "--"}, cmd...)...)
}

// runSSHCommand executes the ssh or scp command, adding the flag provided --ssh-options
func runSSHCommand(cmd string, args ...string) (string, error) {
	if key, err := getPrivateSSHKey(); len(key) != 0 {
		if err != nil {
			klog.Errorf("private SSH key (%s) not found. Check if the SSH key is configured properly:, err: %v", key, err)
			return "", fmt.Errorf("private SSH key (%s) does not exist", key)
		}

		args = append([]string{"-i", key}, args...)
	}
	if env, found := sshOptionsMap[*sshEnv]; found {
		args = append(strings.Split(env, " "), args...)
	}
	if *sshOptions != "" {
		args = append(strings.Split(*sshOptions, " "), args...)
	}
	klog.Infof("Running the command %s, with args: %v", cmd, args)
	output, err := exec.Command(cmd, args...).CombinedOutput()
	if err != nil {
		klog.Errorf("failed to run SSH command: out: %s, err: %v", output, err)
		return string(output), fmt.Errorf("command [%s %s] failed with error: %v", cmd, strings.Join(args, " "), err)
	}
	return string(output), nil
}

// getPrivateSSHKey returns the path to ssh private key
func getPrivateSSHKey() (string, error) {
	if *sshKey != "" {
		if _, err := os.Stat(*sshKey); err != nil {
			return *sshKey, err
		}

		return *sshKey, nil
	}

	if key, found := sshDefaultKeyMap[*sshEnv]; found {
		if _, err := os.Stat(key); err != nil {
			return key, err
		}

		return key, nil
	}

	return "", nil
}
<|endoftext|>