<commit_before>package csv

import (
	"reflect"
	"testing"
	"bytes"
)

type X struct {
	First string
}

type P struct {
	First string
	Last  string
}

func (p P) MarshalCSV() ([]byte, error) {
	return []byte(p.First + " " + p.Last), nil
}

func TestMarshal_without_a_slice(t *testing.T) {
	_, err := Marshal(simple{})

	if err == nil {
		t.Error("Non slice produced no error")
	}
}

func TestEncodeFieldValue(t *testing.T) {
	var encTests = []struct {
		val      interface{}
		expected string
		tag      string
	}{
		// Strings
		{"ABC", "ABC", ""},
		{byte(123), "123", ""},

		// Numerics
		{int(1), "1", ""},
		{float32(3.2), "3.2", ""},
		{uint32(123), "123", ""},
		{complex64(1 + 2i), "(+1+2i)", ""},

		// Boolean
		{true, "Yes", `true:"Yes" false:"No"`},
		{false, "No", `true:"Yes" false:"No"`},

		// TODO Array
		// Interface with Marshaler
		{P{"Jay", "Zee"}, "Jay Zee", ""},

		// Struct without Marshaler will produce nothing
		{X{"Jay"}, "", ""},
	}

	enc := &encoder{}

	for _, test := range encTests {
		fv := reflect.ValueOf(test.val)
		st := reflect.StructTag(test.tag)
		res := enc.encodeCol(fv, st)

		if res != test.expected {
			t.Errorf("%s does not match %s", res, test.expected)
		}
	}

}

func TestMarshalCsvOfStructs(t *testing.T) {

	type ST struct {
		A string
	}

	expected := []byte(`A
ASD
a
` )

	data := []struct{
		A string
	}{{"ASD"}}

	data = append(data , ST{"a"})
	out, err  := Marshal(data)
	if  err  != nil {
		t.Logf("Failed to marshal to csv ")
		t.Fail()
	}
	if ! bytes.Equal(out , expected) {
		t.Logf("Failed to marshal to correct format")
		t.Fail()
	}
}

func TestMarshalCsvOfPointers(t *testing.T) {

	expected := []byte(`A
ASD
` )
	tPointer := []interface{}{}
	tPointer = append(tPointer, struct{A string}{"ASD"})
	out, err  := Marshal(tPointer)
	if  err  != nil {
		t.Logf("Failed to marshal to csv ")
	}

	if ! bytes.Equal(out , expected) {
		t.Logf("Failed to marshal to correct format")
		t.Fail()
	}
}<commit_msg>simplified test<commit_after>package csv

import (
	"reflect"
	"testing"
	"bytes"
)

type X struct {
	First string
}

type P struct {
	First string
	Last  string
}

func (p P) MarshalCSV() ([]byte, error) {
	return []byte(p.First + " " + p.Last), nil
}

func TestMarshal_without_a_slice(t *testing.T) {
	_, err := Marshal(simple{})

	if err == nil {
		t.Error("Non slice produced no error")
	}
}

func TestEncodeFieldValue(t *testing.T) {
	var encTests = []struct {
		val      interface{}
		expected string
		tag      string
	}{
		// Strings
		{"ABC", "ABC", ""},
		{byte(123), "123", ""},

		// Numerics
		{int(1), "1", ""},
		{float32(3.2), "3.2", ""},
		{uint32(123), "123", ""},
		{complex64(1 + 2i), "(+1+2i)", ""},

		// Boolean
		{true, "Yes", `true:"Yes" false:"No"`},
		{false, "No", `true:"Yes" false:"No"`},

		// TODO Array
		// Interface with Marshaler
		{P{"Jay", "Zee"}, "Jay Zee", ""},

		// Struct without Marshaler will produce nothing
		{X{"Jay"}, "", ""},
	}

	enc := &encoder{}

	for _, test := range encTests {
		fv := reflect.ValueOf(test.val)
		st := reflect.StructTag(test.tag)
		res := enc.encodeCol(fv, st)

		if res != test.expected {
			t.Errorf("%s does not match %s", res, test.expected)
		}
	}

}

func TestMarshalCsvOfStructs(t *testing.T) {

	type ST struct {
		A string
	}

	expected := []byte(`A
ASD
a
` )
	data := []ST{{"ASD"}, {"a"}}
	out, err := Marshal(data)
	if err != nil {
		t.Logf("Failed to marshal to csv ")
		t.Fail()
	}
	if ! bytes.Equal(out, expected) {
		t.Logf("Failed to marshal to correct format")
		t.Fail()
	}
}

func TestMarshalCsvOfPointers(t *testing.T) {

	expected := []byte(`A
ASD
` )
	tPointer := []interface{}{}
	tPointer = append(tPointer, struct{ A string }{"ASD"})
	out, err := Marshal(tPointer)

	if err != nil {
		t.Logf("Failed to marshal to csv ")
	}

	if ! bytes.Equal(out, expected) {
		t.Logf("Failed to marshal to correct format")
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/viant/endly/bootstrap"
)

func main() {
	bootstrap.Bootstrap()
}
<commit_msg>added defulat import<commit_after>package main

import (
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	_ "github.com/viant/asc"
	_ "github.com/viant/bgc"
	_ "github.com/viant/endly/static" //load external resource like .csv .json files to mem storage
	_ "github.com/viant/toolbox/storage/aws"
	_ "github.com/viant/toolbox/storage/gs"

	"github.com/viant/endly/bootstrap"
)

func main() {
	bootstrap.Bootstrap()
}
<|endoftext|>


<commit_before>package neoutils

import (
	"github.com/Financial-Times/neo-cypher-runner-go"
	log "github.com/Sirupsen/logrus"
	"github.com/jmcvetta/neoism"
)

// StringerDb wraps neoism Database to provide a String function, which outputs the database URL
type StringerDb struct{ *neoism.Database }

func (sdb StringerDb) String() string {
	return sdb.Url
}

// Check will use the supplied CypherRunner to check connectivity to Neo4j
func Check(cr neocypherrunner.CypherRunner) error {
	results := []struct {
		node interface{}
	}{}

	query := &neoism.CypherQuery{
		Statement: `MATCH (n) RETURN id(n) LIMIT 1`,
		Result:    &results,
	}

	err := cr.CypherBatch([]*neoism.CypherQuery{query})

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

// EnsureIndexes will, for a map of labels and properties, check whether an index exists for a given property on a given label, and if missing will create one
func EnsureIndexes(im IndexManager, indexes map[string]string) error {
	for label, propertyName := range indexes {
		err := ensureIndex(im, label, propertyName)
		if err != nil { // stop as soon as something goes wrong
			return err
		}
	}
	return nil
}

// EnsureConstraints will, for a map of labels and properties, check whether a constraint exists for a given property on a given label, and
// if missing will create one. Creating the unique constraint ensures an index automatically.
func EnsureConstraints(im IndexManager, indexes map[string]string) error {
	for label, propertyName := range indexes {
		err := ensureConstraint(im, label, propertyName)
		if err != nil { // stop as soon as something goes wrong
			return err
		}
	}
	return nil
}

func ensureIndex(im IndexManager, label string, propertyName string) error {

	indexes, err := im.Indexes(label)

	var indexFound bool

	if err != nil && err != neoism.NotFound {
		return err
	}

	for _, index := range indexes {
		if len(index.PropertyKeys) == 1 && index.PropertyKeys[0] == propertyName {
			indexFound = true
			break
		}
	}

	if !indexFound {
		log.Infof("Creating index for type %s on property %s\n", label, propertyName)
		_, err := im.CreateIndex(label, propertyName)
		if err != nil {
			return err
		}
	}
	return nil

}

func ensureConstraint(im IndexManager, label string, propertyName string) error {
	_, err := im.UniqueConstraints(label, propertyName)

	if err != nil {
		if err == neoism.NotFound {
			log.Infof("Creating unique constraint for type %s on property %s\n", label, propertyName)
			_, err = im.CreateUniqueConstraint(label, propertyName)
			if err != nil {
				log.Errorf("Cannot create constraint for type %s on property %s\n", label, propertyName)
				return err
			}
		}
		return err
	}
	return nil

}

func getIndex(im IndexManager, label string, propertyName string) (*neoism.Index, error) {
	indexes, err := im.Indexes(label)

	if err != nil {
		return nil, err
	}

	for _, index := range indexes {
		if len(index.PropertyKeys) == 1 && index.PropertyKeys[0] == propertyName {
			return index, nil
		}
	}
	return nil, nil
}

// IndexManager manages the maintenance of indexes and unique constraints
type IndexManager interface {
	CreateIndex(label string, propertyName string) (*neoism.Index, error)
	Indexes(label string) ([]*neoism.Index, error)
	CreateUniqueConstraint(label string, propertyName string) (*neoism.UniqueConstraint, error)
	UniqueConstraints(label string, propertyName string) ([]*neoism.UniqueConstraint, error)
}
<commit_msg>Update package name for neocypherrunner<commit_after>package neoutils

import (
	"github.com/Financial-Times/neo-cypher-runner-go/neocypherrunner"
	log "github.com/Sirupsen/logrus"
	"github.com/jmcvetta/neoism"
)

// StringerDb wraps neoism Database to provide a String function, which outputs the database URL
type StringerDb struct{ *neoism.Database }

func (sdb StringerDb) String() string {
	return sdb.Url
}

// Check will use the supplied CypherRunner to check connectivity to Neo4j
func Check(cr neocypherrunner.CypherRunner) error {
	results := []struct {
		node interface{}
	}{}

	query := &neoism.CypherQuery{
		Statement: `MATCH (n) RETURN id(n) LIMIT 1`,
		Result:    &results,
	}

	err := cr.CypherBatch([]*neoism.CypherQuery{query})

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

// EnsureIndexes will, for a map of labels and properties, check whether an index exists for a given property on a given label, and if missing will create one
func EnsureIndexes(im IndexManager, indexes map[string]string) error {
	for label, propertyName := range indexes {
		err := ensureIndex(im, label, propertyName)
		if err != nil { // stop as soon as something goes wrong
			return err
		}
	}
	return nil
}

// EnsureConstraints will, for a map of labels and properties, check whether a constraint exists for a given property on a given label, and
// if missing will create one. Creating the unique constraint ensures an index automatically.
func EnsureConstraints(im IndexManager, indexes map[string]string) error {
	for label, propertyName := range indexes {
		err := ensureConstraint(im, label, propertyName)
		if err != nil { // stop as soon as something goes wrong
			return err
		}
	}
	return nil
}

func ensureIndex(im IndexManager, label string, propertyName string) error {

	indexes, err := im.Indexes(label)

	var indexFound bool

	if err != nil && err != neoism.NotFound {
		return err
	}

	for _, index := range indexes {
		if len(index.PropertyKeys) == 1 && index.PropertyKeys[0] == propertyName {
			indexFound = true
			break
		}
	}

	if !indexFound {
		log.Infof("Creating index for type %s on property %s\n", label, propertyName)
		_, err := im.CreateIndex(label, propertyName)
		if err != nil {
			return err
		}
	}
	return nil

}

func ensureConstraint(im IndexManager, label string, propertyName string) error {
	_, err := im.UniqueConstraints(label, propertyName)

	if err != nil {
		if err == neoism.NotFound {
			log.Infof("Creating unique constraint for type %s on property %s\n", label, propertyName)
			_, err = im.CreateUniqueConstraint(label, propertyName)
			if err != nil {
				log.Errorf("Cannot create constraint for type %s on property %s\n", label, propertyName)
				return err
			}
		}
		return err
	}
	return nil

}

func getIndex(im IndexManager, label string, propertyName string) (*neoism.Index, error) {
	indexes, err := im.Indexes(label)

	if err != nil {
		return nil, err
	}

	for _, index := range indexes {
		if len(index.PropertyKeys) == 1 && index.PropertyKeys[0] == propertyName {
			return index, nil
		}
	}
	return nil, nil
}

// IndexManager manages the maintenance of indexes and unique constraints
type IndexManager interface {
	CreateIndex(label string, propertyName string) (*neoism.Index, error)
	Indexes(label string) ([]*neoism.Index, error)
	CreateUniqueConstraint(label string, propertyName string) (*neoism.UniqueConstraint, error)
	UniqueConstraints(label string, propertyName string) ([]*neoism.UniqueConstraint, error)
}
<|endoftext|>


<commit_before>package main

import (
        "database/sql"
        "encoding/xml"
        "flag"
        "fmt"
        _ "github.com/lib/pq"
        "log"
        "os"
        "path/filepath"
        "runtime"
        "sync"
        "time"
)

var count int = 0

type ReportHost struct {
        XMLName        xml.Name       `xml:"ReportHost"`
        HostProperties HostProperties `xml:"HostProperties"`
        ReportItemList []ReportItem   `xml:"ReportItem"`
}

type HostProperties struct {
        XMLName xml.Name   `xml:HostProperties"`
        Info    []InfoList `xml:"tag"`
}

type InfoList struct {
        Key     string  `xml:"name,attr"`
        Value   string  `xml:",chardata"`
}

type ReportItem struct {
        CVE                        []string `xml:"cve"`
        CvssBaseScore              float64  `xml:"cvss_base_score"`
        Desciption                 string   `xml:"description"`
        ExploitAvailable           bool     `xml:"exploit_available"`
        ExploitFrameworkCanvas     bool     `xml:"exploit_framework_canvas"`
        ExploitFrameworkCore       bool     `xml:"exploit_framework_core"`
        ExploitFrameworkMetasploit bool     `xml:"exploit_framework_metasploit"`
        ExploitEase                string   `xml:"exploit_ease"`
        MetasploitName             string   `xml:"metasploit_name"`
        PluginName                 string   `xml:"plugin_name"`
        RiskFactor                 string   `xml:"risk_factor"`
        Solution                   string   `xml:"solution"`
        Synopsis                   string   `xml:"synopsis"`
        PluginOutput               string   `xml:"plugin_output"`
        SeeAlso                    string   `xml:"see_also"`
        ExploitedMalware           bool     `xml:"exploited_by_malware"`
        Severity                   int      `xml:"severity,attr"`
        PluginID                   int      `xml:"pluginID,attr"`
}

type readyData struct {
        ip                string
        os                string
        mac               string
        fqdn              string
        netbios           string
        cve               []string
        cvss              float64
        desc              string
        exploitAvailable  bool
        exploitCanvas     bool
        exploitCore       bool
        exploitMetasploit bool
        exploitEase       string
        metasploitName    string
        pluginName        string
        risk              string
        solution          string
        synopsis          string
        pluginOutput      string
        seeAlso           string
        exploitedMalware  bool
        severity          int
        pluginID          int
}

func dataPrep(prep *sql.Stmt, report *ReportHost, db *sql.DB) {
        hashmap := make(map[string]string)
        for _, host := range report.HostProperties.Info {
                hashmap[host.Key] = host.Value
        }
        var rd readyData
        rd.ip = hashmap["host-ip"]
        rd.os = hashmap["operating-system"]
        rd.mac = hashmap["mac-address"]
        rd.fqdn = hashmap["host-fqdn"]
        rd.netbios = hashmap["netbios-name"]
        for _, host := range report.ReportItemList {
                rd.cve = host.CVE
                rd.cvss = host.CvssBaseScore
                rd.desc = host.Desciption
                rd.exploitAvailable = host.ExploitAvailable
                rd.exploitCanvas = host.ExploitFrameworkCanvas
                rd.exploitCore = host.ExploitFrameworkCore
                rd.exploitMetasploit = host.ExploitFrameworkMetasploit
                rd.exploitEase = host.ExploitEase
                rd.metasploitName = host.MetasploitName
                rd.pluginName = host.PluginName
                rd.risk = host.RiskFactor
                rd.solution = host.Solution
                rd.synopsis = host.Synopsis
                rd.pluginOutput = host.PluginOutput
                rd.seeAlso = host.SeeAlso
                rd.exploitedMalware = host.ExploitedMalware
                rd.severity = host.Severity
                rd.pluginID = host.PluginID
                databaseImport(prep, &rd, db)
        }
        if verbose == 2 {
                fmt.Println("")
        }
}

func xmlParse(wg *sync.WaitGroup, xmlFile *os.File, prep *sql.Stmt, db *sql.DB) {
        decoder := xml.NewDecoder(xmlFile)
        for {
                count = 0
                t, _ := decoder.Token()
                if t == nil {
                        break
                }
                switch se := t.(type) {
                case xml.StartElement:
                        if se.Name.Local == "ReportHost" {
                                var report ReportHost

                                decoder.DecodeElement(&report, &se)
                                dataPrep(prep, &report, db)
                        }
                }

        }
        wg.Done()
}

// Database import
func databaseImport(prep *sql.Stmt, d *readyData, db *sql.DB) {
        month := int(time.Now().Month())
        year := int(time.Now().Year())
        if len(d.cve) == 0 {
                d.cve = append(d.cve, "None")
        }
        for _, cve := range d.cve {
                res, err := prep.Exec(d.ip, d.mac, d.netbios, d.fqdn, d.os,
                        d.pluginName, d.pluginID, d.severity, cve, d.risk,
                        d.desc, d.solution, d.synopsis, d.pluginOutput, d.seeAlso,
                        d.exploitAvailable, d.exploitEase, d.exploitMetasploit,
                        d.metasploitName, d.exploitCanvas, d.exploitCore,
                        d.exploitedMalware, d.cvss, month, year)
                if verbose == 2 {
                        fmt.Printf("Importing %-15s Items:%4d \r", d.ip, count)
                        count += 1
                }
                if err != nil || res == nil {
                        log.Fatal(err)
                }
        }

}

func createTable(db *sql.DB, table *string) {
        var stm string = `create table if not exists ` + *table + `(id serial,
        host text,mac_address text, netbios text, fqdn text,os_name text,
        plugin_name text, plugin_id integer,severity integer, cve text,
        risk text,description text, solution text, synopsis text,
        plugin_output text, see_also text,exploit_available boolean,
        exploit_ease text,metasploit_framework boolean, metasploit_name text,
        canvas_framework boolean, core_framework boolean,
        exploited_malware boolean, cvss float,month integer, year integer)`
        prep, err := db.Prepare(stm)
        if err != nil {
                log.Fatal(err)
        }
        res, err := prep.Exec()
        if err != nil || res == nil {
                log.Fatal(err)
        }

        prep.Close()
}

func dropTable(db *sql.DB, table *string) {
        var stm string = "drop table if exists " + *table
        prep, err := db.Prepare(stm)
        if err != nil {
                log.Fatal(err)
        }
        res, err := prep.Exec()
        if err != nil || res == nil {
                log.Fatal(err)
        }
        prep.Close()
}

func coreCheck(cores *int) {
        if *cores > runtime.NumCPU() || *cores <= 0 {
                log.Fatal(`You don't have that many cores.... you can use up to `,
                        runtime.NumCPU())
        } else {
                runtime.GOMAXPROCS(*cores)
                if verbose == 1 {
                        fmt.Println("Using " + string(runtime.GOMAXPROCS(*cores)) + " cores")
                }
        }
}

var verbose int

func main() {
        fileOpt := flag.String("file", "xmlFile", "file to parse into db")
        dirOpt := flag.String("dir", "directory", "dir of xml files")
        coreOpt := flag.Int("cores", 1, "Number of Cores to use")
        verboseOpt := flag.Int("verbose", 0, "Verbose level 0,1,2")
        userOpt := flag.String("user", "postgres", "User for Postgres")
        passOpt := flag.String("pass", "", "Password for Postgres user")
        dbOpt := flag.String("db", "gotest", "DB to use")
        tableOpt := flag.String("table", "internal_network", "Table to use")
        sslOpt := flag.String("ssl", "disable", "Enable or Disable")
        flag.Parse()
        coreCheck(coreOpt)
        verbose = *verboseOpt
        var con_opts string
        if *passOpt == "" {
                con_opts = "user=" + *userOpt +
                        " dbname=" + *dbOpt + " sslmode=" + *sslOpt
        } else {
                con_opts = "user=" + *userOpt + " password=" + *passOpt +
                        " dbname=" + *dbOpt + " sslmode=" + *sslOpt
        }
        db, err := sql.Open("postgres", con_opts)
        if err != nil {
                fmt.Println("Error Connecting:", err)
                return
        }
        defer db.Close()
        txn, err := db.Begin()
        if err != nil {
                log.Fatal(err)
        }
        //dropTable(db, tableOpt)
        createTable(db, tableOpt)
        prep, err := txn.Prepare(`copy ` + *tableOpt + ` (host, mac_address, netbios,
                fqdn, os_name, plugin_name, plugin_id, severity, cve,
                risk, description, solution, synopsis, plugin_output,
                see_also, exploit_available, exploit_ease, metasploit_framework,
                metasploit_name, canvas_framework, core_framework,
                exploited_malware, cvss, month, year) from stdin`)
        if err != nil {
                log.Fatal(err)
        }
        file := *fileOpt
        dir := *dirOpt
        fmt.Println("Importing..")
        var wg sync.WaitGroup
        if file != "xmlFile" {
                xmlFile, err := os.Open(file)
                defer xmlFile.Close()
                if err != nil {
                        log.Fatal(err)
                        return
                }
                wg.Add(1)
                xmlParse(&wg, xmlFile, prep, db)
        } else if dir != "directory" {
                files, _ := filepath.Glob(dir + "/*")
                for _, file := range files {
                        xmlFile, err := os.Open(file)
                        defer xmlFile.Close()
                        if err != nil {
                                log.Fatal(err)
                                return
                        }
                        if verbose == 1 {
                                fmt.Printf("Parsing %s\n", file)
                        }
                        wg.Add(1)
                        go xmlParse(&wg, xmlFile, prep, db)
                }

        }
        wg.Wait()
        if err != nil {
                log.Fatal(err)
        }
        _, err = prep.Exec()
        if err != nil {
                log.Fatal(err)
        }
        err = prep.Close()
        if err != nil {
                log.Fatal(err)
        }
        err = txn.Commit()
        if err != nil {
                log.Fatal(err)
        }
}
<commit_msg>More Goroutines<commit_after>package main

import (
        "database/sql"
        "encoding/xml"
        "flag"
        "fmt"
        _ "github.com/lib/pq"
        "log"
        "os"
        "path/filepath"
        "runtime"
        "sync"
        "time"
)

var count int = 0

type ReportHost struct {
        XMLName        xml.Name       `xml:"ReportHost"`
        HostProperties HostProperties `xml:"HostProperties"`
        ReportItemList []ReportItem   `xml:"ReportItem"`
}

type HostProperties struct {
        XMLName xml.Name   `xml:HostProperties"`
        Info    []InfoList `xml:"tag"`
}

type InfoList struct {
        Key     string  `xml:"name,attr"`
        Value   string  `xml:",chardata"`
}

type ReportItem struct {
        CVE                        []string `xml:"cve"`
        CvssBaseScore              float64  `xml:"cvss_base_score"`
        Desciption                 string   `xml:"description"`
        ExploitAvailable           bool     `xml:"exploit_available"`
        ExploitFrameworkCanvas     bool     `xml:"exploit_framework_canvas"`
        ExploitFrameworkCore       bool     `xml:"exploit_framework_core"`
        ExploitFrameworkMetasploit bool     `xml:"exploit_framework_metasploit"`
        ExploitEase                string   `xml:"exploit_ease"`
        MetasploitName             string   `xml:"metasploit_name"`
        PluginName                 string   `xml:"plugin_name"`
        RiskFactor                 string   `xml:"risk_factor"`
        Solution                   string   `xml:"solution"`
        Synopsis                   string   `xml:"synopsis"`
        PluginOutput               string   `xml:"plugin_output"`
        SeeAlso                    string   `xml:"see_also"`
        ExploitedMalware           bool     `xml:"exploited_by_malware"`
        Severity                   int      `xml:"severity,attr"`
        PluginID                   int      `xml:"pluginID,attr"`
}

type readyData struct {
        ip                string
        os                string
        mac               string
        fqdn              string
        netbios           string
        cve               []string
        cvss              float64
        desc              string
        exploitAvailable  bool
        exploitCanvas     bool
        exploitCore       bool
        exploitMetasploit bool
        exploitEase       string
        metasploitName    string
        pluginName        string
        risk              string
        solution          string
        synopsis          string
        pluginOutput      string
        seeAlso           string
        exploitedMalware  bool
        severity          int
        pluginID          int
}

func dataPrep(wg *sync.WaitGroup, prep *sql.Stmt, report *ReportHost, db *sql.DB) {
        hashmap := make(map[string]string)
        for _, host := range report.HostProperties.Info {
                hashmap[host.Key] = host.Value
        }
        var rd readyData
        rd.ip = hashmap["host-ip"]
        rd.os = hashmap["operating-system"]
        rd.mac = hashmap["mac-address"]
        rd.fqdn = hashmap["host-fqdn"]
        rd.netbios = hashmap["netbios-name"]
        for _, host := range report.ReportItemList {
                rd.cve = host.CVE
                rd.cvss = host.CvssBaseScore
                rd.desc = host.Desciption
                rd.exploitAvailable = host.ExploitAvailable
                rd.exploitCanvas = host.ExploitFrameworkCanvas
                rd.exploitCore = host.ExploitFrameworkCore
                rd.exploitMetasploit = host.ExploitFrameworkMetasploit
                rd.exploitEase = host.ExploitEase
                rd.metasploitName = host.MetasploitName
                rd.pluginName = host.PluginName
                rd.risk = host.RiskFactor
                rd.solution = host.Solution
                rd.synopsis = host.Synopsis
                rd.pluginOutput = host.PluginOutput
                rd.seeAlso = host.SeeAlso
                rd.exploitedMalware = host.ExploitedMalware
                rd.severity = host.Severity
                rd.pluginID = host.PluginID
                databaseImport(prep, &rd, db)
        }
        if verbose == 2 {
                fmt.Println("")
        }
        wg.Done()
}

func xmlParse(wg *sync.WaitGroup, xmlFile *os.File, prep *sql.Stmt, db *sql.DB) {
        decoder := xml.NewDecoder(xmlFile)
        for {
                count = 0
                t, _ := decoder.Token()
                if t == nil {
                        break
                }
                switch se := t.(type) {
                case xml.StartElement:
                        if se.Name.Local == "ReportHost" {
                                var report ReportHost
                                wg.Add(1)
                                decoder.DecodeElement(&report, &se)
                                go dataPrep(wg, prep, &report, db)
                        }
                }

        }
        wg.Done()
}

// Database import
func databaseImport(prep *sql.Stmt, d *readyData, db *sql.DB) {
        month := int(time.Now().Month())
        year := int(time.Now().Year())
        if len(d.cve) == 0 {
                d.cve = append(d.cve, "None")
        }
        for _, cve := range d.cve {
                res, err := prep.Exec(d.ip, d.mac, d.netbios, d.fqdn, d.os,
                        d.pluginName, d.pluginID, d.severity, cve, d.risk,
                        d.desc, d.solution, d.synopsis, d.pluginOutput, d.seeAlso,
                        d.exploitAvailable, d.exploitEase, d.exploitMetasploit,
                        d.metasploitName, d.exploitCanvas, d.exploitCore,
                        d.exploitedMalware, d.cvss, month, year)
                if verbose == 2 {
                        fmt.Printf("Importing %-15s Items:%4d \r", d.ip, count)
                        count += 1
                }
                if err != nil || res == nil {
                        log.Fatal(err)
                }
        }

}

func createTable(db *sql.DB, table *string) {
        var stm string = `create table if not exists ` + *table + `(id serial,
        host text,mac_address text, netbios text, fqdn text,os_name text,
        plugin_name text, plugin_id integer,severity integer, cve text,
        risk text,description text, solution text, synopsis text,
        plugin_output text, see_also text,exploit_available boolean,
        exploit_ease text,metasploit_framework boolean, metasploit_name text,
        canvas_framework boolean, core_framework boolean,
        exploited_malware boolean, cvss float,month integer, year integer)`
        prep, err := db.Prepare(stm)
        if err != nil {
                log.Fatal(err)
        }
        res, err := prep.Exec()
        if err != nil || res == nil {
                log.Fatal(err)
        }

        prep.Close()
}

func dropTable(db *sql.DB, table *string) {
        var stm string = "drop table if exists " + *table
        prep, err := db.Prepare(stm)
        if err != nil {
                log.Fatal(err)
        }
        res, err := prep.Exec()
        if err != nil || res == nil {
                log.Fatal(err)
        }
        prep.Close()
}

func coreCheck(cores *int) {
        if *cores > runtime.NumCPU() || *cores <= 0 {
                log.Fatal(`You don't have that many cores.... you can use up to `,
                        runtime.NumCPU())
        } else {
                runtime.GOMAXPROCS(*cores)
                if verbose == 1 {
                        fmt.Println("Using " + string(runtime.GOMAXPROCS(*cores)) + " cores")
                }
        }
}

var verbose int

func main() {
        fileOpt := flag.String("file", "xmlFile", "file to parse into db")
        dirOpt := flag.String("dir", "directory", "dir of xml files")
        coreOpt := flag.Int("cores", 1, "Number of Cores to use")
        verboseOpt := flag.Int("verbose", 0, "Verbose level 0,1,2")
        userOpt := flag.String("user", "postgres", "User for Postgres")
        passOpt := flag.String("pass", "", "Password for Postgres user")
        dbOpt := flag.String("db", "gotest", "DB to use")
        tableOpt := flag.String("table", "internal_network", "Table to use")
        sslOpt := flag.String("ssl", "disable", "Enable or Disable")
        flag.Parse()
        coreCheck(coreOpt)
        verbose = *verboseOpt
        var con_opts string
        if *passOpt == "" {
                con_opts = "user=" + *userOpt +
                        " dbname=" + *dbOpt + " sslmode=" + *sslOpt
        } else {
                con_opts = "user=" + *userOpt + " password=" + *passOpt +
                        " dbname=" + *dbOpt + " sslmode=" + *sslOpt
        }
        db, err := sql.Open("postgres", con_opts)
        if err != nil {
                fmt.Println("Error Connecting:", err)
                return
        }
        defer db.Close()
        txn, err := db.Begin()
        if err != nil {
                log.Fatal(err)
        }
        //dropTable(db, tableOpt)
        createTable(db, tableOpt)
        prep, err := txn.Prepare(`copy ` + *tableOpt + ` (host, mac_address, netbios,
                fqdn, os_name, plugin_name, plugin_id, severity, cve,
                risk, description, solution, synopsis, plugin_output,
                see_also, exploit_available, exploit_ease, metasploit_framework,
                metasploit_name, canvas_framework, core_framework,
                exploited_malware, cvss, month, year) from stdin`)
        if err != nil {
                log.Fatal(err)
        }
        file := *fileOpt
        dir := *dirOpt
        fmt.Println("Importing..")
        var wg sync.WaitGroup
        if file != "xmlFile" {
                xmlFile, err := os.Open(file)
                defer xmlFile.Close()
                if err != nil {
                        log.Fatal(err)
                        return
                }
                wg.Add(1)
                xmlParse(&wg, xmlFile, prep, db)
        } else if dir != "directory" {
                files, _ := filepath.Glob(dir + "/*")
                for _, file := range files {
                        xmlFile, err := os.Open(file)
                        defer xmlFile.Close()
                        if err != nil {
                                log.Fatal(err)
                                return
                        }
                        if verbose == 1 {
                                fmt.Printf("Parsing %s\n", file)
                        }
                        wg.Add(1)
                        go xmlParse(&wg, xmlFile, prep, db)
                }

        }
        wg.Wait()
        if err != nil {
                log.Fatal(err)
        }
        _, err = prep.Exec()
        if err != nil {
                log.Fatal(err)
        }
        err = prep.Close()
        if err != nil {
                log.Fatal(err)
        }
        err = txn.Commit()
        if err != nil {
                log.Fatal(err)
        }
}
<|endoftext|>


<commit_before>// (c) 2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package health

import (
	"errors"
	"time"

	"github.com/ava-labs/avalanchego/utils/rpc"
)

var errInvalidNumberOfChecks = errors.New("expected at least 1 check attempt")

// Client for Avalanche Health API Endpoint
type Client struct {
	requester rpc.EndpointRequester
}

// Result represents the output of a health check execution.
type Result struct {
	// the details of task Result - may be nil
	Details interface{} `json:"message,omitempty"`
	// the error returned from a failed health check - an empty string when successful
	Error string `json:"error,omitempty"`
	// the time of the last health check
	Timestamp time.Time `json:"timestamp"`
	// the execution duration of the last check
	Duration time.Duration `json:"duration,omitempty"`
	// the number of failures that occurred in a row
	ContiguousFailures int64 `json:"contiguousFailures"`
	// the time of the initial transitional failure
	TimeOfFirstFailure *time.Time `json:"timeOfFirstFailure"`
}

type APIHealthClientReply struct {
	Checks  map[string]Result `json:"checks"`
	Healthy bool              `json:"healthy"`
}

// NewClient returns a client to interact with Health API endpoint
func NewClient(uri string, requestTimeout time.Duration) *Client {
	return &Client{
		requester: rpc.NewEndpointRequester(uri, "/ext/health", "health", requestTimeout),
	}
}

// Health returns a health check on the Avalanche node
func (c *Client) Health() (*APIHealthClientReply, error) {
	res := &APIHealthClientReply{}
	err := c.requester.SendRequest("health", struct{}{}, res)
	return res, err
}

// AwaitHealthy queries the Health endpoint [checks] times, with a pause of
// [interval] in between checks and returns early if Health returns healthy
func (c *Client) AwaitHealthy(checks int, interval time.Duration) (bool, error) {
	if checks < 1 {
		return false, errInvalidNumberOfChecks
	}

	// Check health once outside the loop to avoid sleeping unnecessarily.
	res, err := c.Health()
	if err == nil && res.Healthy {
		return true, nil
	}

	for i := 1; i < checks; i++ {
		time.Sleep(interval)

		res, err = c.Health()
		if err == nil && res.Healthy {
			return true, nil
		}
	}
	return false, err
}
<commit_msg>fix health client by allowing error struct to be unmarshaled into an error struct<commit_after>// (c) 2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package health

import (
	"errors"
	"time"

	"github.com/ava-labs/avalanchego/utils/rpc"
)

var errInvalidNumberOfChecks = errors.New("expected at least 1 check attempt")

// Client for Avalanche Health API Endpoint
type Client struct {
	requester rpc.EndpointRequester
}

type ErrorMsg struct {
	Message string `json:"message"`
}

// Result represents the output of a health check execution.
type Result struct {
	// the details of task Result - may be nil
	Details interface{} `json:"message,omitempty"`
	// the error returned from a failed health check - an empty string when successful
	Error ErrorMsg `json:"error,omitempty"`
	// the time of the last health check
	Timestamp time.Time `json:"timestamp"`
	// the execution duration of the last check
	Duration time.Duration `json:"duration,omitempty"`
	// the number of failures that occurred in a row
	ContiguousFailures int64 `json:"contiguousFailures"`
	// the time of the initial transitional failure
	TimeOfFirstFailure *time.Time `json:"timeOfFirstFailure"`
}

type APIHealthClientReply struct {
	Checks  map[string]Result `json:"checks"`
	Healthy bool              `json:"healthy"`
}

// NewClient returns a client to interact with Health API endpoint
func NewClient(uri string, requestTimeout time.Duration) *Client {
	return &Client{
		requester: rpc.NewEndpointRequester(uri, "/ext/health", "health", requestTimeout),
	}
}

// Health returns a health check on the Avalanche node
func (c *Client) Health() (*APIHealthClientReply, error) {
	res := &APIHealthClientReply{}
	err := c.requester.SendRequest("health", struct{}{}, res)
	return res, err
}

// AwaitHealthy queries the Health endpoint [checks] times, with a pause of
// [interval] in between checks and returns early if Health returns healthy
func (c *Client) AwaitHealthy(checks int, interval time.Duration) (bool, error) {
	if checks < 1 {
		return false, errInvalidNumberOfChecks
	}

	// Check health once outside the loop to avoid sleeping unnecessarily.
	res, err := c.Health()
	if err == nil && res.Healthy {
		return true, nil
	}

	for i := 1; i < checks; i++ {
		time.Sleep(interval)

		res, err = c.Health()
		if err == nil && res.Healthy {
			return true, nil
		}
	}
	return false, err
}
<|endoftext|>


<commit_before>package paymentmodels

import (
	"socialapi/workers/payment/paymenterrors"
	"time"

	"github.com/koding/bongo"
)

type Subscription struct {
	Id int64 `json:"id,string"`

	// Id of subscription in 3rd payment provider like Stripe.
	ProviderSubscriptionId string `json:"providerSubscriptionId"`

	// Name of provider. Enum:
	//    'stripe', 'paypal'
	Provider string `json:"provider"`

	// Token that was fetched from provider from client after successful
	// credit card validation, which is then used to identify the user.
	ProviderToken string `json:"providerToken"`

	// Account the subscription belongs to, internal account id.
	CustomerId int64 `json:"customerId,string"`

	// Plan the subscription belongs to, internal plan id.
	PlanId int64 `json:"planId,string"`

	// State of the subscription. Enum:
	//    'active', 'expired'
	State string `json:"state"`

	// Cost of plan in cents.
	AmountInCents uint64 `json:"amountInCents"`

	// Timestamps
	CreatedAt          time.Time `json:"createdAt"`
	UpdatedAt          time.Time `json:"updatedAt" `
	ExpiredAt          time.Time `json:"expiredAt"`
	CanceledAt         time.Time `json:"canceled_at"`
	CurrentPeriodStart time.Time `json:"current_period_start"`
	CurrentPeriodEnd   time.Time `json:"current_period_end"`
}

var (
	SubscriptionStateActive = "active"

	// this is when user manually cancels
	SubscriptionStateCanceled = "canceled"

	// this is when user fails to pay
	SubscriptionStateExpired = "expired"
)

func (s *Subscription) UpdateInvoiceCreated(amountInCents uint64, planId, periodStart, periodEnd int64) error {
	if s.PlanId == planId {
		return ErrUpdatingToSamePlan
	}

	s.PlanId = planId
	s.AmountInCents = amountInCents
	s.CurrentPeriodStart = time.Unix(periodStart, 0)
	s.CurrentPeriodEnd = time.Unix(periodEnd, 0)

	// when user downgrades to non-free plan, we set `CanceledAt` till the
	// end of billing cycle and when `invoice.created` is fired we update
	// null this field
	s.CanceledAt = time.Time{}

	return bongo.B.Update(s)
}

func (s *Subscription) UpdatePlan(planId int64, amountInCents uint64) error {
	if s.PlanId == planId {
		return ErrUpdatingToSamePlan
	}

	s.PlanId = planId
	s.AmountInCents = amountInCents

	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) UpdateState(state string) error {
	s.State = state
	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) Cancel() error {
	s.CanceledAt = time.Now()
	s.State = SubscriptionStateCanceled

	return bongo.B.Update(s)
}

func (s *Subscription) Expire() error {
	s.ExpiredAt = time.Now()
	s.State = SubscriptionStateExpired

	return bongo.B.Update(s)
}

func (s *Subscription) UpdateTimeForDowngrade(t time.Time) error {
	s.CanceledAt = t
	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) ByProviderId(providerId, provider string) error {
	selector := map[string]interface{}{
		"provider_subscription_id": providerId,
		"provider":                 provider,
	}
	err := s.Find(selector)

	return err
}

func (s *Subscription) ByCustomerIdAndState(customerId int64, state string) error {
	selector := map[string]interface{}{
		"customer_id": customerId,
		"state":       state,
	}

	err := s.Find(selector)
	if err == bongo.RecordNotFound {
		return paymenterrors.ErrCustomerNotSubscribedToAnyPlans
	}

	return err
}

func (s *Subscription) ByCanceledAtGte(t time.Time) ([]Subscription, error) {
	subscriptions := []Subscription{}

	err := bongo.B.DB.
		Table(s.BongoName()).
		Where(
		"canceled_at < ? AND canceled_at != ?", t, time.Time{},
	).Find(&subscriptions).Error

	return subscriptions, err
}

func (s *Subscription) ByExpiredAtAndNotExpired(t time.Time) ([]Subscription, error) {
	subscriptions := []Subscription{}

	err := bongo.B.DB.
		Table(s.BongoName()).
		Where(
		"expired_at < ? AND expired_at != ? AND state == ?",
		t, time.Time{}, SubscriptionStateActive,
	).Find(&subscriptions).Error

	return subscriptions, err
}

func (s *Subscription) UpdateToExpireTime(t time.Time) error {
	s.ExpiredAt = t
	s.CanceledAt = t

	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) UpdateCurrentPeriods(start, end time.Time) error {
	s.CurrentPeriodStart = start
	s.CurrentPeriodEnd = end

	return bongo.B.Update(s)
}
<commit_msg>payment: minor fix<commit_after>package paymentmodels

import (
	"socialapi/workers/payment/paymenterrors"
	"time"

	"github.com/koding/bongo"
)

type Subscription struct {
	Id int64 `json:"id,string"`

	// Id of subscription in 3rd payment provider like Stripe.
	ProviderSubscriptionId string `json:"providerSubscriptionId"`

	// Name of provider. Enum:
	//    'stripe', 'paypal'
	Provider string `json:"provider"`

	// Token that was fetched from provider from client after successful
	// credit card validation, which is then used to identify the user.
	ProviderToken string `json:"providerToken"`

	// Account the subscription belongs to, internal account id.
	CustomerId int64 `json:"customerId,string"`

	// Plan the subscription belongs to, internal plan id.
	PlanId int64 `json:"planId,string"`

	// State of the subscription. Enum:
	//    'active', 'expired'
	State string `json:"state"`

	// Cost of plan in cents.
	AmountInCents uint64 `json:"amountInCents"`

	// Timestamps
	CreatedAt          time.Time `json:"createdAt"`
	UpdatedAt          time.Time `json:"updatedAt" `
	ExpiredAt          time.Time `json:"expiredAt"`
	CanceledAt         time.Time `json:"canceled_at"`
	CurrentPeriodStart time.Time `json:"current_period_start"`
	CurrentPeriodEnd   time.Time `json:"current_period_end"`
}

var (
	SubscriptionStateActive = "active"

	// this is when user manually cancels
	SubscriptionStateCanceled = "canceled"

	// this is when user fails to pay
	SubscriptionStateExpired = "expired"
)

func (s *Subscription) UpdateInvoiceCreated(amountInCents uint64, planId, periodStart, periodEnd int64) error {
	if s.PlanId == planId {
		return ErrUpdatingToSamePlan
	}

	s.PlanId = planId
	s.AmountInCents = amountInCents
	s.CurrentPeriodStart = time.Unix(periodStart, 0)
	s.CurrentPeriodEnd = time.Unix(periodEnd, 0)

	// when user downgrades to non-free plan, we set `CanceledAt` till the
	// end of billing cycle and when `invoice.created` is fired we update
	// null this field
	s.CanceledAt = time.Time{}

	return bongo.B.Update(s)
}

func (s *Subscription) UpdatePlan(planId int64, amountInCents uint64) error {
	if s.PlanId == planId {
		return ErrUpdatingToSamePlan
	}

	s.PlanId = planId
	s.AmountInCents = amountInCents

	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) UpdateState(state string) error {
	s.State = state
	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) Cancel() error {
	s.CanceledAt = time.Now()
	s.State = SubscriptionStateCanceled

	return bongo.B.Update(s)
}

func (s *Subscription) Expire() error {
	s.ExpiredAt = time.Now()
	s.State = SubscriptionStateExpired

	return bongo.B.Update(s)
}

func (s *Subscription) UpdateTimeForDowngrade(t time.Time) error {
	s.CanceledAt = t
	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) ByProviderId(providerId, provider string) error {
	selector := map[string]interface{}{
		"provider_subscription_id": providerId,
		"provider":                 provider,
	}
	err := s.Find(selector)

	return err
}

func (s *Subscription) ByCustomerIdAndState(customerId int64, state string) error {
	selector := map[string]interface{}{
		"customer_id": customerId,
		"state":       state,
	}

	err := s.Find(selector)
	if err == bongo.RecordNotFound {
		return paymenterrors.ErrCustomerNotSubscribedToAnyPlans
	}

	return err
}

func (s *Subscription) ByCanceledAtGte(t time.Time) ([]Subscription, error) {
	subscriptions := []Subscription{}

	err := bongo.B.DB.
		Table(s.BongoName()).
		Where(
		"canceled_at < ? AND canceled_at != ?", t, time.Time{},
	).Find(&subscriptions).Error

	return subscriptions, err
}

func (s *Subscription) ByExpiredAtAndNotExpired(t time.Time) ([]Subscription, error) {
	subscriptions := []Subscription{}

	err := bongo.B.DB.
		Table(s.BongoName()).
		Where(
		"expired_at < ? AND expired_at != ? AND state = ?",
		t, time.Time{}, SubscriptionStateActive,
	).Find(&subscriptions).Error

	return subscriptions, err
}

func (s *Subscription) UpdateToExpireTime(t time.Time) error {
	s.ExpiredAt = t
	s.CanceledAt = t

	err := bongo.B.Update(s)

	return err
}

func (s *Subscription) UpdateCurrentPeriods(start, end time.Time) error {
	s.CurrentPeriodStart = start
	s.CurrentPeriodEnd = end

	return bongo.B.Update(s)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"text/template"
)

type person struct {
	Fname string
	Mname string
	Lname string
	Age   int
}

const dtmpl = `
<!DOCTYPE html>

	<meta charset="UTF-8">
	<title>dynamic-templateloop</title>

	{{range $key, $value := .}}
	Tenant: {{ $key }}

	{{range $value}}
	Name :{{.Fname}} {{.Mname}}.{{.Lname}}

	Age :{{ .Age }}

	{{end}}
	{{end}}

`

func main() {
	p1 := person{"Praveen", "Kumar", "K", 36}
	p2 := person{"Srinivasa", "Reddy", "M", 36}
	p3 := person{"Mahesh", "Reddy", "M", 36}
	persons := map[string][]person{}
	persons["tenant1"]=[]person{p1,p2}
	persons["tenant2"]=[]person{p3}

	tmpl, err := template.New("dynamic-loop").Parse(dtmpl)
	if err != nil {
		panic(err)
	}

	//var tpl bytes.Buffer
	//tmpl.Execute(&tpl, persons)
	//fmt.Println(tpl.String())
  http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		tmpl.Execute(w, persons)
	})
	http.ListenAndServe(":8080", nil)
}
<commit_msg>Update dynamic-template-loop-map-of-slice.go<commit_after>package main

import (
	"net/http"
	"text/template"
)

type person struct {
	Fname string
	Mname string
	Lname string
	Age   int
}

const dtmpl = `
<!DOCTYPE html>

	<meta charset="UTF-8">
	<title>dynamic-templateloop</title>

	{{range $key, $value := .}}
	Tenant: {{ $key }}

	{{range $value}}
	Name :{{.Fname}} {{.Mname}}.{{.Lname}}

	Age :{{ .Age }}

	{{end}}
	{{end}}

`

func main() {
	p1 := person{"Praveen", "Kumar", "K", 36}
	p2 := person{"Srinivasa", "Reddy", "M", 36}
	p3 := person{"Mahesh", "Reddy", "M", 36}
	persons := map[string][]person{}
	persons["tenant1"] = []person{p1, p2}
	persons["tenant2"] = []person{p3}

	tmpl, err := template.New("dynamic-loop").Parse(dtmpl)
	if err != nil {
		panic(err)
	}

	//var tpl bytes.Buffer
	//tmpl.Execute(&tpl, persons)
	//fmt.Println(tpl.String())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		tmpl.Execute(w, persons)
	})
	http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>package cortexbot

import (
	"fmt"
	"testing"

	"github.com/ilyaglow/go-cortex"
)

func TestObservableConstructor(t *testing.T) {
	var cases = []struct {
		data     string
		tlp      int
		dataType string
	}{
		{"1.1.1.1", 0, "ip"},
		{"https://subdomain.domain.com/route/query?param=val", 1, "url"},
		{"subdomain.domain.com", 2, "domain"},
		{"email.address@domain.com", 3, "mail"},
		{"email+address@domain.com", 3, "mail"},
		{"a94a8fe5ccb19ba61c4c0873d391e987982fbbd3", 0, "hash"},
		{"randomdata", 1, "other"},
	}

	for _, c := range cases {
		a := newArtifact(c.data, c.tlp)
		if a.Description() != c.data {
			t.Errorf("need %s, got %s", c.data, a.Description())
		}

		if a.Type() != c.dataType {
			t.Errorf("need %s, got %s", c.dataType, a.Type())
		}
	}
}

func TestBuildTaxonomies(t *testing.T) {
	var txs []cortex.Taxonomy

	tx1 := cortex.Taxonomy{
		Predicate: "Predicate1",
		Namespace: "Namespace1",
		Value:     "Value1",
		Level:     "safe",
	}
	txs = append(txs, tx1)

	tx2 := cortex.Taxonomy{
		Predicate: "Predicate2",
		Namespace: "Namespace2",
		Value:     "Value2",
		Level:     "info",
	}
	txs = append(txs, tx2)

	expected := fmt.Sprintf("%s:%s = %s, %s:%s = %s",
		tx1.Namespace, tx1.Predicate, tx1.Value,
		tx2.Namespace, tx2.Predicate, tx2.Value,
	)
	if buildTaxonomies(txs) != expected {
		t.Errorf("Wrong taxonomies format:\n%s\n%s", expected, buildTaxonomies(txs))
	}
}
<commit_msg>Fix taxonomies test<commit_after>package cortexbot

import (
	"fmt"
	"testing"

	"github.com/ilyaglow/go-cortex"
)

func TestObservableConstructor(t *testing.T) {
	var cases = []struct {
		data     string
		tlp      int
		dataType string
	}{
		{"1.1.1.1", 0, "ip"},
		{"https://subdomain.domain.com/route/query?param=val", 1, "url"},
		{"subdomain.domain.com", 2, "domain"},
		{"email.address@domain.com", 3, "mail"},
		{"email+address@domain.com", 3, "mail"},
		{"a94a8fe5ccb19ba61c4c0873d391e987982fbbd3", 0, "hash"},
		{"randomdata", 1, "other"},
	}

	for _, c := range cases {
		a := newArtifact(c.data, c.tlp)
		if a.Description() != c.data {
			t.Errorf("need %s, got %s", c.data, a.Description())
		}

		if a.Type() != c.dataType {
			t.Errorf("need %s, got %s", c.dataType, a.Type())
		}
	}
}

func TestBuildTaxonomies(t *testing.T) {
	var txs []cortex.Taxonomy

	tx1 := cortex.Taxonomy{
		Predicate: "Predicate1",
		Namespace: "Namespace1",
		Value:     "Value1",
		Level:     "safe",
	}
	txs = append(txs, tx1)

	tx2 := cortex.Taxonomy{
		Predicate: "Predicate2",
		Namespace: "Namespace2",
		Value:     "Value2",
		Level:     "info",
	}
	txs = append(txs, tx2)

	expected := fmt.Sprintf("%s:%s = %s\n%s:%s = %s",
		tx1.Namespace, tx1.Predicate, tx1.Value,
		tx2.Namespace, tx2.Predicate, tx2.Value,
	)
	if buildTaxonomies(txs) != expected {
		t.Errorf("Wrong taxonomies format:\n%s\n%s", expected, buildTaxonomies(txs))
	}
}
<|endoftext|>


<commit_before>package quic

import (
	"bytes"
	"errors"
	"fmt"

	"github.com/lucas-clemente/quic-go/frames"
	"github.com/lucas-clemente/quic-go/handshake"
	"github.com/lucas-clemente/quic-go/protocol"
)

type packedPacket struct {
	number protocol.PacketNumber
	raw    []byte
	frames []frames.Frame
}

type packetPacker struct {
	connectionID protocol.ConnectionID
	version      protocol.VersionNumber
	cryptoSetup  *handshake.CryptoSetup

	packetNumberGenerator *packetNumberGenerator

	connectionParametersManager *handshake.ConnectionParametersManager

	streamFramer  *streamFramer
	controlFrames []frames.Frame
}

func newPacketPacker(connectionID protocol.ConnectionID, cryptoSetup *handshake.CryptoSetup, connectionParametersHandler *handshake.ConnectionParametersManager, streamFramer *streamFramer, version protocol.VersionNumber) *packetPacker {
	return &packetPacker{
		cryptoSetup:                 cryptoSetup,
		connectionID:                connectionID,
		connectionParametersManager: connectionParametersHandler,
		version:                     version,
		streamFramer:                streamFramer,
		packetNumberGenerator:       newPacketNumberGenerator(protocol.SkipPacketAveragePeriodLength),
	}
}

func (p *packetPacker) PackConnectionClose(frame *frames.ConnectionCloseFrame, leastUnacked protocol.PacketNumber) (*packedPacket, error) {
	return p.packPacket(nil, []frames.Frame{frame}, leastUnacked, true, false)
}

func (p *packetPacker) PackPacket(stopWaitingFrame *frames.StopWaitingFrame, controlFrames []frames.Frame, leastUnacked protocol.PacketNumber, maySendOnlyAck bool) (*packedPacket, error) {
	return p.packPacket(stopWaitingFrame, controlFrames, leastUnacked, false, maySendOnlyAck)
}

func (p *packetPacker) packPacket(stopWaitingFrame *frames.StopWaitingFrame, controlFrames []frames.Frame, leastUnacked protocol.PacketNumber, onlySendOneControlFrame, maySendOnlyAck bool) (*packedPacket, error) {
	if len(controlFrames) > 0 {
		p.controlFrames = append(p.controlFrames, controlFrames...)
	}

	currentPacketNumber := p.packetNumberGenerator.Peek()

	// cryptoSetup needs to be locked here, so that the AEADs are not changed between
	// calling DiversificationNonce() and Seal().
	p.cryptoSetup.LockForSealing()
	defer p.cryptoSetup.UnlockForSealing()

	packetNumberLen := protocol.GetPacketNumberLengthForPublicHeader(currentPacketNumber, leastUnacked)
	responsePublicHeader := &PublicHeader{
		ConnectionID:         p.connectionID,
		PacketNumber:         currentPacketNumber,
		PacketNumberLen:      packetNumberLen,
		TruncateConnectionID: p.connectionParametersManager.TruncateConnectionID(),
		DiversificationNonce: p.cryptoSetup.DiversificationNonce(),
	}

	publicHeaderLength, err := responsePublicHeader.GetLength()
	if err != nil {
		return nil, err
	}

	if stopWaitingFrame != nil {
		stopWaitingFrame.PacketNumber = currentPacketNumber
		stopWaitingFrame.PacketNumberLen = packetNumberLen
	}

	var payloadFrames []frames.Frame
	if onlySendOneControlFrame {
		payloadFrames = []frames.Frame{controlFrames[0]}
	} else {
		payloadFrames, err = p.composeNextPacket(stopWaitingFrame, publicHeaderLength)
		if err != nil {
			return nil, err
		}
	}

	// Check if we have enough frames to send
	if len(payloadFrames) == 0 {
		return nil, nil
	}
	// Don't send out packets that only contain a StopWaitingFrame
	if !onlySendOneControlFrame && len(payloadFrames) == 1 && stopWaitingFrame != nil {
		return nil, nil
	}
	// Don't send out packets that only contain an ACK (plus optional STOP_WAITING), if requested
	if !maySendOnlyAck {
		if len(payloadFrames) == 1 {
			if _, ok := payloadFrames[0].(*frames.AckFrame); ok {
				return nil, nil
			}
		} else if len(payloadFrames) == 2 && stopWaitingFrame != nil {
			if _, ok := payloadFrames[1].(*frames.AckFrame); ok {
				return nil, nil
			}
		}
	}

	raw := getPacketBuffer()
	buffer := bytes.NewBuffer(raw)

	if responsePublicHeader.WritePublicHeader(buffer, p.version) != nil {
		return nil, err
	}

	payloadStartIndex := buffer.Len()

	for _, frame := range payloadFrames {
		err := frame.Write(buffer, p.version)
		if err != nil {
			return nil, err
		}
	}

	if protocol.ByteCount(buffer.Len()+12) > protocol.MaxPacketSize {
		return nil, errors.New("PacketPacker BUG: packet too large")
	}

	raw = raw[0:buffer.Len()]
	p.cryptoSetup.Seal(raw[payloadStartIndex:payloadStartIndex], raw[payloadStartIndex:], currentPacketNumber, raw[:payloadStartIndex])
	raw = raw[0 : buffer.Len()+12]

	num := p.packetNumberGenerator.Pop()
	if num != currentPacketNumber {
		return nil, errors.New("PacketPacker BUG: Peeked and Popped packet numbers do not match.")
	}

	return &packedPacket{
		number: currentPacketNumber,
		raw:    raw,
		frames: payloadFrames,
	}, nil
}

func (p *packetPacker) composeNextPacket(stopWaitingFrame *frames.StopWaitingFrame, publicHeaderLength protocol.ByteCount) ([]frames.Frame, error) {
	var payloadLength protocol.ByteCount
	var payloadFrames []frames.Frame

	maxFrameSize := protocol.MaxFrameAndPublicHeaderSize - publicHeaderLength

	if stopWaitingFrame != nil {
		payloadFrames = append(payloadFrames, stopWaitingFrame)
		minLength, err := stopWaitingFrame.MinLength(p.version)
		if err != nil {
			return nil, err
		}
		payloadLength += minLength
	}

	for len(p.controlFrames) > 0 {
		frame := p.controlFrames[0]
		minLength, _ := frame.MinLength(p.version) // controlFrames does not contain any StopWaitingFrames. So it will *never* return an error
		if payloadLength+minLength > maxFrameSize {
			break
		}
		payloadFrames = append(payloadFrames, frame)
		payloadLength += minLength
		p.controlFrames = p.controlFrames[1:]
	}

	if payloadLength > maxFrameSize {
		return nil, fmt.Errorf("Packet Packer BUG: packet payload (%d) too large (%d)", payloadLength, maxFrameSize)
	}

	// temporarily increase the maxFrameSize by 2 bytes
	// this leads to a properly sized packet in all cases, since we do all the packet length calculations with StreamFrames that have the DataLen set
	// however, for the last StreamFrame in the packet, we can omit the DataLen, thus saving 2 bytes and yielding a packet of exactly the correct size
	maxFrameSize += 2

	fs := p.streamFramer.PopStreamFrames(maxFrameSize - payloadLength)
	if len(fs) != 0 {
		fs[len(fs)-1].DataLenPresent = false
	}

	// TODO: Simplify
	for _, f := range fs {
		payloadFrames = append(payloadFrames, f)
	}

	for b := p.streamFramer.PopBlockedFrame(); b != nil; b = p.streamFramer.PopBlockedFrame() {
		p.controlFrames = append(p.controlFrames, b)
	}

	return payloadFrames, nil
}

func (p *packetPacker) QueueControlFrameForNextPacket(f frames.Frame) {
	p.controlFrames = append(p.controlFrames, f)
}
<commit_msg>use packetPacker.controlFrames as stack, decreasing allocs<commit_after>package quic

import (
	"bytes"
	"errors"
	"fmt"

	"github.com/lucas-clemente/quic-go/frames"
	"github.com/lucas-clemente/quic-go/handshake"
	"github.com/lucas-clemente/quic-go/protocol"
)

type packedPacket struct {
	number protocol.PacketNumber
	raw    []byte
	frames []frames.Frame
}

type packetPacker struct {
	connectionID protocol.ConnectionID
	version      protocol.VersionNumber
	cryptoSetup  *handshake.CryptoSetup

	packetNumberGenerator *packetNumberGenerator

	connectionParametersManager *handshake.ConnectionParametersManager

	streamFramer  *streamFramer
	controlFrames []frames.Frame
}

func newPacketPacker(connectionID protocol.ConnectionID, cryptoSetup *handshake.CryptoSetup, connectionParametersHandler *handshake.ConnectionParametersManager, streamFramer *streamFramer, version protocol.VersionNumber) *packetPacker {
	return &packetPacker{
		cryptoSetup:                 cryptoSetup,
		connectionID:                connectionID,
		connectionParametersManager: connectionParametersHandler,
		version:                     version,
		streamFramer:                streamFramer,
		packetNumberGenerator:       newPacketNumberGenerator(protocol.SkipPacketAveragePeriodLength),
	}
}

func (p *packetPacker) PackConnectionClose(frame *frames.ConnectionCloseFrame, leastUnacked protocol.PacketNumber) (*packedPacket, error) {
	return p.packPacket(nil, []frames.Frame{frame}, leastUnacked, true, false)
}

func (p *packetPacker) PackPacket(stopWaitingFrame *frames.StopWaitingFrame, controlFrames []frames.Frame, leastUnacked protocol.PacketNumber, maySendOnlyAck bool) (*packedPacket, error) {
	return p.packPacket(stopWaitingFrame, controlFrames, leastUnacked, false, maySendOnlyAck)
}

func (p *packetPacker) packPacket(stopWaitingFrame *frames.StopWaitingFrame, controlFrames []frames.Frame, leastUnacked protocol.PacketNumber, onlySendOneControlFrame, maySendOnlyAck bool) (*packedPacket, error) {
	if len(controlFrames) > 0 {
		p.controlFrames = append(p.controlFrames, controlFrames...)
	}

	currentPacketNumber := p.packetNumberGenerator.Peek()

	// cryptoSetup needs to be locked here, so that the AEADs are not changed between
	// calling DiversificationNonce() and Seal().
	p.cryptoSetup.LockForSealing()
	defer p.cryptoSetup.UnlockForSealing()

	packetNumberLen := protocol.GetPacketNumberLengthForPublicHeader(currentPacketNumber, leastUnacked)
	responsePublicHeader := &PublicHeader{
		ConnectionID:         p.connectionID,
		PacketNumber:         currentPacketNumber,
		PacketNumberLen:      packetNumberLen,
		TruncateConnectionID: p.connectionParametersManager.TruncateConnectionID(),
		DiversificationNonce: p.cryptoSetup.DiversificationNonce(),
	}

	publicHeaderLength, err := responsePublicHeader.GetLength()
	if err != nil {
		return nil, err
	}

	if stopWaitingFrame != nil {
		stopWaitingFrame.PacketNumber = currentPacketNumber
		stopWaitingFrame.PacketNumberLen = packetNumberLen
	}

	var payloadFrames []frames.Frame
	if onlySendOneControlFrame {
		payloadFrames = []frames.Frame{controlFrames[0]}
	} else {
		payloadFrames, err = p.composeNextPacket(stopWaitingFrame, publicHeaderLength)
		if err != nil {
			return nil, err
		}
	}

	// Check if we have enough frames to send
	if len(payloadFrames) == 0 {
		return nil, nil
	}
	// Don't send out packets that only contain a StopWaitingFrame
	if !onlySendOneControlFrame && len(payloadFrames) == 1 && stopWaitingFrame != nil {
		return nil, nil
	}
	// Don't send out packets that only contain an ACK (plus optional STOP_WAITING), if requested
	if !maySendOnlyAck {
		if len(payloadFrames) == 1 {
			if _, ok := payloadFrames[0].(*frames.AckFrame); ok {
				return nil, nil
			}
		} else if len(payloadFrames) == 2 && stopWaitingFrame != nil {
			if _, ok := payloadFrames[1].(*frames.AckFrame); ok {
				return nil, nil
			}
		}
	}

	raw := getPacketBuffer()
	buffer := bytes.NewBuffer(raw)

	if responsePublicHeader.WritePublicHeader(buffer, p.version) != nil {
		return nil, err
	}

	payloadStartIndex := buffer.Len()

	for _, frame := range payloadFrames {
		err := frame.Write(buffer, p.version)
		if err != nil {
			return nil, err
		}
	}

	if protocol.ByteCount(buffer.Len()+12) > protocol.MaxPacketSize {
		return nil, errors.New("PacketPacker BUG: packet too large")
	}

	raw = raw[0:buffer.Len()]
	p.cryptoSetup.Seal(raw[payloadStartIndex:payloadStartIndex], raw[payloadStartIndex:], currentPacketNumber, raw[:payloadStartIndex])
	raw = raw[0 : buffer.Len()+12]

	num := p.packetNumberGenerator.Pop()
	if num != currentPacketNumber {
		return nil, errors.New("PacketPacker BUG: Peeked and Popped packet numbers do not match.")
	}

	return &packedPacket{
		number: currentPacketNumber,
		raw:    raw,
		frames: payloadFrames,
	}, nil
}

func (p *packetPacker) composeNextPacket(stopWaitingFrame *frames.StopWaitingFrame, publicHeaderLength protocol.ByteCount) ([]frames.Frame, error) {
	var payloadLength protocol.ByteCount
	var payloadFrames []frames.Frame

	maxFrameSize := protocol.MaxFrameAndPublicHeaderSize - publicHeaderLength

	if stopWaitingFrame != nil {
		payloadFrames = append(payloadFrames, stopWaitingFrame)
		minLength, err := stopWaitingFrame.MinLength(p.version)
		if err != nil {
			return nil, err
		}
		payloadLength += minLength
	}

	for len(p.controlFrames) > 0 {
		frame := p.controlFrames[len(p.controlFrames)-1]
		minLength, _ := frame.MinLength(p.version) // controlFrames does not contain any StopWaitingFrames. So it will *never* return an error
		if payloadLength+minLength > maxFrameSize {
			break
		}
		payloadFrames = append(payloadFrames, frame)
		payloadLength += minLength
		p.controlFrames = p.controlFrames[:len(p.controlFrames)-1]
	}

	if payloadLength > maxFrameSize {
		return nil, fmt.Errorf("Packet Packer BUG: packet payload (%d) too large (%d)", payloadLength, maxFrameSize)
	}

	// temporarily increase the maxFrameSize by 2 bytes
	// this leads to a properly sized packet in all cases, since we do all the packet length calculations with StreamFrames that have the DataLen set
	// however, for the last StreamFrame in the packet, we can omit the DataLen, thus saving 2 bytes and yielding a packet of exactly the correct size
	maxFrameSize += 2

	fs := p.streamFramer.PopStreamFrames(maxFrameSize - payloadLength)
	if len(fs) != 0 {
		fs[len(fs)-1].DataLenPresent = false
	}

	// TODO: Simplify
	for _, f := range fs {
		payloadFrames = append(payloadFrames, f)
	}

	for b := p.streamFramer.PopBlockedFrame(); b != nil; b = p.streamFramer.PopBlockedFrame() {
		p.controlFrames = append(p.controlFrames, b)
	}

	return payloadFrames, nil
}

func (p *packetPacker) QueueControlFrameForNextPacket(f frames.Frame) {
	p.controlFrames = append(p.controlFrames, f)
}
<|endoftext|>


<commit_before>package check

// GoCyclo is the check for the go cyclo command
type GoCyclo struct {
	Dir       string
	Filenames []string
}

// Name returns the name of the display name of the command
func (g GoCyclo) Name() string {
	return "gocyclo"
}

// Weight returns the weight this check has in the overall average
func (g GoCyclo) Weight() float64 {
	return .10
}

// Percentage returns the percentage of .go files that pass gofmt
func (g GoCyclo) Percentage() (float64, []FileSummary, error) {
	return GoTool(g.Dir, g.Filenames, []string{"gometalinter", "--deadline=180s", "--disable-all", "--enable=gocyclo", "--cyclo-over=15"})
}

// Description returns the description of GoCyclo
func (g GoCyclo) Description() string {
	return `Gocyclo calculates cyclomatic complexities of functions in Go source code.

The cyclomatic complexity of a function is calculated according to the following rules:

1 is the base complexity of a function
+1 for each 'if', 'for', 'case', '&&' or '||'`
}
<commit_msg>#184 add note in gocyclo description about threshold<commit_after>package check

// GoCyclo is the check for the go cyclo command
type GoCyclo struct {
	Dir       string
	Filenames []string
}

// Name returns the name of the display name of the command
func (g GoCyclo) Name() string {
	return "gocyclo"
}

// Weight returns the weight this check has in the overall average
func (g GoCyclo) Weight() float64 {
	return .10
}

// Percentage returns the percentage of .go files that pass gofmt
func (g GoCyclo) Percentage() (float64, []FileSummary, error) {
	return GoTool(g.Dir, g.Filenames, []string{"gometalinter", "--deadline=180s", "--disable-all", "--enable=gocyclo", "--cyclo-over=15"})
}

// Description returns the description of GoCyclo
func (g GoCyclo) Description() string {
	return `Gocyclo calculates cyclomatic complexities of functions in Go source code.

The cyclomatic complexity of a function is calculated according to the following rules:

1 is the base complexity of a function
+1 for each 'if', 'for', 'case', '&&' or '||'

Go Report Card warns on functions with cyclomatic complexity > 15.`
}
<|endoftext|>


<commit_before>// Package paccap provides as easy-to-use interface for capturing
// and analysing the packets.
package gopaccap

// TODO: Expose two functions - ReadPcap and LiveCapture
// TODO: implement cache for keeping src IP and destination port, expiring after 5 minutes

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/pcap"
)

/////////////////////////
// Global Declarations //
/////////////////////////
var banner string = `
$$$$$$$\                      $$$$$$\
$$  __$$\                    $$  __$$\
$$ |  $$ |$$$$$$\   $$$$$$$\ $$ /  \__| $$$$$$\   $$$$$$\
$$$$$$$  |\____$$\ $$  _____|$$ |       \____$$\ $$  __$$\
$$  ____/ $$$$$$$ |$$ /      $$ |       $$$$$$$ |$$ /  $$ |
$$ |     $$  __$$ |$$ |      $$ |  $$\ $$  __$$ |$$ |  $$ |
$$ |     \$$$$$$$ |\$$$$$$$\ \$$$$$$  |\$$$$$$$ |$$$$$$$  |
\__|      \_______| \_______| \______/  \_______|$$  ____/
                                                 $$ |
                                                 $$ |
                                                 \__|
`

var (
	snapshot_len int32         = 65535
	promiscuous  bool          = false
	timeout      time.Duration = -1 * time.Second
)

var logger = log.New()

/////////////////////////
//  Exposed Functions  //
/////////////////////////

// ReadPcap reads the pcap files from the specified path and
// logs the packet details.
func ReadPcap(filter, path string) {
	fmt.Println(banner)
	// TODO: Logging
	packetSource, err := readPackets(filter, path)

	if err != nil {
		log.Fatal(err)
	}
	for packet := range packetSource.Packets() {
		// DO something
		// Insert a delay for the information
		fmt.Println(packet)
	}
}

// LiveCapture attaches with the NIC specified and starts capturing
// the packets logging the packet details
func LiveCapture(filter, device string) {
	fmt.Println(banner)

	packetSource, err := getPackets(filter, device)
	if err != nil {
		log.Fatal(err)
	}

	for packet := range packetSource.Packets() {
		// DO something
		fmt.Println(packet)
	}
}

/////////////////////////
//   Helper Functions  //
/////////////////////////

// getPackets returns the PacketSource for analysing the packets.

func readPackets(filter, path string) (*gopacket.PacketSource, error) {
	handle, err := pcap.OpenOffline(path)
	if err != nil {
		return nil, err
	}
	// set the BPF filter
	err = handle.SetBPFFilter(filter)
	if err != nil {
		log.Fatal("Check the filter. Possibly there is a syntax error.")
	}
	packetSource := gopacket.NewPacketSource(
		handle,
		handle.LinkType(),
	)
	return packetSource, nil
}

func getPackets(filter, device string) (*gopacket.PacketSource, error) {
	handle, err := pcap.OpenLive(
		device,
		snapshot_len,
		promiscuous,
		timeout,
	)
	if err != nil {
		return nil, err
	}
	// set the filter to monitor HTTP traffic for now
	err = handle.SetBPFFilter(filter)
	if err != nil {
		log.Fatal("Check the filter. Possibly there is a syntax error.")
	}
	packetSource := gopacket.NewPacketSource(
		handle,
		handle.LinkType(),
	)
	return packetSource, nil
}
<commit_msg>src/dest logging and cache added<commit_after>// Package paccap provides as easy-to-use interface for capturing
// and analysing the packets.
package gopaccap

// TODO: Expose two functions - ReadPcap and LiveCapture
// TODO: implement cache for keeping src IP and destination port, expiring after 5 minutes

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/patrickmn/go-cache"
)

/////////////////////////
// Global Declarations //
/////////////////////////
var banner string = `
$$$$$$$\                      $$$$$$\
$$  __$$\                    $$  __$$\
$$ |  $$ |$$$$$$\   $$$$$$$\ $$ /  \__| $$$$$$\   $$$$$$\
$$$$$$$  |\____$$\ $$  _____|$$ |       \____$$\ $$  __$$\
$$  ____/ $$$$$$$ |$$ /      $$ |       $$$$$$$ |$$ /  $$ |
$$ |     $$  __$$ |$$ |      $$ |  $$\ $$  __$$ |$$ |  $$ |
$$ |     \$$$$$$$ |\$$$$$$$\ \$$$$$$  |\$$$$$$$ |$$$$$$$  |
\__|      \_______| \_______| \______/  \_______|$$  ____/
                                                 $$ |
                                                 $$ |
                                                 \__|
`

var (
	snapshot_len int32         = 65535
	promiscuous  bool          = false
	timeout      time.Duration = -1 * time.Second
)

var logger = log.New()

type paccap struct {
	ipcache *cache.Cache
}

/////////////////////////
//  Exposed Functions  //
/////////////////////////

// PacketCapture creates a new instance of the packet capturing
// device.
func PacketCapture() *paccap {
	// create a new cache with 5 minute expiration
	// and a purge time of 5 minutes
	ipc := cache.New(5*time.Minute, 5*time.Minute)
	// new instance of paccap
	pc := &paccap{ipcache: ipc}
	return pc
}

// ReadPcap reads the pcap files from the specified path and
// logs the packet details.
func (pc *paccap) ReadPcap(filter, path string) {
	fmt.Println(banner)
	logger.Infof("[PacCap ] Starting to read from pcap file...")
	fmt.Println()
	packetSource, err := readPackets(filter, path)
	if err != nil {
		log.Fatal(err)
	}
	for packet := range packetSource.Packets() {
		//TODO: Insert a delay for the streamlined info
		sip, dip := getIPAddresses(packet)
		spn, dpn := getPortAddresses(packet)
		logger.Infof("[PacCap ] Packet Captured! SOURCE %v:%v | DESTINATION %v:%v",
			sip, spn, dip, dpn)
		// adding src ip and destination to the cache
		pc.ipcache.Set(sip, dpn, cache.DefaultExpiration)
		fmt.Println()
	}
}

// LiveCapture attaches with the NIC specified and starts capturing
// the packets logging the packet details
func (pc *paccap) LiveCapture(filter, device string) {
	fmt.Println(banner)
	logger.Info("[PacCap ] Capturing NIC to read the packets...")
	fmt.Println()
	packetSource, err := getPackets(filter, device)
	if err != nil {
		log.Fatal(err)
	}
	for packet := range packetSource.Packets() {
		sip, dip := getIPAddresses(packet)
		spn, dpn := getPortAddresses(packet)
		logger.Infof("[PacCap ] Packet Captured! SOURCE %v:%v | DESTINATION %v:%v",
			sip, spn, dip, dpn)
		// observe cache hits
		// _, found := pc.ipcache.Get(sip)
		// if found {
		// 	logger.Infof("[PacCap ] Cache hit!")
		// }
		pc.ipcache.Set(sip, dpn, cache.DefaultExpiration)
		fmt.Println()
	}
}

/////////////////////////
//   Helper Functions  //
/////////////////////////

// readPackets
func readPackets(filter, path string) (*gopacket.PacketSource, error) {
	handle, err := pcap.OpenOffline(path)
	if err != nil {
		return nil, err
	}
	// set the BPF filter
	err = handle.SetBPFFilter(filter)
	if err != nil {
		log.Fatal("Check the filter. Possibly there is a syntax error.")
	}
	packetSource := gopacket.NewPacketSource(
		handle,
		handle.LinkType(),
	)
	return packetSource, nil
}

// getPackets returns the PacketSource for analysing the packets.
func getPackets(filter, device string) (*gopacket.PacketSource, error) {
	handle, err := pcap.OpenLive(
		device,
		snapshot_len,
		promiscuous,
		timeout,
	)
	if err != nil {
		return nil, err
	}
	// set the filter to monitor HTTP traffic for now
	err = handle.SetBPFFilter(filter)
	if err != nil {
		log.Fatal("Check the filter. Possibly there is a syntax error.")
	}
	packetSource := gopacket.NewPacketSource(
		handle,
		handle.LinkType(),
	)
	return packetSource, nil
}

// GetIPAddresses returns the IP addresses of source and
// Destination
func getIPAddresses(packet gopacket.Packet) (string, string) {
	if ipLayer := packet.Layer(layers.LayerTypeIPv4); ipLayer != nil {
		ip, _ := ipLayer.(*layers.IPv4)
		return ip.SrcIP.String(), ip.DstIP.String()
	} else if ipLayer := packet.Layer(layers.LayerTypeIPv6); ipLayer != nil {
		ip, _ := ipLayer.(*layers.IPv6)
		return ip.SrcIP.String(), ip.DstIP.String()
	}
	log.Errorf("[PacCap ] Couldn't inspect IP payload.")
	return "nil", "nil"
}

// GetPortAddresses returns the Port addresses of source and
// Destination
func getPortAddresses(packet gopacket.Packet) (string, string) {
	if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
		tcp, _ := tcpLayer.(*layers.TCP)
		return tcp.SrcPort.String(), tcp.DstPort.String()
	}
	log.Errorf("[PacCap ] Couldn't inspect TCP payload.")
	return "nil", "nil"
}
<|endoftext|>


<commit_before>package util

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/openvswitch/ovn-kubernetes/go-controller/pkg/config"
)

const (
	ovsCommandTimeout = 5
	ovsVsctlCommand   = "ovs-vsctl"
	ovsOfctlCommand   = "ovs-ofctl"
	ovnNbctlCommand   = "ovn-nbctl"
	osRelease         = "/etc/os-release"
	rhel              = "RHEL"
	ubuntu            = "Ubuntu"
	windowsOS         = "windows"
)

// PathExist checks the path exist or not.
func PathExist(path string) bool {
	_, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		return false
	}
	return true
}

func runningPlatform() (string, error) {
	if runtime.GOOS == windowsOS {
		return windowsOS, nil
	}
	fileContents, err := ioutil.ReadFile(osRelease)
	if err != nil {
		return "", fmt.Errorf("failed to parse file %s (%v)", osRelease, err)
	}

	var platform string
	ss := strings.Split(string(fileContents), "\n")
	for _, pair := range ss {
		keyValue := strings.Split(pair, "=")
		if len(keyValue) == 2 {
			if keyValue[0] == "Name" || keyValue[0] == "NAME" {
				platform = keyValue[1]
				break
			}
		}
	}

	if platform == "" {
		return "", fmt.Errorf("failed to find the platform name")
	}

	if strings.Contains(platform, "Fedora") ||
		strings.Contains(platform, "Red Hat") || strings.Contains(platform, "CentOS") {
		return rhel, nil
	} else if strings.Contains(platform, "Debian") ||
		strings.Contains(platform, ubuntu) {
		return ubuntu, nil
	} else if strings.Contains(platform, "VMware") {
		return "Photon", nil
	}
	return "", fmt.Errorf("Unknown platform")
}

// RunOVSOfctl runs a command via ovs-ofctl.
func RunOVSOfctl(args ...string) (string, string, error) {
	cmdPath, err := exec.LookPath(ovsOfctlCommand)
	if err != nil {
		return "", "", err
	}

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}
	cmd := exec.Command(cmdPath, args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	err = cmd.Run()
	return strings.Trim(stdout.String(), "\" \n"), stderr.String(), err
}

// RunOVSVsctl runs a command via ovs-vsctl.
func RunOVSVsctl(args ...string) (string, string, error) {
	cmdPath, err := exec.LookPath(ovsVsctlCommand)
	if err != nil {
		return "", "", err
	}

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}
	cmdArgs := []string{fmt.Sprintf("--timeout=%d", ovsCommandTimeout)}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command(cmdPath, cmdArgs...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	err = cmd.Run()
	return strings.Trim(strings.TrimSpace(stdout.String()), "\""), stderr.String(), err
}

// RunOVNNbctlWithTimeout runs command via ovs-nbctl with a specific timeout
func RunOVNNbctlWithTimeout(timeout int, args ...string) (string, string,
	error) {
	cmdPath, err := exec.LookPath(ovnNbctlCommand)
	if err != nil {
		return "", "", err
	}

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}

	var cmd *exec.Cmd
	if config.OvnNorth.ClientAuth.Scheme == config.OvnDBSchemeSSL {
		cmdArgs := []string{
			fmt.Sprintf("--private-key=%s", config.OvnNorth.ClientAuth.PrivKey),
			fmt.Sprintf("--certificate=%s", config.OvnNorth.ClientAuth.Cert),
			fmt.Sprintf("--bootstrap-ca-cert=%s", config.OvnNorth.ClientAuth.CACert),
			fmt.Sprintf("--db=%s", config.OvnNorth.ClientAuth.GetURL()),
			fmt.Sprintf("--timeout=%d", timeout),
		}
		cmdArgs = append(cmdArgs, args...)
		cmd = exec.Command(cmdPath, cmdArgs...)
	} else {
		cmdArgs := []string{
			fmt.Sprintf("--db=%s", config.OvnNorth.ClientAuth.GetURL()),
			fmt.Sprintf("--timeout=%d", timeout),
		}
		cmdArgs = append(cmdArgs, args...)
		cmd = exec.Command(cmdPath, cmdArgs...)
	}

	cmd.Stdout = stdout
	cmd.Stderr = stderr

	err = cmd.Run()
	return strings.Trim(strings.TrimSpace(stdout.String()), "\""), stderr.String(), err
}

// RunOVNNbctl runs a command via ovs-nbctl.
func RunOVNNbctl(args ...string) (string, string, error) {
	return RunOVNNbctlWithTimeout(ovsCommandTimeout, args...)
}
<commit_msg>util: New function for unix domain socket calls.<commit_after>package util

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/openvswitch/ovn-kubernetes/go-controller/pkg/config"
)

const (
	ovsCommandTimeout = 5
	ovsVsctlCommand   = "ovs-vsctl"
	ovsOfctlCommand   = "ovs-ofctl"
	ovnNbctlCommand   = "ovn-nbctl"
	osRelease         = "/etc/os-release"
	rhel              = "RHEL"
	ubuntu            = "Ubuntu"
	windowsOS         = "windows"
)

// PathExist checks the path exist or not.
func PathExist(path string) bool {
	_, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		return false
	}
	return true
}

func runningPlatform() (string, error) {
	if runtime.GOOS == windowsOS {
		return windowsOS, nil
	}
	fileContents, err := ioutil.ReadFile(osRelease)
	if err != nil {
		return "", fmt.Errorf("failed to parse file %s (%v)", osRelease, err)
	}

	var platform string
	ss := strings.Split(string(fileContents), "\n")
	for _, pair := range ss {
		keyValue := strings.Split(pair, "=")
		if len(keyValue) == 2 {
			if keyValue[0] == "Name" || keyValue[0] == "NAME" {
				platform = keyValue[1]
				break
			}
		}
	}

	if platform == "" {
		return "", fmt.Errorf("failed to find the platform name")
	}

	if strings.Contains(platform, "Fedora") ||
		strings.Contains(platform, "Red Hat") || strings.Contains(platform, "CentOS") {
		return rhel, nil
	} else if strings.Contains(platform, "Debian") ||
		strings.Contains(platform, ubuntu) {
		return ubuntu, nil
	} else if strings.Contains(platform, "VMware") {
		return "Photon", nil
	}
	return "", fmt.Errorf("Unknown platform")
}

// RunOVSOfctl runs a command via ovs-ofctl.
func RunOVSOfctl(args ...string) (string, string, error) {
	cmdPath, err := exec.LookPath(ovsOfctlCommand)
	if err != nil {
		return "", "", err
	}

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}
	cmd := exec.Command(cmdPath, args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	err = cmd.Run()
	return strings.Trim(stdout.String(), "\" \n"), stderr.String(), err
}

// RunOVSVsctl runs a command via ovs-vsctl.
func RunOVSVsctl(args ...string) (string, string, error) {
	cmdPath, err := exec.LookPath(ovsVsctlCommand)
	if err != nil {
		return "", "", err
	}

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}
	cmdArgs := []string{fmt.Sprintf("--timeout=%d", ovsCommandTimeout)}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command(cmdPath, cmdArgs...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	err = cmd.Run()
	return strings.Trim(strings.TrimSpace(stdout.String()), "\""), stderr.String(), err
}

// RunOVNNbctlUnix runs command via ovn-nbctl, with ovn-nbctl using the unix
// domain sockets to connect to the ovsdb-server backing the OVN NB database.
func RunOVNNbctlUnix(args ...string) (string, string, error) {
	cmdPath, err := exec.LookPath(ovnNbctlCommand)
	if err != nil {
		return "", "", err
	}

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}

	cmdArgs := []string{fmt.Sprintf("--timeout=%d", ovsCommandTimeout)}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command(cmdPath, cmdArgs...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	err = cmd.Run()
	return strings.Trim(strings.TrimSpace(stdout.String()), "\""), stderr.String(), err
}

// RunOVNNbctlWithTimeout runs command via ovn-nbctl with a specific timeout
func RunOVNNbctlWithTimeout(timeout int, args ...string) (string, string,
	error) {
	cmdPath, err := exec.LookPath(ovnNbctlCommand)
	if err != nil {
		return "", "", err
	}

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}

	var cmd *exec.Cmd
	if config.OvnNorth.ClientAuth.Scheme == config.OvnDBSchemeSSL {
		cmdArgs := []string{
			fmt.Sprintf("--private-key=%s", config.OvnNorth.ClientAuth.PrivKey),
			fmt.Sprintf("--certificate=%s", config.OvnNorth.ClientAuth.Cert),
			fmt.Sprintf("--bootstrap-ca-cert=%s", config.OvnNorth.ClientAuth.CACert),
			fmt.Sprintf("--db=%s", config.OvnNorth.ClientAuth.GetURL()),
			fmt.Sprintf("--timeout=%d", timeout),
		}
		cmdArgs = append(cmdArgs, args...)
		cmd = exec.Command(cmdPath, cmdArgs...)
	} else {
		cmdArgs := []string{
			fmt.Sprintf("--db=%s", config.OvnNorth.ClientAuth.GetURL()),
			fmt.Sprintf("--timeout=%d", timeout),
		}
		cmdArgs = append(cmdArgs, args...)
		cmd = exec.Command(cmdPath, cmdArgs...)
	}

	cmd.Stdout = stdout
	cmd.Stderr = stderr

	err = cmd.Run()
	return strings.Trim(strings.TrimSpace(stdout.String()), "\""), stderr.String(), err
}

// RunOVNNbctl runs a command via ovn-nbctl.
func RunOVNNbctl(args ...string) (string, string, error) {
	return RunOVNNbctlWithTimeout(ovsCommandTimeout, args...)
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"github.com/jonboulle/clockwork"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	"strings"
	"testing"
	"time"
)

type flakeyRooterAPI struct {
	orig     libkb.ExternalAPI
	flakeOut bool
}

func newFlakeyRooterAPI(x libkb.ExternalAPI) *flakeyRooterAPI {
	return &flakeyRooterAPI{
		orig: x,
	}
}

func (e *flakeyRooterAPI) GetText(arg libkb.APIArg) (*libkb.ExternalTextRes, error) {
	return e.orig.GetText(arg)
}

func (e *flakeyRooterAPI) Get(arg libkb.APIArg) (res *libkb.ExternalAPIRes, err error) {

	// Show an error if we're in flakey mode
	if e.flakeOut && strings.Contains(arg.Endpoint, "rooter") {
		return &libkb.ExternalAPIRes{HTTPStatus: 429}, &libkb.APIError{Status: "Ratelimited", StatusCode: 429}
	}

	return e.orig.Get(arg)
}

func (e *flakeyRooterAPI) GetHTML(arg libkb.APIArg) (res *libkb.ExternalHTMLRes, err error) {
	return e.orig.GetHTML(arg)
}

func (e *flakeyRooterAPI) Post(arg libkb.APIArg) (res *libkb.ExternalAPIRes, err error) {
	return e.orig.Post(arg)
}

func (e *flakeyRooterAPI) PostHTML(arg libkb.APIArg) (res *libkb.ExternalHTMLRes, err error) {
	return e.orig.PostHTML(arg)
}

func TestSoftSnooze(t *testing.T) {
	tc := SetupEngineTest(t, "track")
	defer tc.Cleanup()
	fu := CreateAndSignupFakeUser(tc, "track")

	fakeClock := clockwork.NewFakeClock()
	tc.G.Clock = fakeClock
	// to pick up the new clock...
	tc.G.ResetLoginState()

	flakeyAPI := flakeyRooterAPI{orig: tc.G.XAPI, flakeOut: false}
	tc.G.XAPI = &flakeyAPI

	idUI := &FakeIdentifyUI{}
	username := "t_tracy"
	arg := &keybase1.Identify2Arg{
		UserAssertion: username,
		NeedProofSet:  true,
	}
	ctx := &Context{
		LogUI:      tc.G.UI.GetLogUI(),
		IdentifyUI: idUI,
		SecretUI:   fu.NewSecretUI(),
	}

	// Identify tracy; all proofs should work
	eng := NewResolveThenIdentify2(tc.G, arg)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}
	targ := TrackTokenArg{
		Token:   idUI.Token,
		Options: keybase1.TrackOptions{BypassConfirm: true},
	}

	// Track tracy
	teng := NewTrackToken(&targ, tc.G)
	if err := RunEngine(teng, ctx); err != nil {
		t.Fatal(err)
	}

	// Now make her Rooter proof flakey / fail with a 429
	flakeyAPI.flakeOut = true
	idUI = &FakeIdentifyUI{}
	ctx.IdentifyUI = idUI

	// Advance so that our previous cached success is out of
	// cache on its own, but still can override a 429-like soft failure.
	fakeClock.Advance(tc.G.Env.GetProofCacheMediumDur() + time.Minute)

	eng = NewResolveThenIdentify2(tc.G, arg)
	eng.testArgs = &Identify2WithUIDTestArgs{noCache: true}
	// Should not get an error
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}
	result, found := idUI.ProofResults["rooter"]
	if !found {
		t.Fatal("Failed to find a rooter proof")
	}
	if pe := libkb.ImportProofError(result.SnoozedResult); pe == nil {
		t.Fatal("expected a snoozed error result")
	}

	// Now time out the success that allowed us to circumvent
	// the soft failure.
	fakeClock.Advance(tc.G.Env.GetProofCacheLongDur())
	eng = NewResolveThenIdentify2(tc.G, arg)
	eng.testArgs = &Identify2WithUIDTestArgs{noCache: true}
	if err := RunEngine(eng, ctx); err == nil {
		t.Fatal("Expected a failure in our proof")
	}

	result, found = idUI.ProofResults["rooter"]
	if !found {
		t.Fatal("Failed to find a rooter proof")
	}
	if pe := libkb.ImportProofError(result.ProofResult); pe == nil {
		t.Fatal("expected a Rooter error result")
	}

	defer runUntrack(tc.G, fu, username)
	assertTracking(tc, username)
}
<commit_msg>whoops, fix test<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"github.com/jonboulle/clockwork"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	"strings"
	"testing"
	"time"
)

type flakeyRooterAPI struct {
	orig     libkb.ExternalAPI
	flakeOut bool
}

func newFlakeyRooterAPI(x libkb.ExternalAPI) *flakeyRooterAPI {
	return &flakeyRooterAPI{
		orig: x,
	}
}

func (e *flakeyRooterAPI) GetText(arg libkb.APIArg) (*libkb.ExternalTextRes, error) {
	return e.orig.GetText(arg)
}

func (e *flakeyRooterAPI) Get(arg libkb.APIArg) (res *libkb.ExternalAPIRes, err error) {

	// Show an error if we're in flakey mode
	if e.flakeOut && strings.Contains(arg.Endpoint, "rooter") {
		return &libkb.ExternalAPIRes{HTTPStatus: 429}, &libkb.APIError{Msg: "Ratelimited", Code: 429}
	}

	return e.orig.Get(arg)
}

func (e *flakeyRooterAPI) GetHTML(arg libkb.APIArg) (res *libkb.ExternalHTMLRes, err error) {
	return e.orig.GetHTML(arg)
}

func (e *flakeyRooterAPI) Post(arg libkb.APIArg) (res *libkb.ExternalAPIRes, err error) {
	return e.orig.Post(arg)
}

func (e *flakeyRooterAPI) PostHTML(arg libkb.APIArg) (res *libkb.ExternalHTMLRes, err error) {
	return e.orig.PostHTML(arg)
}

func TestSoftSnooze(t *testing.T) {
	tc := SetupEngineTest(t, "track")
	defer tc.Cleanup()
	fu := CreateAndSignupFakeUser(tc, "track")

	fakeClock := clockwork.NewFakeClock()
	tc.G.Clock = fakeClock
	// to pick up the new clock...
	tc.G.ResetLoginState()

	flakeyAPI := flakeyRooterAPI{orig: tc.G.XAPI, flakeOut: false}
	tc.G.XAPI = &flakeyAPI

	idUI := &FakeIdentifyUI{}
	username := "t_tracy"
	arg := &keybase1.Identify2Arg{
		UserAssertion: username,
		NeedProofSet:  true,
	}
	ctx := &Context{
		LogUI:      tc.G.UI.GetLogUI(),
		IdentifyUI: idUI,
		SecretUI:   fu.NewSecretUI(),
	}

	// Identify tracy; all proofs should work
	eng := NewResolveThenIdentify2(tc.G, arg)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}
	targ := TrackTokenArg{
		Token:   idUI.Token,
		Options: keybase1.TrackOptions{BypassConfirm: true},
	}

	// Track tracy
	teng := NewTrackToken(&targ, tc.G)
	if err := RunEngine(teng, ctx); err != nil {
		t.Fatal(err)
	}

	// Now make her Rooter proof flakey / fail with a 429
	flakeyAPI.flakeOut = true
	idUI = &FakeIdentifyUI{}
	ctx.IdentifyUI = idUI

	// Advance so that our previous cached success is out of
	// cache on its own, but still can override a 429-like soft failure.
	fakeClock.Advance(tc.G.Env.GetProofCacheMediumDur() + time.Minute)

	eng = NewResolveThenIdentify2(tc.G, arg)
	eng.testArgs = &Identify2WithUIDTestArgs{noCache: true}
	// Should not get an error
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}
	result, found := idUI.ProofResults["rooter"]
	if !found {
		t.Fatal("Failed to find a rooter proof")
	}
	if pe := libkb.ImportProofError(result.SnoozedResult); pe == nil {
		t.Fatal("expected a snoozed error result")
	}

	// Now time out the success that allowed us to circumvent
	// the soft failure.
	fakeClock.Advance(tc.G.Env.GetProofCacheLongDur())
	eng = NewResolveThenIdentify2(tc.G, arg)
	eng.testArgs = &Identify2WithUIDTestArgs{noCache: true}
	if err := RunEngine(eng, ctx); err == nil {
		t.Fatal("Expected a failure in our proof")
	}

	result, found = idUI.ProofResults["rooter"]
	if !found {
		t.Fatal("Failed to find a rooter proof")
	}
	if pe := libkb.ImportProofError(result.ProofResult); pe == nil {
		t.Fatal("expected a Rooter error result")
	}

	defer runUntrack(tc.G, fu, username)
	assertTracking(tc, username)
}
<|endoftext|>


<commit_before>package para

import "sync"

type Fetcher struct {
	wg             sync.WaitGroup
	cleanupWaiting sync.Once
	initClose      sync.Once
	resultChan     chan *Result
	closeChan      chan struct{}
	// Public channels
	Get chan *Result
	// optional channel to notify outside when clean up is complete
	Done chan struct{}
}

type Result struct {
	err      error
	response string
}

func NewFetcher() *Fetcher {
	f := &Fetcher{
		wg:             sync.WaitGroup{},
		cleanupWaiting: sync.Once{},
		initClose:      sync.Once{},
		closeChan:      make(chan struct{}),
		resultChan:     make(chan *Result, 10),

		Get:  make(chan *Result),
		Done: make(chan struct{}),
	}
	go f.Listen()
	return f
}

// Fetch is called by the thread initiating the fanout.
// It "forks" and starts up the described worker.
//
// The Worker should implement its own timeout functionality
//
// TODO what should be the Format for the worker func?
func (f *Fetcher) Fetch(worker func() (string, error)) {
	f.wg.Add(1)
	f.cleanupWaiting.Do(func() { go f.waitAndCleanup(f.Done) })
	go func(worker func() (string, error)) {
		r, e := worker()
		select {
		case <-f.closeChan:
			f.wg.Done()
		default:
			f.resultChan <- &Result{
				err:      e,
				response: r,
			}
		}
	}(worker)
}

func (f *Fetcher) waitAndCleanup(doneChan chan struct{}) {
	// Wait until all workers have returned or a successful result has come
	// and workers can shutdown
	f.wg.Wait()
	// then tell the listen loop to close
	close(f.resultChan)
	// and tell the outside we are cleaned up.
	close(doneChan)
}

// Listen validates the results and initiates shutdown if there is a valid result.
// If no valid result is found it emits the last failing result.
func (f *Fetcher) Listen() {
	var lastResult *Result
	for r := range f.resultChan {
		f.wg.Done()
		lastResult = r
		if r.err == nil {
			f.initClose.Do(func() {
				close(f.closeChan)
				f.Get <- r
			})
		}
	}
	f.Get <- lastResult
}
<commit_msg>fixed race and smoothed out result channel buffering problem<commit_after>package para

import "sync"

type Fetcher struct {
	wg             sync.WaitGroup
	cleanupWaiting sync.Once
	initClose      sync.Once
	resultChan     chan *Result
	closeChan      chan struct{}
	// Public channels
	Get chan *Result
	// optional channel to notify outside when clean up is complete
	Done chan struct{}
}

type Result struct {
	err      error
	response string
}

func NewFetcher() *Fetcher {
	f := &Fetcher{
		wg:             sync.WaitGroup{},
		cleanupWaiting: sync.Once{},
		initClose:      sync.Once{},
		closeChan:      make(chan struct{}),
		resultChan:     make(chan *Result, 1),

		Get:  make(chan *Result),
		Done: make(chan struct{}),
	}
	go f.Listen()
	return f
}

// Fetch is called by the thread initiating the fanout.
// It "forks" and starts up the described worker.
//
// The Worker should implement its own timeout functionality
//
// TODO what should be the Format for the worker func?
func (f *Fetcher) Fetch(worker func() (string, error)) {
	f.wg.Add(1)
	f.cleanupWaiting.Do(func() { go f.waitAndCleanup(f.Done) })
	go func(worker func() (string, error)) {
		r, e := worker()
		select {
		case <-f.closeChan:
			f.wg.Done()
		case f.resultChan <- &Result{
			err:      e,
			response: r,
		}:
		}
	}(worker)
}

func (f *Fetcher) waitAndCleanup(doneChan chan struct{}) {
	// Wait until all workers have returned or a successful result has come
	// and workers can shutdown
	f.wg.Wait()
	// then tell the listen loop to close
	close(f.resultChan)
	// and tell the outside we are cleaned up.
	close(doneChan)
}

// Listen validates the results and initiates shutdown if there is a valid result.
// If no valid result is found it emits the last failing result.
func (f *Fetcher) Listen() {
	var lastResult *Result
	for r := range f.resultChan {
		f.wg.Done()
		lastResult = r
		if r.err == nil {
			f.initClose.Do(func() {
				close(f.closeChan)
				f.Get <- r
			})
		}
	}
	f.Get <- lastResult
}
<|endoftext|>