<commit_before>// Copyright 2021 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 configs_test

import (
	"path/filepath"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	"google.golang.org/protobuf/encoding/prototext"
	"google.golang.org/protobuf/testing/protocmp"
	"bitbucket.org/creachadair/stringset"
	cpb "github.com/google/localtoast/scannerlib/proto/compliance_go_proto"
	spb "github.com/google/localtoast/scannerlib/proto/compliance_go_proto"
	apb "github.com/google/localtoast/scannerlib/proto/api_go_proto"
	sipb "github.com/google/localtoast/scannerlib/proto/scan_instructions_go_proto"
)

var configFileNames = stringset.New(
	"vm_image_scanning.textproto", "container_image_scanning.textproto", "instance_scanning.textproto")

// Check if the serialized config uses the new reduced per-OS format.
// The test will skip these so that the migration to the reduced format can
// happen in several steps.
// TODO(b/200859368): Remove the usage of this function and update the tests
// once all configs have been migrated.
func isReducedConfigFile(configBytes []byte) bool {
	config := &apb.PerOsBenchmarkConfig{}
	err := prototext.Unmarshal(configBytes, config)
	return err == nil
}

// Validate behavior across all the configs.
func TestRequiredAttributes(t *testing.T) {
	fullNoteIDmap := make(map[string]*cpb.ComplianceNote)
	cpeVersionMap := make(map[string]*stringset.Set)
	for filePath, configBytes := range scanConfigs {
		if isReducedConfigFile(configBytes) {
			continue
		}
		config := &apb.ScanConfig{}
		if err := prototext.Unmarshal(configBytes, config); err != nil {
			t.Errorf("error reading %s: %v", filePath, err)
		}
		for _, b := range config.GetBenchmarkConfigs() {
			noteID := b.GetId()
			note := b.GetComplianceNote()
			if note.GetTitle() == "" {
				t.Errorf(`%s GetTitle(): got "", want non-empty string`, noteID)
			}
			if note.GetDescription() == "" {
				t.Errorf(`%s GetDescription(): got "", want non-empty string`, noteID)
			}
			if note.GetRationale() == "" {
				t.Errorf(`%s GetRationale(): got "", want non-empty string`, noteID)
			}
			if note.GetRemediation() == "" {
				t.Errorf(`%s GetRemediation(): got "", want non-empty string`, noteID)
			}
			if len(note.GetVersion()) == 0 {
				t.Errorf(`%s len(GetVersion()): got 0, want at least 1`, noteID)
			}
			for _, version := range note.GetVersion() {
				if version.GetVersion() == "" {
					t.Errorf(`%s version.GetVersion(): got "", want non-empty string`, noteID)
				}
				if version.GetCpeUri() != "fallback" && version.GetCpeUri() != "cpe:/distribution_independent_linux" && len(strings.Split(version.GetCpeUri(), ":")) != 5 {
					t.Errorf("%s version.getCpeUri(): got %s, want a valid CPE", noteID, version.GetCpeUri())
				}
			}
			if note.GetCisBenchmark().GetProfileLevel() < 1 || note.GetCisBenchmark().GetProfileLevel() > 3 {
				t.Errorf("%s GetCisBenchmark().GetProfilelevel(): got %d, want 1 <= level <= 3", noteID, note.GetCisBenchmark().GetProfileLevel())
			}
			if note.GetCisBenchmark().GetSeverity() == spb.Severity_SEVERITY_UNSPECIFIED {
				t.Errorf("%s GetCisBenchmark().GetSeverity(): got %s, want any specified severity", noteID, note.GetCisBenchmark().GetSeverity())
			}
			scanInstructions := &sipb.BenchmarkScanInstruction{}
			if err := prototext.Unmarshal(note.GetScanInstructions(), scanInstructions); err != nil {
				t.Errorf("%s could not parse scan instructions: %v", noteID, err)
			}
			// Ensure that each note ID has an identical note, excluding version number, and scan instructions.
			// Image and instance scanning notes will check for the same thing, so they have mostly the same
			// information (note ID, description, etc...).
			// However, image notes have shorter scan instructions because we can't scan for everything that
			// instance scanning does.
			if val, ok := fullNoteIDmap[noteID]; ok {
				if diff := cmp.Diff(val, note, protocmp.Transform(), protocmp.IgnoreFields(&cpb.ComplianceNote{}, "version", "scan_instructions")); diff != "" {
					t.Errorf("%s checking note uniqueness: got %v, wanted no diff", noteID, diff)
				}
			} else {
				fullNoteIDmap[noteID] = note
			}
			for _, noteVer := range note.GetVersion() {
				if val, ok := cpeVersionMap[noteVer.GetCpeUri()]; ok {
					val.Add(noteVer.GetVersion())
				} else {
					s := stringset.New(noteVer.GetVersion())
					cpeVersionMap[noteVer.GetCpeUri()] = &s
				}
			}
		}
	}
	for cpe, versions := range cpeVersionMap {
		if versions.Len() > 1 {
			t.Errorf("CPE %s checking number of versions, got %v, wanted a single version", cpe, versions)
		}
	}
}

func TestFilesHaveSupportedName(t *testing.T) {
	for filePath := range scanConfigs {
		fileName := filepath.Base(filePath)
		if !configFileNames.Contains(fileName) {
			t.Errorf("checking file names: got %q, want image_scanning.textproto or instance_scanning.textproto", fileName)
		}
	}
}

func TestScanInstructionsHaveDisplayCommandAndNonComplianceReason(t *testing.T) {
	for filePath, configBytes := range scanConfigs {
		if isReducedConfigFile(configBytes) {
			continue
		}
		config := &apb.ScanConfig{}
		if err := prototext.Unmarshal(configBytes, config); err != nil {
			t.Errorf("error reading %s: %v", filePath, err)
		}
		for _, b := range config.GetBenchmarkConfigs() {
			noteID := b.GetId()
			scanInstructions := &sipb.BenchmarkScanInstruction{}
			if err := prototext.Unmarshal(b.GetComplianceNote().GetScanInstructions(), scanInstructions); err != nil {
				t.Errorf("%s could not parse scan instructions: %v", noteID, err)
			}
			for _, a := range scanInstructions.GetCheckAlternatives() {
				for _, f := range a.GetFileChecks() {
					if len(f.GetFileDisplayCommand()) > 0 && len(f.GetNonComplianceMsg()) == 0 {
						t.Errorf("check for benchmark %s has a file display command set but no non-compliance message", noteID)
					}
				}
			}
		}
	}
}

func TestScanInstructionsSameForFileSets(t *testing.T) {
	for _, fileName := range configFileNames.Elements() {
		scanInstructionForNoteID := make(map[string]*sipb.BenchmarkScanInstruction)
		for filePath, configBytes := range scanConfigs {
			if filepath.Base(filePath) != fileName {
				continue
			}
			if isReducedConfigFile(configBytes) {
				continue
			}
			config := &apb.ScanConfig{}
			if err := prototext.Unmarshal(configBytes, config); err != nil {
				t.Errorf("error reading %s: %v", filePath, err)
			}
			for _, b := range config.GetBenchmarkConfigs() {
				noteID := b.GetId()
				note := b.GetComplianceNote()
				// Ensure that each note ID for images have identical scan instructions.
				scanInstructions := &sipb.BenchmarkScanInstruction{}
				if err := prototext.Unmarshal(note.GetScanInstructions(), scanInstructions); err != nil {
					t.Errorf("%s could not parse scan instructions: %v", noteID, err)
				}
				if val, ok := scanInstructionForNoteID[noteID]; ok {
					if diff := cmp.Diff(val, scanInstructions, protocmp.Transform()); diff != "" {
						t.Errorf("%s checking scan instructions unique: got %v, wanted no diff", noteID, diff)
					}
				} else {
					scanInstructionForNoteID[noteID] = scanInstructions
				}
			}
		}
	}
}

func TestFallbackNotesHaveExpectedIdFormat(t *testing.T) {
	for filePath, configBytes := range scanConfigs {
		if !strings.Contains(filePath, "fallback") {
			continue
		}
		if isReducedConfigFile(configBytes) {
			continue
		}
		config := &apb.ScanConfig{}
		if err := prototext.Unmarshal(configBytes, config); err != nil {
			t.Errorf("error reading %s: %v", filePath, err)
		}
		for _, b := range config.GetBenchmarkConfigs() {
			noteID := b.GetId()
			if !strings.HasSuffix(noteID, "-fallback") {
				t.Errorf("Fallback note ID %q should end with -fallback", noteID)
			}
		}
	}
}
<commit_msg>Change the config test to test the new reduced format.<commit_after>// Copyright 2021 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 configs_test

import (
	"path/filepath"
	"strings"
	"testing"

	"google.golang.org/protobuf/encoding/prototext"
	"bitbucket.org/creachadair/stringset"
	spb "github.com/google/localtoast/scannerlib/proto/compliance_go_proto"
	apb "github.com/google/localtoast/scannerlib/proto/api_go_proto"
	ipb "github.com/google/localtoast/scannerlib/proto/scan_instructions_go_proto"
	sipb "github.com/google/localtoast/scannerlib/proto/scan_instructions_go_proto"
)

var configFileNames = stringset.New(
	"vm_image_scanning.textproto", "container_image_scanning.textproto", "instance_scanning.textproto")

// Check if the serialized config uses the new reduced per-OS format.
// The test will skip these so that the migration to the reduced format can
// happen in several steps.
// TODO(b/200859368): Remove the usage of this function and update the tests
// once all configs have been migrated.
func isReducedConfigFile(configBytes []byte) bool {
	config := &apb.PerOsBenchmarkConfig{}
	err := prototext.Unmarshal(configBytes, config)
	return err == nil
}

// Validate behavior across all the configs.
func TestRequiredAttributes(t *testing.T) {
	for filePath, configBytes := range scanConfigDefs {
		config := &apb.ScanConfig{}
		if err := prototext.Unmarshal(configBytes, config); err != nil {
			t.Errorf("error reading %s: %v", filePath, err)
		}
		for _, b := range config.GetBenchmarkConfigs() {
			noteID := b.GetId()
			note := b.GetComplianceNote()
			if note.GetTitle() == "" {
				t.Errorf(`%s GetTitle(): got "", want non-empty string`, noteID)
			}
			if note.GetDescription() == "" {
				t.Errorf(`%s GetDescription(): got "", want non-empty string`, noteID)
			}
			if note.GetRationale() == "" {
				t.Errorf(`%s GetRationale(): got "", want non-empty string`, noteID)
			}
			if note.GetRemediation() == "" {
				t.Errorf(`%s GetRemediation(): got "", want non-empty string`, noteID)
			}
			if len(note.GetVersion()) == 0 {
				t.Errorf(`%s len(GetVersion()): got 0, want at least 1`, noteID)
			}
			for _, version := range note.GetVersion() {
				if version.GetVersion() == "" {
					t.Errorf(`%s version.GetVersion(): got "", want non-empty string`, noteID)
				}
				if version.GetCpeUri() != "fallback" && version.GetCpeUri() != "cpe:/distribution_independent_linux" && len(strings.Split(version.GetCpeUri(), ":")) != 5 {
					t.Errorf("%s version.getCpeUri(): got %s, want a valid CPE", noteID, version.GetCpeUri())
				}
			}
			if note.GetCisBenchmark().GetProfileLevel() < 1 || note.GetCisBenchmark().GetProfileLevel() > 3 {
				t.Errorf("%s GetCisBenchmark().GetProfilelevel(): got %d, want 1 <= level <= 3", noteID, note.GetCisBenchmark().GetProfileLevel())
			}
			if note.GetCisBenchmark().GetSeverity() == spb.Severity_SEVERITY_UNSPECIFIED {
				t.Errorf("%s GetCisBenchmark().GetSeverity(): got %s, want any specified severity", noteID, note.GetCisBenchmark().GetSeverity())
			}
			scanInstructions := &sipb.BenchmarkScanInstructionDef{}
			if err := prototext.Unmarshal(note.GetScanInstructions(), scanInstructions); err != nil {
				t.Errorf("%s could not parse scan instructions: %v", noteID, err)
			}
		}
	}
}

func TestFilesHaveSupportedName(t *testing.T) {
	for filePath := range reducedScanConfigs {
		fileName := filepath.Base(filePath)
		if !configFileNames.Contains(fileName) {
			t.Errorf("checking file names: got %q, want image_scanning.textproto or instance_scanning.textproto", fileName)
		}
	}
}

func TestScanInstructionsHaveDisplayCommandAndNonComplianceReason(t *testing.T) {
	for filePath, configBytes := range scanConfigDefs {
		config := &apb.ScanConfig{}
		if err := prototext.Unmarshal(configBytes, config); err != nil {
			t.Fatalf("error reading %s: %v", filePath, err)
		}
		for _, b := range config.GetBenchmarkConfigs() {
			noteID := b.GetId()
			instructionDef := &ipb.BenchmarkScanInstructionDef{}
			if err := prototext.Unmarshal(b.GetComplianceNote().GetScanInstructions(), instructionDef); err != nil {
				t.Errorf("%s could not parse scan instructions: %v", noteID, err)
			}
			var scanInstructions []*ipb.BenchmarkScanInstruction
			if instructionDef.GetGeneric() != nil {
				scanInstructions = []*ipb.BenchmarkScanInstruction{instructionDef.GetGeneric()}
			} else if instructionDef.GetScanTypeSpecific() != nil {
				scanInstructions = []*ipb.BenchmarkScanInstruction{
					instructionDef.GetScanTypeSpecific().InstanceScanning,
					instructionDef.GetScanTypeSpecific().ImageScanning,
				}
			} else {
				t.Fatalf("benchmark %s has invalid instruction def %v", noteID, instructionDef)
			}
			for _, i := range scanInstructions {
				for _, a := range i.GetCheckAlternatives() {
					for _, f := range a.GetFileChecks() {
						if len(f.GetFileDisplayCommand()) > 0 && len(f.GetNonComplianceMsg()) == 0 {
							t.Errorf("check for benchmark %s has a file display command set but no non-compliance message: %v", noteID, f)
						}
					}
				}
			}
		}
	}
}

func TestFallbackBenchmarkDefsHaveExpectedIdFormat(t *testing.T) {
	for filePath, configBytes := range scanConfigDefs {
		if filePath != "fallback.textproto" {
			continue
		}
		if isReducedConfigFile(configBytes) {
			continue
		}
		config := &apb.ScanConfig{}
		if err := prototext.Unmarshal(configBytes, config); err != nil {
			t.Errorf("error reading %s: %v", filePath, err)
		}
		for _, b := range config.GetBenchmarkConfigs() {
			noteID := b.GetId()
			if !strings.HasSuffix(noteID, "-fallback") {
				t.Errorf("Fallback benchmark ID %q should end with -fallback", noteID)
			}
		}
	}
}

func TestFallbackPerOsBenchmarksHaveExpectedIdFormat(t *testing.T) {
	for filePath, configBytes := range reducedScanConfigs {
		if !strings.Contains(filePath, "/fallback/") {
			continue
		}
		config := &apb.PerOsBenchmarkConfig{}
		if err := prototext.Unmarshal(configBytes, config); err != nil {
			t.Errorf("error reading %s: %v", filePath, err)
		}
		for _, id := range config.BenchmarkId {
			if !strings.HasSuffix(id, "-fallback") {
				t.Errorf("Fallback benchmark ID %q should end with -fallback", id)
			}
		}
	}
}
<|endoftext|>


<commit_before>package congress

import (
	"fmt"
	"github.com/sunlightlabs/go-sunlight/internal"
)

/**
 * Congress Legislator object. Issued for members of the Congress
 */
type Legislator struct {
	BioguideId  string `json:"bioguide_id"`
	Birthday    string `json:"birthday"`
	Chamber     string `json:"chamber"`
	ContactForm string `json:"contact_form"`
	CrpId       string `json:"crp_id"`
	District    int    `json:"district"`
	FacebookId  string `json:"facebook_id"`
	Fax         string `json:"fax"`

	FecIds     []string `json:"fec_ids"`
	FirstName  string   `json:"first_name"`
	MiddleName string   `json:"middle_name"`
	LastName   string   `json:"last_name"`
	Gender     string   `json:"gender"`

	GovtrackId string `json:"govtrack_id"`
	ICPSRId    int    `json:"icpsr_id"`
	InOffice   bool   `json:"in_office"`

	NameSuffix string `json:"name_suffix"`
	Nickname   string `json:"nickname"`
	OcEmail    string `json:"oc_email"`
	Office     string `json:"office"`
	Party      string `json:"party"`

	Phone     string `json:"phone"`
	State     string `json:"state"`
	StateName string `json:"state_name"`
	TermEnd   string `json:"term_end"`
	TermStart string `json:"term_start"`
	ThomasId  string `json:"thomas_id"`
	Title     string `json:"title"`

	TwitterId   string `json:"twitter_id"`
	VotesmartId int    `json:"votesmart_id"`
	Website     string `json:"website"`
	YoutubeId   string `json:"youtube_id"`
}

type LegislatorResult struct {
	Result
	Results []Legislator `json:"results"`
}

func GetLegislators(criteria map[string]string) (*LegislatorResult, error) {
	l := LegislatorResult{}
	err := internal.GetURL(&l, congressRoot, criteria, "legislators")
	if err != nil {
		return nil, err
	}
	return &l, nil
}

func GetLegislatorsByLatLon(lat float32, lon float32) (*LegislatorResult, error) {
	var params = map[string]string{
		"latitude":  fmt.Sprintf("%f", lat),
		"longitude": fmt.Sprintf("%f", lon),
	}

	l := LegislatorResult{}
	err := internal.GetURL(&l, congressRoot, params, "legislators", "locate")
	if err != nil {
		return nil, err
	}
	return &l, nil
}
<commit_msg>Add a GetBills method on the Legislator struct<commit_after>package congress

import (
	"fmt"
	"github.com/sunlightlabs/go-sunlight/internal"
)

/**
 * Congress Legislator object. Issued for members of the Congress
 */
type Legislator struct {
	BioguideId  string `json:"bioguide_id"`
	Birthday    string `json:"birthday"`
	Chamber     string `json:"chamber"`
	ContactForm string `json:"contact_form"`
	CrpId       string `json:"crp_id"`
	District    int    `json:"district"`
	FacebookId  string `json:"facebook_id"`
	Fax         string `json:"fax"`

	FecIds     []string `json:"fec_ids"`
	FirstName  string   `json:"first_name"`
	MiddleName string   `json:"middle_name"`
	LastName   string   `json:"last_name"`
	Gender     string   `json:"gender"`

	GovtrackId string `json:"govtrack_id"`
	ICPSRId    int    `json:"icpsr_id"`
	InOffice   bool   `json:"in_office"`

	NameSuffix string `json:"name_suffix"`
	Nickname   string `json:"nickname"`
	OcEmail    string `json:"oc_email"`
	Office     string `json:"office"`
	Party      string `json:"party"`

	Phone     string `json:"phone"`
	State     string `json:"state"`
	StateName string `json:"state_name"`
	TermEnd   string `json:"term_end"`
	TermStart string `json:"term_start"`
	ThomasId  string `json:"thomas_id"`
	Title     string `json:"title"`

	TwitterId   string `json:"twitter_id"`
	VotesmartId int    `json:"votesmart_id"`
	Website     string `json:"website"`
	YoutubeId   string `json:"youtube_id"`
}

type LegislatorResult struct {
	Result
	Results []Legislator `json:"results"`
}

func (l Legislator) GetBills() (*BillResult, error) {
	bills, err := BillSearch(map[string]string{
		"sponsor_id": l.BioguideId,
	})
	if err != nil {
		return nil, err
	}
	return bills, nil
}

func GetLegislators(criteria map[string]string) (*LegislatorResult, error) {
	l := LegislatorResult{}
	err := internal.GetURL(&l, congressRoot, criteria, "legislators")
	if err != nil {
		return nil, err
	}
	return &l, nil
}

func GetLegislatorsByLatLon(lat float32, lon float32) (*LegislatorResult, error) {
	var params = map[string]string{
		"latitude":  fmt.Sprintf("%f", lat),
		"longitude": fmt.Sprintf("%f", lon),
	}

	l := LegislatorResult{}
	err := internal.GetURL(&l, congressRoot, params, "legislators", "locate")
	if err != nil {
		return nil, err
	}
	return &l, nil
}
<|endoftext|>


<commit_before>package hoverfly

import (
	"github.com/SpectoLabs/hoverfly/core/testutil"
	"net/http"
	"testing"
)

func TestSynthesizeResponse(t *testing.T) {

	req, err := http.NewRequest("GET", "http://example.com", nil)
	testutil.Expect(t, err, nil)

	sr, err := SynthesizeResponse(req, "./examples/middleware/synthetic_service/synthetic.py")
	testutil.Expect(t, err, nil)

	testutil.Expect(t, sr.StatusCode, 200)
}

func TestSynthesizeResponseWOMiddleware(t *testing.T) {

	req, err := http.NewRequest("GET", "http://example.com", nil)
	testutil.Expect(t, err, nil)

	_, err = SynthesizeResponse(req, "")
	testutil.Refute(t, err, nil)

	testutil.Expect(t, err.Error(), "Synthesize failed, middleware not provided")
}

func TestSynthesizeMiddlewareFailure(t *testing.T) {

	req, err := http.NewRequest("GET", "http://example.com", nil)
	testutil.Expect(t, err, nil)

	_, err = SynthesizeResponse(req, "./examples/middleware/this_is_not_there.py")
	testutil.Refute(t, err, nil)
}
<commit_msg>Cleaning up the synthesize_test by replacing testutils with Gomega<commit_after>package hoverfly

import (
	. "github.com/onsi/gomega"
	"net/http"
	"testing"
)

func TestSynthesizeResponse(t *testing.T) {
	RegisterTestingT(t)

	req, err := http.NewRequest("GET", "http://example.com", nil)
	Expect(err).To(BeNil())

	sr, err := SynthesizeResponse(req, "./examples/middleware/synthetic_service/synthetic.py")
	Expect(err).To(BeNil())

	Expect(sr.StatusCode).To(Equal(http.StatusOK))
}

func TestSynthesizeResponseWOMiddleware(t *testing.T) {
	RegisterTestingT(t)

	req, err := http.NewRequest("GET", "http://example.com", nil)
	Expect(err).To(BeNil())

	_, err = SynthesizeResponse(req, "")
	Expect(err).ToNot(BeNil())

	Expect(err).To(MatchError("Synthesize failed, middleware not provided"))
}

func TestSynthesizeMiddlewareFailure(t *testing.T) {
	RegisterTestingT(t)

	req, err := http.NewRequest("GET", "http://example.com", nil)
	Expect(err).To(BeNil())

	_, err = SynthesizeResponse(req, "./examples/middleware/this_is_not_there.py")
	Expect(err).ToNot(BeNil())
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/temoto/robotstxt"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// Get the depth of a file path
func pathDepth(filePath string) int {
	cleanPath := strings.Trim(path.Clean(filePath), "/")

	// Special case: / should have length 0
	if cleanPath == "" {
		return 0
	}

	return len(strings.Split(cleanPath, "/"))
}

func HttpGetUnsafe(reqUrl string) (resp *http.Response, err error) {
	tr := http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}
	client := &http.Client{Transport: &tr}
	return client.Get(reqUrl)
}

type HttpCrawlerConfig struct {
	BodySizeLimit int64         `json:"maxBodySize"`
	Entry         string        `json:"entry"`
	MaxPathDepth  int           `json:"maxPathDepth"`
	RateLimit     time.Duration `json:"maxRequestPerSecond"`
	RobotName     string        `json:"robotName"`
	ObeyRobotsTxt bool          `json:"obeyRobotsTxt"`
}

type HttpCrawler struct {
	Config HttpCrawlerConfig

	Entry *url.URL

	RobotsTestAgent *robotstxt.Group
	Ticker          <-chan time.Time
}

func CreateHttpCrawler(rawConfig *json.RawMessage) (crawler *HttpCrawler, err error) {
	// Create a new instance
	crawler = &HttpCrawler{}

	// Parse config while providing default values
	config := HttpCrawlerConfig{
		BodySizeLimit: 10 * 1000 * 1000, // 1 MB
		MaxPathDepth:  20,
		RateLimit:     0,
		RobotName:     DEFAULT_BOTNAME,
		ObeyRobotsTxt: true,
	}
	err = json.Unmarshal(*rawConfig, &config)
	if err != nil {
		return
	}

	crawler.Config = config

	// Parse the entry URL
	entry, err := url.Parse(crawler.Config.Entry)
	if err != nil {
		return
	}
	crawler.Entry = entry

	// Initialize the RateLimit throttle ticker if needed
	if crawler.Config.RateLimit > 0 {
		crawler.Ticker = time.Tick(crawler.Config.RateLimit)
	}

	// Try to parse /robots.txt
	if crawler.Config.ObeyRobotsTxt {
		robotsURL := url.URL{}
		robotsURL = *entry
		robotsURL.Path = "/robots.txt"

		robotsRes, robotsErr := HttpGetUnsafe(robotsURL.String())
		if robotsErr == nil {
			defer robotsRes.Body.Close()

			var robots *robotstxt.RobotsData
			robots, err = robotstxt.FromResponse(robotsRes)
			if err != nil {
				return
			}

			crawler.RobotsTestAgent = robots.FindGroup(crawler.Config.RobotName)
		}
	}

	return
}

func (crawler *HttpCrawler) walker(entry *url.URL, fn WalkFunction) (err error) {
	entryStr := entry.String()

	// Check if this file is allowed to be crawled by robots.txt rules
	if crawler.RobotsTestAgent != nil && !crawler.RobotsTestAgent.Test(entryStr) {
		return
	}

	// Throttle requests as specified by the RateLimit if needed
	if crawler.Config.RateLimit > 0 {
		<-crawler.Ticker
	}

	// Do a standard HTTP GET request, but only download the first few kilobytes
	// of body data. Reasons:
	// 1. Save a request as otherwise we would have to create one HEAD request
	//    before every GET to check for the content type and content length
	// 2. Work around broken HEAD implementations
	resp, err := HttpGetUnsafe(entryStr)
	if err != nil {
		return
	}

	defer resp.Body.Close()

	// Determine the content length in bytes
	var contentLength int64
	if len(resp.Header.Get("Content-Length")) > 0 {
		contentLength, err = strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 0)
		if err != nil {
			return
		}
	}

	// We only continue walking on Content-Type: text/html files and call the
	// WalkFunction on all other files
	// TODO implement fallback mime sniffing (magic numbers parsing)
	mime, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
	if err != nil {
		return
	}

	if mime != "text/html" {
		var modTime time.Time
		lastModified := resp.Header.Get("Last-Modified")
		if lastModified != "" {
			modTime, err = http.ParseTime(resp.Header.Get("Last-Modified"))
			if err != nil {
				fmt.Println(resp.Header.Get("Last-Modified"))
				return
			}
		}

		fn(entryStr, FileInfo{
			URL:      entry,
			Size:     contentLength,
			MimeType: mime,
			ModTime:  modTime,
		})

		return
	}

	// Limit the amount of downloaded data
	body, err := ioutil.ReadAll(io.LimitReader(resp.Body, crawler.Config.BodySizeLimit))
	if err != nil {
		return
	}

	if int64(len(body)) < contentLength {
		err = fmt.Errorf("BodySizeLimit exceeded")
		return
	}

	// Search for all anchor tags
	tokenizer := html.NewTokenizer(bytes.NewBuffer(body))

	for {
		tt := tokenizer.Next()

		switch tt {
		case html.ErrorToken:
			if tErr := tokenizer.Err(); tErr != io.EOF {
				err = tErr
			}
			return
		case html.StartTagToken:
			token := tokenizer.Token()

			if strings.ToLower(token.Data) == "a" {
				for _, a := range token.Attr {
					if a.Key == "href" {
						var u *url.URL
						u, err = url.Parse(a.Val)
						if err != nil {
							return
						}

						nextUrl := entry.ResolveReference(u)

						// Stop if we are going to leave the server
						if nextUrl.Hostname() != entry.Hostname() {
							break
						}

						// Ignore if the path depth is decreasing (e.g. href is .. )
						if pathDepth(nextUrl.Path) < pathDepth(entry.Path) {
							break
						}

						// Ignore if the url is not changing (e.g. href is . )
						if nextUrl.String() == entry.String() {
							break
						}

						// Stop if we have reached the maximum path depth
						if pathDepth(nextUrl.Path) > crawler.Config.MaxPathDepth {
							err = fmt.Errorf("MaxPathDepth exceeded")
							return
						}

						// Ignore Apache dir list sort links
						var match bool
						match, err = regexp.MatchString("^C=(.*)(&|;)O=(.*)$", nextUrl.RawQuery)
						if match || err != nil {
							break
						}

						// Errors are bubbled up
						err = crawler.walker(nextUrl, fn)
						if err != nil {
							return
						}

						break
					}
				}
			}
		}
	}

	return
}

func (crawler *HttpCrawler) Walk(fn WalkFunction) error {
	return crawler.walker(crawler.Entry, fn)
}

func (crawler *HttpCrawler) Close() {

}
<commit_msg>HTTP: Share an http.Client for each host so we can use keep alive<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/temoto/robotstxt"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// Get the depth of a file path
func pathDepth(filePath string) int {
	cleanPath := strings.Trim(path.Clean(filePath), "/")

	// Special case: / should have length 0
	if cleanPath == "" {
		return 0
	}

	return len(strings.Split(cleanPath, "/"))
}

type HttpCrawlerConfig struct {
	BodySizeLimit int64         `json:"maxBodySize"`
	Entry         string        `json:"entry"`
	MaxPathDepth  int           `json:"maxPathDepth"`
	RateLimit     time.Duration `json:"maxRequestPerSecond"`
	RobotName     string        `json:"robotName"`
	ObeyRobotsTxt bool          `json:"obeyRobotsTxt"`
}

type HttpCrawler struct {
	Config HttpCrawlerConfig

	Entry *url.URL

	RobotsTestAgent *robotstxt.Group
	Ticker          <-chan time.Time
	HttpClient      *http.Client
}

func CreateHttpCrawler(rawConfig *json.RawMessage) (crawler *HttpCrawler, err error) {
	// Create a new instance
	crawler = &HttpCrawler{
		// Share an http.Client so we can keep alive connections
		HttpClient: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					MaxIdleConnsPerHost: 1024,
					InsecureSkipVerify: true,
				},
			},
		},
	}

	// Parse config while providing default values
	config := HttpCrawlerConfig{
		BodySizeLimit: 10 * 1000 * 1000, // 1 MB
		MaxPathDepth:  20,
		RateLimit:     0,
		RobotName:     DEFAULT_BOTNAME,
		ObeyRobotsTxt: true,
	}
	err = json.Unmarshal(*rawConfig, &config)
	if err != nil {
		return
	}

	crawler.Config = config

	// Parse the entry URL
	entry, err := url.Parse(crawler.Config.Entry)
	if err != nil {
		return
	}
	crawler.Entry = entry

	// Initialize the RateLimit throttle ticker if needed
	if crawler.Config.RateLimit > 0 {
		crawler.Ticker = time.Tick(crawler.Config.RateLimit)
	}

	// Try to parse /robots.txt
	if crawler.Config.ObeyRobotsTxt {
		robotsURL := url.URL{}
		robotsURL = *entry
		robotsURL.Path = "/robots.txt"

		robotsRes, robotsErr := crawler.httpGet(robotsURL.String())
		if robotsErr == nil {
			defer robotsRes.Body.Close()

			var robots *robotstxt.RobotsData
			robots, err = robotstxt.FromResponse(robotsRes)
			if err != nil {
				return
			}

			crawler.RobotsTestAgent = robots.FindGroup(crawler.Config.RobotName)
		}
	}

	return
}

func (crawler *HttpCrawler) httpGet(reqUrl string) (resp *http.Response, err error) {
	return crawler.HttpClient.Get(reqUrl)
}

func (crawler *HttpCrawler) walker(entry *url.URL, fn WalkFunction) (err error) {
	entryStr := entry.String()

	// Check if this file is allowed to be crawled by robots.txt rules
	if crawler.RobotsTestAgent != nil && !crawler.RobotsTestAgent.Test(entryStr) {
		return
	}

	// Throttle requests as specified by the RateLimit if needed
	if crawler.Config.RateLimit > 0 {
		<-crawler.Ticker
	}

	// Do a standard HTTP GET request, but only download the first few kilobytes
	// of body data. Reasons:
	// 1. Save a request as otherwise we would have to create one HEAD request
	//    before every GET to check for the content type and content length
	// 2. Work around broken HEAD implementations
	resp, err := crawler.httpGet(entryStr)
	if err != nil {
		return
	}

	defer resp.Body.Close()

	// Determine the content length in bytes
	var contentLength int64
	if len(resp.Header.Get("Content-Length")) > 0 {
		contentLength, err = strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 0)
		if err != nil {
			return
		}
	}

	// We only continue walking on Content-Type: text/html files and call the
	// WalkFunction on all other files
	// TODO implement fallback mime sniffing (magic numbers parsing)
	mime, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
	if err != nil {
		return
	}

	if mime != "text/html" {
		var modTime time.Time
		lastModified := resp.Header.Get("Last-Modified")
		if lastModified != "" {
			modTime, err = http.ParseTime(resp.Header.Get("Last-Modified"))
			if err != nil {
				fmt.Println(resp.Header.Get("Last-Modified"))
				return
			}
		}

		fn(entryStr, FileInfo{
			URL:      entry,
			Size:     contentLength,
			MimeType: mime,
			ModTime:  modTime,
		})

		return
	}

	// Limit the amount of downloaded data
	body, err := ioutil.ReadAll(io.LimitReader(resp.Body, crawler.Config.BodySizeLimit))
	if err != nil {
		return
	}

	if int64(len(body)) < contentLength {
		err = fmt.Errorf("BodySizeLimit exceeded")
		return
	}

	// Search for all anchor tags
	tokenizer := html.NewTokenizer(bytes.NewBuffer(body))

	for {
		tt := tokenizer.Next()

		switch tt {
		case html.ErrorToken:
			if tErr := tokenizer.Err(); tErr != io.EOF {
				err = tErr
			}
			return
		case html.StartTagToken:
			token := tokenizer.Token()

			if strings.ToLower(token.Data) == "a" {
				for _, a := range token.Attr {
					if a.Key == "href" {
						var u *url.URL
						u, err = url.Parse(a.Val)
						if err != nil {
							return
						}

						nextUrl := entry.ResolveReference(u)

						// Stop if we are going to leave the server
						if nextUrl.Hostname() != entry.Hostname() {
							break
						}

						// Ignore if the path depth is decreasing (e.g. href is .. )
						if pathDepth(nextUrl.Path) < pathDepth(entry.Path) {
							break
						}

						// Ignore if the url is not changing (e.g. href is . )
						if nextUrl.String() == entry.String() {
							break
						}

						// Stop if we have reached the maximum path depth
						if pathDepth(nextUrl.Path) > crawler.Config.MaxPathDepth {
							err = fmt.Errorf("MaxPathDepth exceeded")
							return
						}

						// Ignore Apache dir list sort links
						var match bool
						match, err = regexp.MatchString("^C=(.*)(&|;)O=(.*)$", nextUrl.RawQuery)
						if match || err != nil {
							break
						}

						// Errors are bubbled up
						err = crawler.walker(nextUrl, fn)
						if err != nil {
							return
						}

						break
					}
				}
			}
		}
	}

	return
}

func (crawler *HttpCrawler) Walk(fn WalkFunction) error {
	return crawler.walker(crawler.Entry, fn)
}

func (crawler *HttpCrawler) Close() {

}
<|endoftext|>


<commit_before>/*
Copyright 2015 Google Inc. All rights reserved.

Use of this source code is governed by a BSD-style
license that can be found in the LICENSE file or at
https://developers.google.com/open-source/licenses/bsd
*/
package cups

import (
	"fmt"
	"reflect"
	"sort"
	"strconv"
	"strings"

	"github.com/golang/glog"
	"github.com/google/cups-connector/cdd"
)

// translateAttrs extracts a PrinterDescriptionSection, PrinterStateSection, name, default diplay name, UUID, and tags from maps of tags (CUPS attributes)
func translateAttrs(printerTags map[string][]string) (*cdd.PrinterDescriptionSection, *cdd.PrinterStateSection, string, string, string, map[string]string) {
	var name, info string
	if n, ok := printerTags[attrPrinterName]; ok && len(n) > 0 {
		name = n[0]
	}
	if i, ok := printerTags[attrPrinterInfo]; ok && len(i) > 0 {
		info = i[0]
	}
	uuid := getUUID(printerTags)

	var desc = cdd.PrinterDescriptionSection{VendorCapability: &[]cdd.VendorCapability{}}
	var state cdd.PrinterStateSection

	desc.SupportedContentType = convertSupportedContentType(printerTags)
	desc.Marker, state.MarkerState = convertMarkers(printerTags)
	desc.PageOrientation = convertPageOrientation(printerTags)
	desc.Copies = convertCopies(printerTags)
	if vc := convertPagesPerSheet(printerTags); vc != nil {
		*desc.VendorCapability = append(*desc.VendorCapability, *vc)
	}

	state.State = getState(printerTags)
	state.VendorState = getVendorState(printerTags)

	tags := make(map[string]string, len(printerTags))
	for k, v := range printerTags {
		tags[k] = strings.Join(v, ",")
	}

	return &desc, &state, name, info, uuid, tags
}

func getUUID(printerTags map[string][]string) string {
	var uuid string
	if u, ok := printerTags[attrPrinterUUID]; ok {
		uuid = u[0]
		uuid = strings.TrimPrefix(uuid, "urn:")
		uuid = strings.TrimPrefix(uuid, "uuid:")
	}
	return uuid
}

func getState(printerTags map[string][]string) cdd.CloudDeviceStateType {
	if s, ok := printerTags[attrPrinterState]; ok {
		switch s[0] {
		case "3":
			return cdd.CloudDeviceStateIdle
		case "4":
			return cdd.CloudDeviceStateProcessing
		case "5":
			return cdd.CloudDeviceStateStopped
		default:
			return cdd.CloudDeviceStateIdle
		}
	}
	return cdd.CloudDeviceStateIdle
}

func getVendorState(printerTags map[string][]string) *cdd.VendorState {
	reasons, exists := printerTags[attrPrinterStateReasons]
	if !exists || len(reasons) < 1 {
		return nil
	}

	sort.Strings(reasons)
	vendorState := &cdd.VendorState{Item: make([]cdd.VendorStateItem, len(reasons))}
	for i, reason := range reasons {
		vs := cdd.VendorStateItem{DescriptionLocalized: cdd.NewLocalizedString(reason)}
		if strings.HasSuffix(reason, "-error") {
			vs.State = cdd.VendorStateError
		} else if strings.HasSuffix(reason, "-warning") {
			vs.State = cdd.VendorStateWarning
		} else if strings.HasSuffix(reason, "-report") {
			vs.State = cdd.VendorStateInfo
		} else {
			vs.State = cdd.VendorStateError
		}
		vendorState.Item[i] = vs
	}

	return vendorState
}

func getAdobeVersionRange(pdfVersionsSupported []string) (string, string) {
	var min, max string
	for _, pdfVersion := range pdfVersionsSupported {
		var v string
		if _, err := fmt.Sscanf(pdfVersion, "adobe-%s", &v); err != nil || len(v) < 1 {
			continue
		}
		if min == "" || v < min {
			min = v
		}
		if max == "" || v > max {
			max = v
		}
	}
	return min, max
}

func convertSupportedContentType(printerTags map[string][]string) *[]cdd.SupportedContentType {
	pdfMin, pdfMax := getAdobeVersionRange(printerTags[attrPDFVersionsSupported])
	pdf := cdd.SupportedContentType{ContentType: "application/pdf"}
	if pdfMin != "" && pdfMax != "" {
		pdf.MinVersion = pdfMin
		pdf.MaxVersion = pdfMax
	}

	mimeTypes, exists := printerTags[attrDocumentFormatSupported]
	if !exists || len(mimeTypes) < 1 {
		return nil
	}

	// TODO: Filter these.

	// Preferred order:
	//  1) PDF because it's small.
	//  2) Postscript because it's a vector format.
	//  3) Any "native" formats so that they don't need conversion.
	//  4) PWG-Raster because it should work any time, but it's huge.

	sct := append(make([]cdd.SupportedContentType, 0, len(mimeTypes)),
		pdf, cdd.SupportedContentType{ContentType: "application/postscript"})
	for i := range mimeTypes {
		if mimeTypes[i] == "application/octet-stream" || // Avoid random byte blobs.
			mimeTypes[i] == "application/pdf" ||
			mimeTypes[i] == "application/postscript" ||
			mimeTypes[i] == "image/pwg-raster" {
			continue
		}
		sct = append(sct, cdd.SupportedContentType{ContentType: mimeTypes[i]})
	}
	/*
		TODO: Consider adding pwg-raster with config option to enable/disable.
		- All clients authored by Google do not create PWG Raster jobs.
		- cups-filters only supports pwg-raster input in recent versions.
		  https://www.cups.org/pipermail/cups/2015-July/026927.html
	*/
	return &sct
}

var cupsMarkerNameToGCP map[string]cdd.MarkerColorType = map[string]cdd.MarkerColorType{
	"black":        cdd.MarkerColorBlack,
	"color":        cdd.MarkerColorColor,
	"cyan":         cdd.MarkerColorCyan,
	"magenta":      cdd.MarkerColorMagenta,
	"yellow":       cdd.MarkerColorYellow,
	"lightcyan":    cdd.MarkerColorLightCyan,
	"lightmagenta": cdd.MarkerColorLightMagenta,
	"gray":         cdd.MarkerColorGray,
	"lightgray":    cdd.MarkerColorLightGray,
	"pigmentblack": cdd.MarkerColorPigmentBlack,
	"matteblack":   cdd.MarkerColorMatteBlack,
	"photocyan":    cdd.MarkerColorPhotoCyan,
	"photomagenta": cdd.MarkerColorPhotoMagenta,
	"photoyellow":  cdd.MarkerColorPhotoYellow,
	"photogray":    cdd.MarkerColorPhotoGray,
	"red":          cdd.MarkerColorRed,
	"green":        cdd.MarkerColorGreen,
	"blue":         cdd.MarkerColorBlue,
}

// convertMarkers converts CUPS marker-(names|types|levels) to *[]cdd.Marker and *cdd.MarkerState.
//
// Normalizes marker type: toner(Cartridge|-cartridge) => toner,
// ink(Cartridge|-cartridge|Ribbon|-ribbon) => ink
func convertMarkers(printerTags map[string][]string) (*[]cdd.Marker, *cdd.MarkerState) {
	names, types, levels := printerTags[attrMarkerNames], printerTags[attrMarkerTypes], printerTags[attrMarkerLevels]
	if len(names) == 0 || len(types) == 0 || len(levels) == 0 {
		return nil, nil
	}
	if len(names) != len(types) || len(types) != len(levels) {
		glog.Warningf("Received badly-formatted markers from CUPS: %s, %s, %s",
			strings.Join(names, ";"), strings.Join(types, ";"), strings.Join(levels, ";"))
		return nil, nil
	}

	markers := make([]cdd.Marker, 0, len(names))
	states := cdd.MarkerState{make([]cdd.MarkerStateItem, 0, len(names))}
	for i := 0; i < len(names); i++ {
		if len(names[i]) == 0 {
			return nil, nil
		}
		var markerType cdd.MarkerType
		switch strings.ToLower(types[i]) {
		case "toner", "tonercartridge", "toner-cartridge":
			markerType = cdd.MarkerToner
		case "ink", "inkcartridge", "ink-cartridge", "ink-ribbon", "inkribbon":
			markerType = cdd.MarkerInk
		case "staples":
			markerType = cdd.MarkerStaples
		default:
			continue
		}

		var color *cdd.MarkerColor
		if markerType == cdd.MarkerToner || markerType == cdd.MarkerInk {
			nameStripped := strings.Replace(strings.Replace(strings.ToLower(names[i]), " ", "", -1), "-", "", -1)
			colorType := cdd.MarkerColorCustom
			for k, v := range cupsMarkerNameToGCP {
				if strings.HasPrefix(nameStripped, k) {
					colorType = v
					break
				}
			}
			color = &cdd.MarkerColor{Type: colorType}
			if colorType == cdd.MarkerColorCustom {
				name := names[i]
				name = strings.TrimSuffix(name, " Cartridge")
				name = strings.TrimSuffix(name, " cartridge")
				name = strings.TrimSuffix(name, " Ribbon")
				name = strings.TrimSuffix(name, " ribbon")
				name = strings.TrimSuffix(name, " Toner")
				name = strings.TrimSuffix(name, " toner")
				name = strings.TrimSuffix(name, " Ink")
				name = strings.TrimSuffix(name, " ink")
				name = strings.Replace(name, "-", " ", -1)
				color.CustomDisplayNameLocalized = cdd.NewLocalizedString(name)
			}
		}

		marker := cdd.Marker{
			VendorID: names[i],
			Type:     markerType,
			Color:    color,
		}

		level, err := strconv.ParseInt(levels[i], 10, 32)
		if err != nil {
			glog.Warningf("Failed to parse CUPS marker state %s=%s: %s", names[i], levels[i], err)
			return nil, nil
		}
		if level > 100 {
			// Lop off extra (proprietary?) bits.
			level = level & 0x7f
		}
		if level < 0 || level > 100 {
			return nil, nil
		}

		var state cdd.MarkerStateType
		if level > 10 {
			state = cdd.MarkerStateOK
		} else {
			state = cdd.MarkerStateExhausted
		}
		level32 := int32(level)
		markerState := cdd.MarkerStateItem{
			VendorID:     names[i],
			State:        state,
			LevelPercent: &level32,
		}

		markers = append(markers, marker)
		states.Item = append(states.Item, markerState)
	}

	return &markers, &states
}

func convertPagesPerSheet(printerTags map[string][]string) *cdd.VendorCapability {
	numberUpSupported, exists := printerTags[attrNumberUpSupported]
	if !exists {
		return nil
	}

	c := cdd.VendorCapability{
		ID:                   attrNumberUp,
		Type:                 cdd.VendorCapabilitySelect,
		SelectCap:            &cdd.SelectCapability{},
		DisplayNameLocalized: cdd.NewLocalizedString("Pages per sheet"),
	}

	def, exists := printerTags[attrNumberUpDefault]
	if !exists {
		def = []string{"1"}
	}

	for _, number := range numberUpSupported {
		option := cdd.SelectCapabilityOption{
			Value:                number,
			IsDefault:            reflect.DeepEqual(number, def[0]),
			DisplayNameLocalized: cdd.NewLocalizedString(number),
		}
		c.SelectCap.Option = append(c.SelectCap.Option, option)
	}

	return &c
}

var (
	pageOrientationByValue map[string]cdd.PageOrientationType = map[string]cdd.PageOrientationType{
		"3":    cdd.PageOrientationPortrait,
		"4":    cdd.PageOrientationLandscape,
		"auto": cdd.PageOrientationAuto, // custom value, not CUPS standard
	}
	orientationValueByType map[cdd.PageOrientationType]string = map[cdd.PageOrientationType]string{
		cdd.PageOrientationPortrait:  "3",
		cdd.PageOrientationLandscape: "4",
	}
)

func convertPageOrientation(printerTags map[string][]string) *cdd.PageOrientation {
	orientationDefault, exists := printerTags[attrOrientationRequestedDefault]
	if !exists || len(orientationDefault) != 1 {
		orientationDefault = []string{"auto"}
	}

	orientationSupported, exists := printerTags[attrOrientationRequestedSupported]
	if !exists {
		return nil
	}

	pageOrientation := cdd.PageOrientation{}
	for _, orientation := range append([]string{"auto"}, orientationSupported...) {
		if po, exists := pageOrientationByValue[orientation]; exists {
			option := cdd.PageOrientationOption{
				Type:      po,
				IsDefault: orientation == orientationDefault[0],
			}
			pageOrientation.Option = append(pageOrientation.Option, option)
		}
	}
	return &pageOrientation
}

func convertCopies(printerTags map[string][]string) *cdd.Copies {
	var err error
	var def int64
	if copiesDefault, exists := printerTags[attrCopiesDefault]; !exists || len(copiesDefault) != 1 {
		def = 1
	} else {
		def, err = strconv.ParseInt(copiesDefault[0], 10, 32)
		if err != nil {
			def = 1
		}
	}

	var max int64
	if copiesSupported, exists := printerTags[attrCopiesSupported]; !exists || len(copiesSupported) != 1 {
		return nil
	} else {
		c := strings.SplitN(copiesSupported[0], "~", 2)
		max, err = strconv.ParseInt(c[1], 10, 32)
		if err != nil {
			return nil
		}
	}

	return &cdd.Copies{
		Default: int32(def),
		Max:     int32(max),
	}
}
<commit_msg>Only allow a subset of the MIME types that CUPS supports.<commit_after>/*
Copyright 2015 Google Inc. All rights reserved.

Use of this source code is governed by a BSD-style
license that can be found in the LICENSE file or at
https://developers.google.com/open-source/licenses/bsd
*/
package cups

import (
	"fmt"
	"reflect"
	"sort"
	"strconv"
	"strings"

	"github.com/golang/glog"
	"github.com/google/cups-connector/cdd"
)

// translateAttrs extracts a PrinterDescriptionSection, PrinterStateSection, name, default diplay name, UUID, and tags from maps of tags (CUPS attributes)
func translateAttrs(printerTags map[string][]string) (*cdd.PrinterDescriptionSection, *cdd.PrinterStateSection, string, string, string, map[string]string) {
	var name, info string
	if n, ok := printerTags[attrPrinterName]; ok && len(n) > 0 {
		name = n[0]
	}
	if i, ok := printerTags[attrPrinterInfo]; ok && len(i) > 0 {
		info = i[0]
	}
	uuid := getUUID(printerTags)

	var desc = cdd.PrinterDescriptionSection{VendorCapability: &[]cdd.VendorCapability{}}
	var state cdd.PrinterStateSection

	desc.SupportedContentType = convertSupportedContentType(printerTags)
	desc.Marker, state.MarkerState = convertMarkers(printerTags)
	desc.PageOrientation = convertPageOrientation(printerTags)
	desc.Copies = convertCopies(printerTags)
	if vc := convertPagesPerSheet(printerTags); vc != nil {
		*desc.VendorCapability = append(*desc.VendorCapability, *vc)
	}

	state.State = getState(printerTags)
	state.VendorState = getVendorState(printerTags)

	tags := make(map[string]string, len(printerTags))
	for k, v := range printerTags {
		tags[k] = strings.Join(v, ",")
	}

	return &desc, &state, name, info, uuid, tags
}

func getUUID(printerTags map[string][]string) string {
	var uuid string
	if u, ok := printerTags[attrPrinterUUID]; ok {
		uuid = u[0]
		uuid = strings.TrimPrefix(uuid, "urn:")
		uuid = strings.TrimPrefix(uuid, "uuid:")
	}
	return uuid
}

func getState(printerTags map[string][]string) cdd.CloudDeviceStateType {
	if s, ok := printerTags[attrPrinterState]; ok {
		switch s[0] {
		case "3":
			return cdd.CloudDeviceStateIdle
		case "4":
			return cdd.CloudDeviceStateProcessing
		case "5":
			return cdd.CloudDeviceStateStopped
		default:
			return cdd.CloudDeviceStateIdle
		}
	}
	return cdd.CloudDeviceStateIdle
}

func getVendorState(printerTags map[string][]string) *cdd.VendorState {
	reasons, exists := printerTags[attrPrinterStateReasons]
	if !exists || len(reasons) < 1 {
		return nil
	}

	sort.Strings(reasons)
	vendorState := &cdd.VendorState{Item: make([]cdd.VendorStateItem, len(reasons))}
	for i, reason := range reasons {
		vs := cdd.VendorStateItem{DescriptionLocalized: cdd.NewLocalizedString(reason)}
		if strings.HasSuffix(reason, "-error") {
			vs.State = cdd.VendorStateError
		} else if strings.HasSuffix(reason, "-warning") {
			vs.State = cdd.VendorStateWarning
		} else if strings.HasSuffix(reason, "-report") {
			vs.State = cdd.VendorStateInfo
		} else {
			vs.State = cdd.VendorStateError
		}
		vendorState.Item[i] = vs
	}

	return vendorState
}

func getAdobeVersionRange(pdfVersionsSupported []string) (string, string) {
	var min, max string
	for _, pdfVersion := range pdfVersionsSupported {
		var v string
		if _, err := fmt.Sscanf(pdfVersion, "adobe-%s", &v); err != nil || len(v) < 1 {
			continue
		}
		if min == "" || v < min {
			min = v
		}
		if max == "" || v > max {
			max = v
		}
	}
	return min, max
}

// CUPS can accept some unsafe types, like application/x-perl, so limit to these.
var mimeTypesAllowed = map[string]struct{}{
	"application/pdf":        struct{}{},
	"application/postscript": struct{}{},
	"image/pwg-raster":       struct{}{},
	"image/gif":              struct{}{},
	"image/jp2":              struct{}{},
	"image/jpeg":             struct{}{},
	"image/png":              struct{}{},
	"image/tiff":             struct{}{},
	"text/plain":             struct{}{},
	"text/rtf":               struct{}{},
}

func convertSupportedContentType(printerTags map[string][]string) *[]cdd.SupportedContentType {
	pdfMin, pdfMax := getAdobeVersionRange(printerTags[attrPDFVersionsSupported])
	pdf := cdd.SupportedContentType{ContentType: "application/pdf"}
	if pdfMin != "" && pdfMax != "" {
		pdf.MinVersion = pdfMin
		pdf.MaxVersion = pdfMax
	}
	sct := []cdd.SupportedContentType{pdf, cdd.SupportedContentType{ContentType: "application/postscript"}}

	mimeTypes, exists := printerTags[attrDocumentFormatSupported]
	if !exists || len(mimeTypes) < 1 {
		return &sct
	}

	// Preferred order:
	//  1) PDF (vector and small).
	//  2) Postscript (vector).
	//  3) Any CUPS-supported formats (don't need conversion in client or cloud).
	//  4) PWG-Raster (all clients support but it's huge).

	for _, mimeType := range mimeTypes {
		if mimeType == "application/pdf" || // Already added.
			mimeType == "application/postscript" || // Already added.
			mimeType == "image/pwg-raster" { // Added last, if at all.
			continue
		}
		if _, exists = mimeTypesAllowed[mimeType]; !exists {
			continue
		}
		sct = append(sct, cdd.SupportedContentType{ContentType: mimeType})
	}
	/*
		TODO: Consider adding pwg-raster with config option to enable/disable.
		- All clients authored by Google do not create PWG Raster jobs.
		- cups-filters only supports pwg-raster input in recent versions.
		  https://www.cups.org/pipermail/cups/2015-July/026927.html
	*/
	return &sct
}

var cupsMarkerNameToGCP map[string]cdd.MarkerColorType = map[string]cdd.MarkerColorType{
	"black":        cdd.MarkerColorBlack,
	"color":        cdd.MarkerColorColor,
	"cyan":         cdd.MarkerColorCyan,
	"magenta":      cdd.MarkerColorMagenta,
	"yellow":       cdd.MarkerColorYellow,
	"lightcyan":    cdd.MarkerColorLightCyan,
	"lightmagenta": cdd.MarkerColorLightMagenta,
	"gray":         cdd.MarkerColorGray,
	"lightgray":    cdd.MarkerColorLightGray,
	"pigmentblack": cdd.MarkerColorPigmentBlack,
	"matteblack":   cdd.MarkerColorMatteBlack,
	"photocyan":    cdd.MarkerColorPhotoCyan,
	"photomagenta": cdd.MarkerColorPhotoMagenta,
	"photoyellow":  cdd.MarkerColorPhotoYellow,
	"photogray":    cdd.MarkerColorPhotoGray,
	"red":          cdd.MarkerColorRed,
	"green":        cdd.MarkerColorGreen,
	"blue":         cdd.MarkerColorBlue,
}

// convertMarkers converts CUPS marker-(names|types|levels) to *[]cdd.Marker and *cdd.MarkerState.
//
// Normalizes marker type: toner(Cartridge|-cartridge) => toner,
// ink(Cartridge|-cartridge|Ribbon|-ribbon) => ink
func convertMarkers(printerTags map[string][]string) (*[]cdd.Marker, *cdd.MarkerState) {
	names, types, levels := printerTags[attrMarkerNames], printerTags[attrMarkerTypes], printerTags[attrMarkerLevels]
	if len(names) == 0 || len(types) == 0 || len(levels) == 0 {
		return nil, nil
	}
	if len(names) != len(types) || len(types) != len(levels) {
		glog.Warningf("Received badly-formatted markers from CUPS: %s, %s, %s",
			strings.Join(names, ";"), strings.Join(types, ";"), strings.Join(levels, ";"))
		return nil, nil
	}

	markers := make([]cdd.Marker, 0, len(names))
	states := cdd.MarkerState{make([]cdd.MarkerStateItem, 0, len(names))}
	for i := 0; i < len(names); i++ {
		if len(names[i]) == 0 {
			return nil, nil
		}
		var markerType cdd.MarkerType
		switch strings.ToLower(types[i]) {
		case "toner", "tonercartridge", "toner-cartridge":
			markerType = cdd.MarkerToner
		case "ink", "inkcartridge", "ink-cartridge", "ink-ribbon", "inkribbon":
			markerType = cdd.MarkerInk
		case "staples":
			markerType = cdd.MarkerStaples
		default:
			continue
		}

		var color *cdd.MarkerColor
		if markerType == cdd.MarkerToner || markerType == cdd.MarkerInk {
			nameStripped := strings.Replace(strings.Replace(strings.ToLower(names[i]), " ", "", -1), "-", "", -1)
			colorType := cdd.MarkerColorCustom
			for k, v := range cupsMarkerNameToGCP {
				if strings.HasPrefix(nameStripped, k) {
					colorType = v
					break
				}
			}
			color = &cdd.MarkerColor{Type: colorType}
			if colorType == cdd.MarkerColorCustom {
				name := names[i]
				name = strings.TrimSuffix(name, " Cartridge")
				name = strings.TrimSuffix(name, " cartridge")
				name = strings.TrimSuffix(name, " Ribbon")
				name = strings.TrimSuffix(name, " ribbon")
				name = strings.TrimSuffix(name, " Toner")
				name = strings.TrimSuffix(name, " toner")
				name = strings.TrimSuffix(name, " Ink")
				name = strings.TrimSuffix(name, " ink")
				name = strings.Replace(name, "-", " ", -1)
				color.CustomDisplayNameLocalized = cdd.NewLocalizedString(name)
			}
		}

		marker := cdd.Marker{
			VendorID: names[i],
			Type:     markerType,
			Color:    color,
		}

		level, err := strconv.ParseInt(levels[i], 10, 32)
		if err != nil {
			glog.Warningf("Failed to parse CUPS marker state %s=%s: %s", names[i], levels[i], err)
			return nil, nil
		}
		if level > 100 {
			// Lop off extra (proprietary?) bits.
			level = level & 0x7f
		}
		if level < 0 || level > 100 {
			return nil, nil
		}

		var state cdd.MarkerStateType
		if level > 10 {
			state = cdd.MarkerStateOK
		} else {
			state = cdd.MarkerStateExhausted
		}
		level32 := int32(level)
		markerState := cdd.MarkerStateItem{
			VendorID:     names[i],
			State:        state,
			LevelPercent: &level32,
		}

		markers = append(markers, marker)
		states.Item = append(states.Item, markerState)
	}

	return &markers, &states
}

func convertPagesPerSheet(printerTags map[string][]string) *cdd.VendorCapability {
	numberUpSupported, exists := printerTags[attrNumberUpSupported]
	if !exists {
		return nil
	}

	c := cdd.VendorCapability{
		ID:                   attrNumberUp,
		Type:                 cdd.VendorCapabilitySelect,
		SelectCap:            &cdd.SelectCapability{},
		DisplayNameLocalized: cdd.NewLocalizedString("Pages per sheet"),
	}

	def, exists := printerTags[attrNumberUpDefault]
	if !exists {
		def = []string{"1"}
	}

	for _, number := range numberUpSupported {
		option := cdd.SelectCapabilityOption{
			Value:                number,
			IsDefault:            reflect.DeepEqual(number, def[0]),
			DisplayNameLocalized: cdd.NewLocalizedString(number),
		}
		c.SelectCap.Option = append(c.SelectCap.Option, option)
	}

	return &c
}

var (
	pageOrientationByValue map[string]cdd.PageOrientationType = map[string]cdd.PageOrientationType{
		"3":    cdd.PageOrientationPortrait,
		"4":    cdd.PageOrientationLandscape,
		"auto": cdd.PageOrientationAuto, // custom value, not CUPS standard
	}
	orientationValueByType map[cdd.PageOrientationType]string = map[cdd.PageOrientationType]string{
		cdd.PageOrientationPortrait:  "3",
		cdd.PageOrientationLandscape: "4",
	}
)

func convertPageOrientation(printerTags map[string][]string) *cdd.PageOrientation {
	orientationDefault, exists := printerTags[attrOrientationRequestedDefault]
	if !exists || len(orientationDefault) != 1 {
		orientationDefault = []string{"auto"}
	}

	orientationSupported, exists := printerTags[attrOrientationRequestedSupported]
	if !exists {
		return nil
	}

	pageOrientation := cdd.PageOrientation{}
	for _, orientation := range append([]string{"auto"}, orientationSupported...) {
		if po, exists := pageOrientationByValue[orientation]; exists {
			option := cdd.PageOrientationOption{
				Type:      po,
				IsDefault: orientation == orientationDefault[0],
			}
			pageOrientation.Option = append(pageOrientation.Option, option)
		}
	}
	return &pageOrientation
}

func convertCopies(printerTags map[string][]string) *cdd.Copies {
	var err error
	var def int64
	if copiesDefault, exists := printerTags[attrCopiesDefault]; !exists || len(copiesDefault) != 1 {
		def = 1
	} else {
		def, err = strconv.ParseInt(copiesDefault[0], 10, 32)
		if err != nil {
			def = 1
		}
	}

	var max int64
	if copiesSupported, exists := printerTags[attrCopiesSupported]; !exists || len(copiesSupported) != 1 {
		return nil
	} else {
		c := strings.SplitN(copiesSupported[0], "~", 2)
		max, err = strconv.ParseInt(c[1], 10, 32)
		if err != nil {
			return nil
		}
	}

	return &cdd.Copies{
		Default: int32(def),
		Max:     int32(max),
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Authors of Cilium
//
// 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 RuntimeTest

import (
	"context"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
)

var _ = Describe("RuntimeKVStoreTest", func() {

	var vm *helpers.SSHMeta

	BeforeAll(func() {
		vm = helpers.InitRuntimeHelper(helpers.Runtime, logger)
		ExpectCiliumReady(vm)
	})

	containers := func(option string) {
		switch option {
		case helpers.Create:
			vm.NetworkCreate(helpers.CiliumDockerNetwork, "")
			vm.ContainerCreate(helpers.Client, helpers.NetperfImage, helpers.CiliumDockerNetwork, "-l id.client")
		case helpers.Delete:
			vm.ContainerRm(helpers.Client)

		}
	}

	BeforeEach(func() {
		vm.Exec("sudo systemctl stop cilium")
	}, 150)

	AfterEach(func() {
		containers(helpers.Delete)
		vm.Exec("sudo systemctl start cilium")
	})

	JustAfterEach(func() {
		vm.ValidateNoErrorsOnLogs(CurrentGinkgoTestDescription().Duration)
	})

	AfterFailed(func() {
		vm.ReportFailed("cilium status")
	})

	It("Consul KVStore", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		vm.ExecContext(
			ctx,
			"sudo cilium-agent --kvstore consul --kvstore-opt consul.address=127.0.0.1:8500 --debug")
		err := vm.WaitUntilReady(150)
		Expect(err).Should(BeNil())

		vm.Exec("sudo systemctl restart cilium-docker")
		helpers.Sleep(2)
		containers(helpers.Create)
		vm.WaitEndpointsReady()
		eps, err := vm.GetEndpointsNames()
		Expect(err).Should(BeNil())
		Expect(len(eps)).To(Equal(1))
	})

	It("Etcd KVStore", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		vm.ExecContext(
			ctx,
			"sudo cilium-agent --kvstore etcd --kvstore-opt etcd.address=127.0.0.1:4001 2>&1 | logger -t cilium")
		err := vm.WaitUntilReady(150)
		Expect(err).Should(BeNil())

		vm.Exec("sudo systemctl restart cilium-docker")
		helpers.Sleep(2)
		containers(helpers.Create)

		vm.WaitEndpointsReady()

		eps, err := vm.GetEndpointsNames()
		Expect(err).Should(BeNil())
		Expect(len(eps)).To(Equal(1))
	})
})
<commit_msg>test/runtime/kvstore: Wait for Cilium to be ready after restarting<commit_after>// Copyright 2017 Authors of Cilium
//
// 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 RuntimeTest

import (
	"context"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
)

var _ = Describe("RuntimeKVStoreTest", func() {

	var vm *helpers.SSHMeta

	BeforeAll(func() {
		vm = helpers.InitRuntimeHelper(helpers.Runtime, logger)
		ExpectCiliumReady(vm)
	})

	containers := func(option string) {
		switch option {
		case helpers.Create:
			vm.NetworkCreate(helpers.CiliumDockerNetwork, "")
			vm.ContainerCreate(helpers.Client, helpers.NetperfImage, helpers.CiliumDockerNetwork, "-l id.client")
		case helpers.Delete:
			vm.ContainerRm(helpers.Client)

		}
	}

	BeforeEach(func() {
		res := vm.ExecWithSudo("systemctl stop cilium")
		res.ExpectSuccess()
	}, 150)

	AfterEach(func() {
		containers(helpers.Delete)
		err := vm.RestartCilium()
		Expect(err).Should(BeNil())
	})

	JustAfterEach(func() {
		vm.ValidateNoErrorsOnLogs(CurrentGinkgoTestDescription().Duration)
	})

	AfterFailed(func() {
		vm.ReportFailed("cilium status")
	})

	It("Consul KVStore", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		vm.ExecContext(
			ctx,
			"sudo cilium-agent --kvstore consul --kvstore-opt consul.address=127.0.0.1:8500 --debug")
		err := vm.WaitUntilReady(150)
		Expect(err).Should(BeNil())

		vm.Exec("sudo systemctl restart cilium-docker")
		helpers.Sleep(2)
		containers(helpers.Create)
		vm.WaitEndpointsReady()
		eps, err := vm.GetEndpointsNames()
		Expect(err).Should(BeNil())
		Expect(len(eps)).To(Equal(1))
	})

	It("Etcd KVStore", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		vm.ExecContext(
			ctx,
			"sudo cilium-agent --kvstore etcd --kvstore-opt etcd.address=127.0.0.1:4001 2>&1 | logger -t cilium")
		err := vm.WaitUntilReady(150)
		Expect(err).Should(BeNil())

		vm.Exec("sudo systemctl restart cilium-docker")
		helpers.Sleep(2)
		containers(helpers.Create)

		vm.WaitEndpointsReady()

		eps, err := vm.GetEndpointsNames()
		Expect(err).Should(BeNil())
		Expect(len(eps)).To(Equal(1))
	})
})
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/kch42/gomcmap/mcmap"
	"github.com/mattn/go-gtk/gdk"
	"github.com/mattn/go-gtk/gdkpixbuf"
	"github.com/mattn/go-gtk/glib"
	"github.com/mattn/go-gtk/gtk"
	"os"
	"strconv"
)

type biomeEditFrame struct {
	*gtk.Frame
	applyBtn                          *gtk.Button
	idInput, snowLineInput, nameInput *gtk.Entry
	colorInput                        *gtk.ColorButton
}

func newBiomeEditFrame() *biomeEditFrame {
	frm := &biomeEditFrame{
		Frame:         gtk.NewFrame("Edit Biome"),
		applyBtn:      gtk.NewButtonWithLabel("Apply"),
		idInput:       gtk.NewEntry(),
		snowLineInput: gtk.NewEntry(),
		nameInput:     gtk.NewEntry(),
		colorInput:    gtk.NewColorButton(),
	}

	frm.idInput.SetSizeRequest(40, -1)
	frm.snowLineInput.SetSizeRequest(40, -1)

	frm.idInput.Connect("changed", frm.unlockApply)
	frm.nameInput.Connect("changed", frm.unlockApply)
	frm.snowLineInput.Connect("changed", frm.unlockApply)
	frm.applyBtn.SetSensitive(false)

	vbox := gtk.NewVBox(false, 0)
	hbox := gtk.NewHBox(false, 0)

	hbox.PackStart(gtk.NewLabel("Color:"), false, false, 0)
	hbox.PackStart(frm.colorInput, false, false, 3)
	hbox.PackStart(gtk.NewLabel("ID:"), false, false, 0)
	hbox.PackStart(frm.idInput, false, false, 3)
	hbox.PackStart(gtk.NewLabel("Snowline:"), false, false, 0)
	hbox.PackStart(frm.snowLineInput, false, false, 3)
	hbox.PackStart(gtk.NewLabel("Name:"), false, false, 0)
	hbox.PackStart(frm.nameInput, true, true, 3)

	vbox.PackStart(hbox, false, false, 0)
	vbox.PackStart(frm.applyBtn, false, false, 3)
	frm.Add(vbox)

	return frm
}

func (frm *biomeEditFrame) setBiomeInfo(info BiomeInfo) {
	frm.colorInput.SetColor(gdk.NewColor(info.Color))
	frm.idInput.SetText(strconv.FormatInt(int64(info.ID), 10))
	frm.snowLineInput.SetText(strconv.FormatInt(int64(info.SnowLine), 10))
	frm.nameInput.SetText(info.Name)
}

func (frm *biomeEditFrame) getBiomeInfo() (BiomeInfo, bool) {
	id, err := strconv.ParseUint(frm.idInput.GetText(), 10, 8)
	if err != nil {
		return BiomeInfo{}, false
	}

	snow, err := strconv.ParseInt(frm.snowLineInput.GetText(), 10, 32)
	if err != nil {
		return BiomeInfo{}, false
	}
	if (snow > mcmap.ChunkSizeY) || (snow < 0) {
		snow = mcmap.ChunkSizeY
	}

	name := frm.nameInput.GetText()
	if name == "" {
		return BiomeInfo{}, false
	}

	col := frm.colorInput.GetColor()

	return BiomeInfo{
		ID:       mcmap.Biome(id),
		SnowLine: int(snow),
		Name:     name,
		Color:    fmt.Sprintf("#%02x%02x%02x", col.Red()<<8, col.Green()<<8, col.Blue()<<8),
	}, true
}

func (frm *biomeEditFrame) checkOK() bool {
	_, ok := frm.getBiomeInfo()
	return ok
}

func (frm *biomeEditFrame) unlockApply() {
	frm.applyBtn.SetSensitive(frm.checkOK())
}

type biomeList struct {
	*gtk.HBox
	treeview *gtk.TreeView
	lStore   *gtk.ListStore
	biomes   []BiomeInfo
}

func newBiomeList() *biomeList {
	bl := &biomeList{
		HBox:     gtk.NewHBox(false, 0),
		treeview: gtk.NewTreeView(),
		lStore:   gtk.NewListStore(gdkpixbuf.GetType(), glib.G_TYPE_STRING, glib.G_TYPE_STRING, glib.G_TYPE_STRING),
	}

	scroll := gtk.NewScrolledWindow(nil, nil)
	scroll.SetPolicy(gtk.POLICY_NEVER, gtk.POLICY_AUTOMATIC)
	scroll.Add(bl.treeview)
	bl.PackStart(scroll, true, true, 3)

	bl.treeview.SetModel(bl.lStore)
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("Color", gtk.NewCellRendererPixbuf(), "pixbuf", 0))
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("ID", gtk.NewCellRendererText(), "text", 1))
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("Snowline", gtk.NewCellRendererText(), "text", 2))
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("Name", gtk.NewCellRendererText(), "text", 3))

	vbox := gtk.NewVBox(false, 0)

	addBtn := gtk.NewButton()
	addBtn.Add(gtk.NewImageFromStock(gtk.STOCK_ADD, gtk.ICON_SIZE_SMALL_TOOLBAR))
	delBtn := gtk.NewButton()
	delBtn.Add(gtk.NewImageFromStock(gtk.STOCK_DELETE, gtk.ICON_SIZE_SMALL_TOOLBAR))
	upBtn := gtk.NewButton()
	upBtn.Add(gtk.NewImageFromStock(gtk.STOCK_GO_UP, gtk.ICON_SIZE_SMALL_TOOLBAR))
	downBtn := gtk.NewButton()
	downBtn.Add(gtk.NewImageFromStock(gtk.STOCK_GO_DOWN, gtk.ICON_SIZE_SMALL_TOOLBAR))

	addBtn.Connect("clicked", bl.onAdd)
	delBtn.Connect("clicked", bl.onDel)
	upBtn.Connect("clicked", bl.onUp)
	downBtn.Connect("clicked", bl.onDown)

	vbox.PackStart(addBtn, false, false, 3)
	vbox.PackStart(delBtn, false, false, 3)
	vbox.PackStart(upBtn, false, false, 3)
	vbox.PackStart(downBtn, false, false, 3)

	bl.PackStart(vbox, false, false, 0)

	return bl
}

func (bl *biomeList) SetBiomes(biomes []BiomeInfo) {
	bl.biomes = biomes
	// TODO: Update View
}

func (bl *biomeList) Biomes() []BiomeInfo { return bl.biomes }

func (bl *biomeList) onAdd()  {} // TODO
func (bl *biomeList) onDel()  {} // TODO
func (bl *biomeList) onUp()   {} // TODO
func (bl *biomeList) onDown() {} // TODO

type BiomeInfoEditor struct {
	*gtk.Dialog
	biolist *biomeList
}

func NewBiomeInfoEditor(biomes []BiomeInfo) *BiomeInfoEditor {
	ed := &BiomeInfoEditor{
		Dialog:  gtk.NewDialog(),
		biolist: newBiomeList(),
	}

	ed.SetModal(true)

	vbox := ed.GetVBox()

	btnHBox := gtk.NewHBox(true, 0)

	resetBtn := gtk.NewButtonWithLabel("Reset to defaults")
	resetBtn.Connect("clicked", ed.reset)
	loadBtn := gtk.NewButtonWithLabel("Load from file ...")
	loadBtn.Connect("clicked", ed.load)
	saveBtn := gtk.NewButtonWithLabel("Save to file ...")
	saveBtn.Connect("clicked", ed.save)

	btnHBox.PackStart(resetBtn, true, true, 3)
	btnHBox.PackStart(loadBtn, true, true, 3)
	btnHBox.PackStart(saveBtn, true, true, 3)
	vbox.PackStart(btnHBox, false, false, 3)

	ed.biolist.SetBiomes(biomes)
	vbox.PackStart(ed.biolist, false, false, 3)

	editFrame := newBiomeEditFrame()
	vbox.PackStart(editFrame, false, false, 3)

	ed.AddButton("Cancel", gtk.RESPONSE_CANCEL)
	ed.AddButton("OK", gtk.RESPONSE_OK)
	ed.ShowAll()
	return ed
}

func (ed *BiomeInfoEditor) reset() {
	ed.biolist.SetBiomes(ReadDefaultBiomes())
}

func mkBioFFilters() (*gtk.FileFilter, *gtk.FileFilter) {
	f1 := gtk.NewFileFilter()
	f1.AddPattern("*.biomes")
	f1.SetName("Biome Infos (.biomes)")

	f2 := gtk.NewFileFilter()
	f2.AddPattern("*")
	f2.SetName("All files")

	return f1, f2
}

func errdlg(msg string, params ...interface{}) {
	dlg := gtk.NewMessageDialog(nil, gtk.DIALOG_MODAL|gtk.DIALOG_DESTROY_WITH_PARENT, gtk.MESSAGE_ERROR, gtk.BUTTONS_OK, msg, params...)
	dlg.Run()
	dlg.Destroy()
}

func (ed *BiomeInfoEditor) load() {
	dlg := gtk.NewFileChooserDialog("Load", nil, gtk.FILE_CHOOSER_ACTION_OPEN, "OK", gtk.RESPONSE_OK, "Cancel", gtk.RESPONSE_CANCEL)
	defer dlg.Destroy()
askFilename:
	if dlg.Run() == gtk.RESPONSE_OK {
		path := dlg.GetFilename()

		f, err := os.Open(path)
		if err != nil {
			errdlg("Could not load biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}
		defer f.Close()

		infos, err := ReadBiomeInfos(f)
		if err != nil {
			errdlg("Could not load biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}

		ed.biolist.SetBiomes(infos)
	}
}

func (ed *BiomeInfoEditor) save() {
	dlg := gtk.NewFileChooserDialog("Save", nil, gtk.FILE_CHOOSER_ACTION_SAVE, "OK", gtk.RESPONSE_OK, "Cancel", gtk.RESPONSE_CANCEL)
	defer dlg.Destroy()
askFilename:
	if dlg.Run() == gtk.RESPONSE_OK {
		path := dlg.GetFilename()

		if _, err := os.Stat(path); err == nil {
			qdlg := gtk.NewMessageDialog(nil, gtk.DIALOG_MODAL, gtk.MESSAGE_QUESTION, gtk.BUTTONS_YES_NO, "File %s already exists. Overwrite?", path)
			resp := qdlg.Run()
			qdlg.Destroy()

			if resp != gtk.RESPONSE_YES {
				goto askFilename
			}
		}

		f, err := os.Create(path)
		if err != nil {
			errdlg("Could not save biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}
		defer f.Close()

		if err := WriteBiomeInfos(ed.biolist.Biomes(), f); err != nil {
			errdlg("Could not save biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}
	}
}

func (ed *BiomeInfoEditor) Biomes() []BiomeInfo { return ed.biolist.Biomes() }
<commit_msg>TreeView is working :)<commit_after>package main

import (
	"fmt"
	"github.com/kch42/gomcmap/mcmap"
	"github.com/mattn/go-gtk/gdk"
	"github.com/mattn/go-gtk/glib"
	"github.com/mattn/go-gtk/gtk"
	"os"
	"strconv"
)

type biomeEditFrame struct {
	*gtk.Frame
	applyBtn                          *gtk.Button
	idInput, snowLineInput, nameInput *gtk.Entry
	colorInput                        *gtk.ColorButton
}

func newBiomeEditFrame() *biomeEditFrame {
	frm := &biomeEditFrame{
		Frame:         gtk.NewFrame("Edit Biome"),
		applyBtn:      gtk.NewButtonWithLabel("Apply"),
		idInput:       gtk.NewEntry(),
		snowLineInput: gtk.NewEntry(),
		nameInput:     gtk.NewEntry(),
		colorInput:    gtk.NewColorButton(),
	}

	frm.idInput.SetSizeRequest(40, -1)
	frm.snowLineInput.SetSizeRequest(40, -1)

	frm.idInput.Connect("changed", frm.unlockApply)
	frm.nameInput.Connect("changed", frm.unlockApply)
	frm.snowLineInput.Connect("changed", frm.unlockApply)
	frm.applyBtn.SetSensitive(false)

	vbox := gtk.NewVBox(false, 0)
	hbox := gtk.NewHBox(false, 0)

	hbox.PackStart(gtk.NewLabel("Color:"), false, false, 0)
	hbox.PackStart(frm.colorInput, false, false, 3)
	hbox.PackStart(gtk.NewLabel("ID:"), false, false, 0)
	hbox.PackStart(frm.idInput, false, false, 3)
	hbox.PackStart(gtk.NewLabel("Snowline:"), false, false, 0)
	hbox.PackStart(frm.snowLineInput, false, false, 3)
	hbox.PackStart(gtk.NewLabel("Name:"), false, false, 0)
	hbox.PackStart(frm.nameInput, true, true, 3)

	vbox.PackStart(hbox, false, false, 0)
	vbox.PackStart(frm.applyBtn, false, false, 3)
	frm.Add(vbox)

	return frm
}

func (frm *biomeEditFrame) setBiomeInfo(info BiomeInfo) {
	frm.colorInput.SetColor(gdk.NewColor(info.Color))
	frm.idInput.SetText(strconv.FormatInt(int64(info.ID), 10))
	frm.snowLineInput.SetText(strconv.FormatInt(int64(info.SnowLine), 10))
	frm.nameInput.SetText(info.Name)
}

func (frm *biomeEditFrame) getBiomeInfo() (BiomeInfo, bool) {
	id, err := strconv.ParseUint(frm.idInput.GetText(), 10, 8)
	if err != nil {
		return BiomeInfo{}, false
	}

	snow, err := strconv.ParseInt(frm.snowLineInput.GetText(), 10, 32)
	if err != nil {
		return BiomeInfo{}, false
	}
	if (snow > mcmap.ChunkSizeY) || (snow < 0) {
		snow = mcmap.ChunkSizeY
	}

	name := frm.nameInput.GetText()
	if name == "" {
		return BiomeInfo{}, false
	}

	col := frm.colorInput.GetColor()

	return BiomeInfo{
		ID:       mcmap.Biome(id),
		SnowLine: int(snow),
		Name:     name,
		Color:    fmt.Sprintf("#%02x%02x%02x", col.Red()<<8, col.Green()<<8, col.Blue()<<8),
	}, true
}

func (frm *biomeEditFrame) checkOK() bool {
	_, ok := frm.getBiomeInfo()
	return ok
}

func (frm *biomeEditFrame) unlockApply() {
	frm.applyBtn.SetSensitive(frm.checkOK())
}

type biomeList struct {
	*gtk.HBox
	treeview *gtk.TreeView
	lStore   *gtk.ListStore
	biomes   []BiomeInfo
}

func newBiomeList() *biomeList {
	bl := &biomeList{
		HBox:     gtk.NewHBox(false, 0),
		treeview: gtk.NewTreeView(),
		lStore:   gtk.NewListStore(glib.G_TYPE_STRING, glib.G_TYPE_STRING, glib.G_TYPE_STRING, glib.G_TYPE_STRING),
	}

	scroll := gtk.NewScrolledWindow(nil, nil)
	scroll.SetPolicy(gtk.POLICY_NEVER, gtk.POLICY_AUTOMATIC)
	scroll.Add(bl.treeview)
	bl.PackStart(scroll, true, true, 3)

	bl.treeview.SetModel(bl.lStore)
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("Color", gtk.NewCellRendererText(), "background", 0))
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("ID", gtk.NewCellRendererText(), "text", 1))
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("Snowline", gtk.NewCellRendererText(), "text", 2))
	bl.treeview.AppendColumn(gtk.NewTreeViewColumnWithAttributes("Name", gtk.NewCellRendererText(), "text", 3))

	bl.treeview.Connect("cursor-changed", bl.onCursorChanged)

	vbox := gtk.NewVBox(false, 0)

	addBtn := gtk.NewButton()
	addBtn.Add(gtk.NewImageFromStock(gtk.STOCK_ADD, gtk.ICON_SIZE_SMALL_TOOLBAR))
	delBtn := gtk.NewButton()
	delBtn.Add(gtk.NewImageFromStock(gtk.STOCK_DELETE, gtk.ICON_SIZE_SMALL_TOOLBAR))
	upBtn := gtk.NewButton()
	upBtn.Add(gtk.NewImageFromStock(gtk.STOCK_GO_UP, gtk.ICON_SIZE_SMALL_TOOLBAR))
	downBtn := gtk.NewButton()
	downBtn.Add(gtk.NewImageFromStock(gtk.STOCK_GO_DOWN, gtk.ICON_SIZE_SMALL_TOOLBAR))

	addBtn.Connect("clicked", bl.onAdd)
	delBtn.Connect("clicked", bl.onDel)
	upBtn.Connect("clicked", bl.onUp)
	downBtn.Connect("clicked", bl.onDown)

	vbox.PackStart(addBtn, false, false, 3)
	vbox.PackStart(delBtn, false, false, 3)
	vbox.PackStart(upBtn, false, false, 3)
	vbox.PackStart(downBtn, false, false, 3)

	bl.PackStart(vbox, false, false, 0)

	return bl
}

func (bl *biomeList) setBiome(iter *gtk.TreeIter, biome BiomeInfo) {
	bl.lStore.Set(iter, biome.Color, strconv.FormatInt(int64(biome.ID), 10), strconv.FormatInt(int64(biome.SnowLine), 10), biome.Name)
}

func (bl *biomeList) SetBiomes(biomes []BiomeInfo) {
	bl.biomes = biomes

	bl.lStore.Clear()
	var iter gtk.TreeIter
	for _, bio := range biomes {
		bl.lStore.Append(&iter)
		bl.setBiome(&iter, bio)
	}
}

func (bl *biomeList) Biomes() []BiomeInfo { return bl.biomes }

func (bl *biomeList) onCursorChanged() {
	// TODO
}

func (bl *biomeList) onAdd()  {} // TODO
func (bl *biomeList) onDel()  {} // TODO
func (bl *biomeList) onUp()   {} // TODO
func (bl *biomeList) onDown() {} // TODO

type BiomeInfoEditor struct {
	*gtk.Dialog
	biolist *biomeList
}

func NewBiomeInfoEditor(biomes []BiomeInfo) *BiomeInfoEditor {
	ed := &BiomeInfoEditor{
		Dialog:  gtk.NewDialog(),
		biolist: newBiomeList(),
	}

	ed.SetModal(true)

	vbox := ed.GetVBox()

	btnHBox := gtk.NewHBox(true, 0)

	resetBtn := gtk.NewButtonWithLabel("Reset to defaults")
	resetBtn.Connect("clicked", ed.reset)
	loadBtn := gtk.NewButtonWithLabel("Load from file ...")
	loadBtn.Connect("clicked", ed.load)
	saveBtn := gtk.NewButtonWithLabel("Save to file ...")
	saveBtn.Connect("clicked", ed.save)

	btnHBox.PackStart(resetBtn, true, true, 3)
	btnHBox.PackStart(loadBtn, true, true, 3)
	btnHBox.PackStart(saveBtn, true, true, 3)
	vbox.PackStart(btnHBox, false, false, 3)

	ed.biolist.SetBiomes(biomes)
	vbox.PackStart(ed.biolist, true, true, 3)

	editFrame := newBiomeEditFrame()
	vbox.PackStart(editFrame, false, false, 3)

	ed.AddButton("Cancel", gtk.RESPONSE_CANCEL)
	ed.AddButton("OK", gtk.RESPONSE_OK)
	ed.ShowAll()
	return ed
}

func (ed *BiomeInfoEditor) reset() {
	ed.biolist.SetBiomes(ReadDefaultBiomes())
}

func mkBioFFilters() (*gtk.FileFilter, *gtk.FileFilter) {
	f1 := gtk.NewFileFilter()
	f1.AddPattern("*.biomes")
	f1.SetName("Biome Infos (.biomes)")

	f2 := gtk.NewFileFilter()
	f2.AddPattern("*")
	f2.SetName("All files")

	return f1, f2
}

func errdlg(msg string, params ...interface{}) {
	dlg := gtk.NewMessageDialog(nil, gtk.DIALOG_MODAL|gtk.DIALOG_DESTROY_WITH_PARENT, gtk.MESSAGE_ERROR, gtk.BUTTONS_OK, msg, params...)
	dlg.Run()
	dlg.Destroy()
}

func (ed *BiomeInfoEditor) load() {
	dlg := gtk.NewFileChooserDialog("Load", nil, gtk.FILE_CHOOSER_ACTION_OPEN, "OK", gtk.RESPONSE_OK, "Cancel", gtk.RESPONSE_CANCEL)
	defer dlg.Destroy()
askFilename:
	if dlg.Run() == gtk.RESPONSE_OK {
		path := dlg.GetFilename()

		f, err := os.Open(path)
		if err != nil {
			errdlg("Could not load biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}
		defer f.Close()

		infos, err := ReadBiomeInfos(f)
		if err != nil {
			errdlg("Could not load biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}

		ed.biolist.SetBiomes(infos)
	}
}

func (ed *BiomeInfoEditor) save() {
	dlg := gtk.NewFileChooserDialog("Save", nil, gtk.FILE_CHOOSER_ACTION_SAVE, "OK", gtk.RESPONSE_OK, "Cancel", gtk.RESPONSE_CANCEL)
	defer dlg.Destroy()
askFilename:
	if dlg.Run() == gtk.RESPONSE_OK {
		path := dlg.GetFilename()

		if _, err := os.Stat(path); err == nil {
			qdlg := gtk.NewMessageDialog(nil, gtk.DIALOG_MODAL, gtk.MESSAGE_QUESTION, gtk.BUTTONS_YES_NO, "File %s already exists. Overwrite?", path)
			resp := qdlg.Run()
			qdlg.Destroy()

			if resp != gtk.RESPONSE_YES {
				goto askFilename
			}
		}

		f, err := os.Create(path)
		if err != nil {
			errdlg("Could not save biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}
		defer f.Close()

		if err := WriteBiomeInfos(ed.biolist.Biomes(), f); err != nil {
			errdlg("Could not save biome infos %s:\n%s", path, err.Error())
			goto askFilename
		}
	}
}

func (ed *BiomeInfoEditor) Biomes() []BiomeInfo { return ed.biolist.Biomes() }
<|endoftext|>


<commit_before>package blitzkrieg

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"runtime"
	"strings"
	"time"
	"strconv"
)

const (
	VERSION = "1.0"
)

var (
	count          int    // Number of requests per client
	clients        int    // Number of clients to simulate
	duration       int    // Duration of the test
	rate           int    // Rate limit
	url            string // URL
	urlsFilePath   string // Input file containing Urls
	keepAlive      bool   // Http Keep-alive on/off flag
	gzip           bool   // Accept gzip compression
	needLogin      bool   // Login on/off flag - if enabled the first url from urlsFilePath is used for login
	connectTimeout int    // Connect timeout in milliseconds
	readTimeout    int    // Read timeout in milliseconds
	writeTimeout   int    // Write timeout in milliseconds
	showErr        bool   // Show errors
	outFormat      string // Output Format
	version        bool   // Display version
	help           bool   // Display help
)

type blitzRequest struct {
	url    string
	method string
	header http.Header
	body   string
}

func (req *blitzRequest) getHttpRequest() (hReq *http.Request) {
	hReq, _ = http.NewRequest(req.method, req.url, strings.NewReader(req.body))
	hReq.Header = req.header
	return
}

func init() {
	flag.IntVar(&count, "n", -1, "Number of requests")
	flag.IntVar(&count, "number", -1, "")
	flag.IntVar(&clients, "c", 100, "Number of clients to simulate")
	flag.IntVar(&clients, "clients", 100, "")
	flag.IntVar(&duration, "d", -1, "Duration of the test in seconds")
	flag.IntVar(&duration, "duration", -1, "")
	flag.IntVar(&rate, "r", 0, "Rate limit")
	flag.IntVar(&rate, "rate", 0, "")
	flag.StringVar(&url, "u", "", "URL to test")
	flag.StringVar(&url, "url", "", "")
	flag.StringVar(&urlsFilePath, "f", "", "URLs file")
	flag.StringVar(&urlsFilePath, "file", "", "")
	flag.BoolVar(&keepAlive, "k", true, "Do keep HTTP keep-alive on")
	flag.BoolVar(&keepAlive, "keep", true, "")
	flag.BoolVar(&gzip, "g", true, "Accept Gzip")
	flag.BoolVar(&gzip, "gzip", true, "")
	flag.BoolVar(&needLogin, "l", false, "Login on/off flag")
	flag.BoolVar(&needLogin, "login", false, "")
	flag.BoolVar(&showErr, "e", false, "")
	flag.BoolVar(&showErr, "err", false, "Display Errors")
	flag.IntVar(&connectTimeout, "tc", 5000, "Connect timeout in ms")
	flag.IntVar(&connectTimeout, "timeoutcon", 5000, "")
	flag.IntVar(&readTimeout, "tr", 5000, "Read timeout in ms")
	flag.IntVar(&readTimeout, "timeoutread", 5000, "")
	flag.IntVar(&writeTimeout, "tw", 5000, "Write timeout in ms")
	flag.IntVar(&writeTimeout, "timeoutwrite", 5000, "")
	flag.StringVar(&outFormat, "o", "", "")
	flag.StringVar(&outFormat, "output", "", "Output Format")
	flag.BoolVar(&version, "v", false, "Prints the version number")
	flag.BoolVar(&help, "h", false, "Show Help")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: blitz [options]\n")
		fmt.Fprintf(os.Stderr, "Options:\n")
		fmt.Fprintf(os.Stderr, "-c,  -clients        Clients           Number of clients to simulate[default 100].\n")
		fmt.Fprintf(os.Stderr, "-n,  -number         Number            Number of requests\n")
		fmt.Fprintf(os.Stderr, "-d,  -duration       Duration          Duration of the test in seconds.\n")
		fmt.Fprintf(os.Stderr, "-r,  -rate           Rate              Rate limit.\n")
		fmt.Fprintf(os.Stderr, "-u,  -url            URL               URL to test.\n")
		fmt.Fprintf(os.Stderr, "-f,  -file           URLs File         URLs file.\n")
		fmt.Fprintf(os.Stderr, "-k,  -keep           KeepAlive         HTTP keep-alive on/off [default true].\n")
		fmt.Fprintf(os.Stderr, "-g,  -gzip           GZip              Accept Gzip Compression [default true].\n")
		fmt.Fprintf(os.Stderr, "-l,  -login          Login             Do login on/off.\n")
		fmt.Fprintf(os.Stderr, "-tc, -timeoutcon     ConnectTimeout    Connect timeout in ms [default 5000].\n")
		fmt.Fprintf(os.Stderr, "-tr, -timeoutread    ReadTimeout       Read timeout in ms [default 5000].\n")
		fmt.Fprintf(os.Stderr, "-tw, -timeoutwrite   WriteTimeout      Write timeout in ms [default 5000].\n")
		fmt.Fprintf(os.Stderr, "-o,  -output         OutputFormat      [graph].\n")
		fmt.Fprintf(os.Stderr, "-e,  -err            ShowErr           Display Errors.\n")
		fmt.Fprintf(os.Stderr, "-v,  -version        Version           Prints the version number.\n")
		fmt.Fprintf(os.Stderr, "-h,  -help           Help              Prints this output.\n")
	}
}

// showVersion displays the version number
func showVersion() {
	fmt.Println("blitz", VERSION)
	fmt.Println("This is free software. There is NO warranty")
}

// NewBlitz returns a new Blitz after parsing/processing
// the command line parameters
func NewBlitz() (blitz *Blitz) {
	flag.Parse()
	if version {
		showVersion()
		os.Exit(0)
	}
	if help {
		flag.Usage()
		os.Exit(0)
	}
	if urlsFilePath == "" && url == "" {
		flag.Usage()
		os.Exit(0)
	}
	if count == -1 && duration == -1 {
		flag.Usage()
		os.Exit(0)
	}
	blitz = &Blitz{
		requests:       make([]*blitzRequest, 0),
		count:          math.MaxInt32,
		clients:        clients,
		rate:           rate,
		keepAlive:      keepAlive,
		gzip:           gzip,
		connectTimeout: connectTimeout,
		readTimeout:    readTimeout,
		writeTimeout:   writeTimeout,
	}
	if count != -1 {
		blitz.count = count
	}

	if urlsFilePath != "" {
		requests, err := readFile(urlsFilePath)
		if err != nil {
			log.Fatalf("Error reading file:%s Error: ", urlsFilePath, err)
			os.Exit(0)
		}
		if len(requests) == 0 {
			log.Fatalf("File has insufficient number of lines")
			os.Exit(0)
		}
		blitz.requests = requests
	}

	if url != "" {
		blitz.requests = append(blitz.requests, &blitzRequest{url: url, method: "GET"})
	}

	if duration != -1 {
		blitz.duration = duration
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}
	return
}

// readFile reads a file containing urls and returns an array of
// http.Request elements
func readFile(path string) (requests []*blitzRequest, err error) {
	var (
		file         *os.File
		line         string
		arr          []string
		length       int
		req          *blitzRequest
		count        int
		loginCookies string
		tmpCookie    string
	)
	if file, err = os.Open(path); err != nil {
		return
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line = strings.TrimSpace(scanner.Text())
		if len(line) == 0 {
			continue
		}
		arr = strings.Split(line, "\t")
		length = len(arr)
		req = &blitzRequest{url: arr[0], method: "GET", header: make(http.Header)}
		if length > 1 {
			req.method = arr[1]
		}
		switch req.method {
		case "POST":
			switch length {
			case 4:
				req.header = parseHeaders(arr[3])
				fallthrough
			case 3:
				req.header.Add("Content-Type", "application/x-www-form-urlencoded")
    			req.header.Add("Content-Length", strconv.Itoa(len(arr[2])))
				req.body = arr[2]
			}
		case "GET":
			if length > 2 {
				req.header = parseHeaders(arr[2])
			}
		}
		req.header.Set("User-Agent", "blitz "+VERSION)

		if loginCookies != "" {
			tmpCookie = req.header.Get("Cookie")
			if tmpCookie != "" {
				req.header.Set("Cookie", loginCookies+tmpCookie)
			} else {
				req.header.Set("Cookie", loginCookies)
			}
		}
		count++
		if needLogin && (count == 1) {
			var buffer bytes.Buffer
			lCookies, _ := doLogin(req)
			for _, lCookie := range lCookies {
				buffer.WriteString(lCookie.Name)
				buffer.WriteString("=")
				buffer.WriteString(lCookie.Value)
				buffer.WriteString("; ")
			}
			loginCookies = buffer.String()
		} else {
			requests = append(requests, req)
		}
	}
	err = scanner.Err()
	return
}

// parseHeaders parses the header string and returns an
// http.Header
func parseHeaders(headerStr string) (header http.Header) {
	var (
		hArr []string
	)
	header = make(http.Header)
	splitArr := strings.Split(strings.TrimLeft(headerStr, "-H"), "-H")
	for _, str := range splitArr {
		str = strings.Trim(str, " '")
		hArr = strings.SplitN(str, ":", 2)
		if len(hArr) > 1 {
			header.Set(strings.TrimSpace(hArr[0]), strings.TrimSpace(hArr[1]))
		} else {
			log.Fatalf("Error parsing header string: %s", headerStr)
			os.Exit(0)
		}
	}
	return
}

func doLogin(req *blitzRequest) ([]*http.Cookie, error) {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	tr.Dial = func(network string, address string) (conn net.Conn, err error) {
		return net.DialTimeout(network, address, time.Duration(connectTimeout)*time.Millisecond)
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Do(req.getHttpRequest())
	if resp != nil && err == nil {
		return resp.Cookies(), nil
	}
	return nil, err

}
<commit_msg>added support for cookie jars<commit_after>package blitzkrieg

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"runtime"
	"strings"
	"time"
	"strconv"
	"code.google.com/p/go.net/publicsuffix"
	"net/http/cookiejar"
)

const (
	VERSION = "1.0"
)

var (
	count          int    // Number of requests per client
	clients        int    // Number of clients to simulate
	duration       int    // Duration of the test
	rate           int    // Rate limit
	url            string // URL
	urlsFilePath   string // Input file containing Urls
	keepAlive      bool   // Http Keep-alive on/off flag
	gzip           bool   // Accept gzip compression
	needLogin      bool   // Login on/off flag - if enabled the first url from urlsFilePath is used for login
	connectTimeout int    // Connect timeout in milliseconds
	readTimeout    int    // Read timeout in milliseconds
	writeTimeout   int    // Write timeout in milliseconds
	showErr        bool   // Show errors
	outFormat      string // Output Format
	version        bool   // Display version
	help           bool   // Display help
)

type blitzRequest struct {
	url    string
	method string
	header http.Header
	body   string
}

func (req *blitzRequest) getHttpRequest() (hReq *http.Request) {
	hReq, _ = http.NewRequest(req.method, req.url, strings.NewReader(req.body))
	hReq.Header = req.header
	return
}

func init() {
	flag.IntVar(&count, "n", -1, "Number of requests")
	flag.IntVar(&count, "number", -1, "")
	flag.IntVar(&clients, "c", 100, "Number of clients to simulate")
	flag.IntVar(&clients, "clients", 100, "")
	flag.IntVar(&duration, "d", -1, "Duration of the test in seconds")
	flag.IntVar(&duration, "duration", -1, "")
	flag.IntVar(&rate, "r", 0, "Rate limit")
	flag.IntVar(&rate, "rate", 0, "")
	flag.StringVar(&url, "u", "", "URL to test")
	flag.StringVar(&url, "url", "", "")
	flag.StringVar(&urlsFilePath, "f", "", "URLs file")
	flag.StringVar(&urlsFilePath, "file", "", "")
	flag.BoolVar(&keepAlive, "k", true, "Do keep HTTP keep-alive on")
	flag.BoolVar(&keepAlive, "keep", true, "")
	flag.BoolVar(&gzip, "g", true, "Accept Gzip")
	flag.BoolVar(&gzip, "gzip", true, "")
	flag.BoolVar(&needLogin, "l", false, "Login on/off flag")
	flag.BoolVar(&needLogin, "login", false, "")
	flag.BoolVar(&showErr, "e", false, "")
	flag.BoolVar(&showErr, "err", false, "Display Errors")
	flag.IntVar(&connectTimeout, "tc", 5000, "Connect timeout in ms")
	flag.IntVar(&connectTimeout, "timeoutcon", 5000, "")
	flag.IntVar(&readTimeout, "tr", 5000, "Read timeout in ms")
	flag.IntVar(&readTimeout, "timeoutread", 5000, "")
	flag.IntVar(&writeTimeout, "tw", 5000, "Write timeout in ms")
	flag.IntVar(&writeTimeout, "timeoutwrite", 5000, "")
	flag.StringVar(&outFormat, "o", "", "")
	flag.StringVar(&outFormat, "output", "", "Output Format")
	flag.BoolVar(&version, "v", false, "Prints the version number")
	flag.BoolVar(&help, "h", false, "Show Help")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: blitz [options]\n")
		fmt.Fprintf(os.Stderr, "Options:\n")
		fmt.Fprintf(os.Stderr, "-c,  -clients        Clients           Number of clients to simulate[default 100].\n")
		fmt.Fprintf(os.Stderr, "-n,  -number         Number            Number of requests\n")
		fmt.Fprintf(os.Stderr, "-d,  -duration       Duration          Duration of the test in seconds.\n")
		fmt.Fprintf(os.Stderr, "-r,  -rate           Rate              Rate limit.\n")
		fmt.Fprintf(os.Stderr, "-u,  -url            URL               URL to test.\n")
		fmt.Fprintf(os.Stderr, "-f,  -file           URLs File         URLs file.\n")
		fmt.Fprintf(os.Stderr, "-k,  -keep           KeepAlive         HTTP keep-alive on/off [default true].\n")
		fmt.Fprintf(os.Stderr, "-g,  -gzip           GZip              Accept Gzip Compression [default true].\n")
		fmt.Fprintf(os.Stderr, "-l,  -login          Login             Do login on/off.\n")
		fmt.Fprintf(os.Stderr, "-tc, -timeoutcon     ConnectTimeout    Connect timeout in ms [default 5000].\n")
		fmt.Fprintf(os.Stderr, "-tr, -timeoutread    ReadTimeout       Read timeout in ms [default 5000].\n")
		fmt.Fprintf(os.Stderr, "-tw, -timeoutwrite   WriteTimeout      Write timeout in ms [default 5000].\n")
		fmt.Fprintf(os.Stderr, "-o,  -output         OutputFormat      [graph].\n")
		fmt.Fprintf(os.Stderr, "-e,  -err            ShowErr           Display Errors.\n")
		fmt.Fprintf(os.Stderr, "-v,  -version        Version           Prints the version number.\n")
		fmt.Fprintf(os.Stderr, "-h,  -help           Help              Prints this output.\n")
	}
}

// showVersion displays the version number
func showVersion() {
	fmt.Println("blitz", VERSION)
	fmt.Println("This is free software. There is NO warranty")
}

// NewBlitz returns a new Blitz after parsing/processing
// the command line parameters
func NewBlitz() (blitz *Blitz) {
	flag.Parse()
	if version {
		showVersion()
		os.Exit(0)
	}
	if help {
		flag.Usage()
		os.Exit(0)
	}
	if urlsFilePath == "" && url == "" {
		flag.Usage()
		os.Exit(0)
	}
	if count == -1 && duration == -1 {
		flag.Usage()
		os.Exit(0)
	}
	blitz = &Blitz{
		requests:       make([]*blitzRequest, 0),
		count:          math.MaxInt32,
		clients:        clients,
		rate:           rate,
		keepAlive:      keepAlive,
		gzip:           gzip,
		connectTimeout: connectTimeout,
		readTimeout:    readTimeout,
		writeTimeout:   writeTimeout,
	}
	if count != -1 {
		blitz.count = count
	}

	if urlsFilePath != "" {
		requests, err := readFile(urlsFilePath)
		if err != nil {
			log.Fatalf("Error reading file:%s Error: ", urlsFilePath, err)
			os.Exit(0)
		}
		if len(requests) == 0 {
			log.Fatalf("File has insufficient number of lines")
			os.Exit(0)
		}
		blitz.requests = requests
	}

	if url != "" {
		blitz.requests = append(blitz.requests, &blitzRequest{url: url, method: "GET"})
	}

	if duration != -1 {
		blitz.duration = duration
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}
	return
}

// readFile reads a file containing urls and returns an array of
// http.Request elements
func readFile(path string) (requests []*blitzRequest, err error) {
	var (
		file         *os.File
		line         string
		arr          []string
		length       int
		req          *blitzRequest
		count        int
		loginCookies string
		tmpCookie    string
	)
	if file, err = os.Open(path); err != nil {
		return
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line = strings.TrimSpace(scanner.Text())
		if len(line) == 0 {
			continue
		}
		arr = strings.Split(line, "\t")
		length = len(arr)
		req = &blitzRequest{url: arr[0], method: "GET", header: make(http.Header)}
		if length > 1 {
			req.method = arr[1]
		}
		switch req.method {
		case "POST":
			switch length {
			case 4:
				req.header = parseHeaders(arr[3])
				fallthrough
			case 3:
				req.header.Add("Content-Type", "application/x-www-form-urlencoded")
    			req.header.Add("Content-Length", strconv.Itoa(len(arr[2])))
				req.body = arr[2]
			}
		case "GET":
			if length > 2 {
				req.header = parseHeaders(arr[2])
			}
		}
		req.header.Set("User-Agent", "blitz "+VERSION)

		if loginCookies != "" {
			tmpCookie = req.header.Get("Cookie")
			if tmpCookie != "" {
				req.header.Set("Cookie", loginCookies+tmpCookie)
			} else {
				req.header.Set("Cookie", loginCookies)
			}
		}
		count++
		if needLogin && (count == 1) {
			var buffer bytes.Buffer
			lCookies, _ := doLogin(req)
			for _, lCookie := range lCookies {
				buffer.WriteString(lCookie.Name)
				buffer.WriteString("=")
				buffer.WriteString(lCookie.Value)
				buffer.WriteString("; ")
			}
			loginCookies = buffer.String()
		} else {
			requests = append(requests, req)
		}
	}
	err = scanner.Err()
	return
}

// parseHeaders parses the header string and returns an
// http.Header
func parseHeaders(headerStr string) (header http.Header) {
	var (
		hArr []string
	)
	header = make(http.Header)
	splitArr := strings.Split(strings.TrimLeft(headerStr, "-H"), "-H")
	for _, str := range splitArr {
		str = strings.Trim(str, " '")
		hArr = strings.SplitN(str, ":", 2)
		if len(hArr) > 1 {
			header.Set(strings.TrimSpace(hArr[0]), strings.TrimSpace(hArr[1]))
		} else {
			log.Fatalf("Error parsing header string: %s", headerStr)
			os.Exit(0)
		}
	}
	return
}

func doLogin(req *blitzRequest) ([]*http.Cookie, error) {
	options := cookiejar.Options{
		PublicSuffixList: publicsuffix.List,
	}
	jar, cerr := cookiejar.New(&options)
	if cerr != nil {
		log.Fatal(cerr)
	}
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	tr.Dial = func(network string, address string) (conn net.Conn, err error) {
		return net.DialTimeout(network, address, time.Duration(connectTimeout)*time.Millisecond)
	}
	client := &http.Client{Transport: tr, Jar:jar}
	resp, err := client.Do(req.getHttpRequest())
	if resp != nil && err == nil {
		return resp.Cookies(), nil
	}
	return nil, err
}
<|endoftext|>


<commit_before>package prefab

import (
	"fmt"
	"log"
	"os/exec"
)

type DatabaseExtension struct {
	Name string `json:"name"`
}

type Database struct {
	Name       string              `json:"name"`
	Extensions []DatabaseExtension `json:"extensions"`
}

func (d Database) Create() error {
	log.Println("Create database:", d.Name)
	log.Println("Extensions: ", d.Extensions)

	_, err := exec.Command("sudo", "-u", "postgres", "createdb", d.Name).CombinedOutput()

	if err != nil {
		log.Println("Database exists")
	}

	for _, extension := range d.Extensions {
		sql_command := fmt.Sprintf("create extension %s", extension.Name)
		_, err := exec.Command("psql", d.Name, "-c", sql_command).CombinedOutput()
		if err != nil {
			log.Println("Extension exists")
		}
	}

	return nil
}
<commit_msg>Fix database extensions.<commit_after>package prefab

import (
	"fmt"
	"log"
	"os/exec"
)

type DatabaseExtension struct {
	Name string `json:"name"`
}

type Database struct {
	Name       string              `json:"name"`
	Extensions []DatabaseExtension `json:"extensions"`
}

func (d Database) Create() error {
	log.Println("Create database:", d.Name)

	_, err := exec.Command("sudo", "-u", "postgres", "createdb", d.Name).CombinedOutput()

	if err != nil {
		log.Println("Database exists")
	}

	for _, extension := range d.Extensions {
		sql_command := fmt.Sprintf("psql %s -c 'create extension %s'", d.Name, extension.Name)
		_, err := exec.Command("su", "postgres", "-c", sql_command).CombinedOutput()
		if err != nil {
			log.Println("Extension exists")
		}
	}

	return nil
}
<|endoftext|>