<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	env "github.com/MattAitchison/envconfig"
	"github.com/fsouza/go-dockerclient"
	"github.com/gliderlabs/connectable/pkg/lookup"

	_ "github.com/gliderlabs/connectable/pkg/lookup/dns"
)

var Version string

var (
	endpoint = env.String("docker_host", "unix:///var/run/docker.sock", "docker endpoint")
	port     = env.String("port", "10000", "primary listen port")

	self *docker.Container
)

func assert(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func runNetCmd(container, image string, cmd string) error {
	client, err := docker.NewClient(endpoint)
	if err != nil {
		return err
	}
	c, err := client.CreateContainer(docker.CreateContainerOptions{
		Config: &docker.Config{
			Image:      image,
			Cmd:        []string{cmd},
			Entrypoint: []string{"/bin/sh", "-c"},
		},
		HostConfig: &docker.HostConfig{
			Privileged:  true,
			NetworkMode: fmt.Sprintf("container:%s", container),
		},
	})
	if err != nil {
		return err
	}
	if err := client.StartContainer(c.ID, nil); err != nil {
		return err
	}
	status, err := client.WaitContainer(c.ID)
	if err != nil {
		return err
	}
	if status != 0 {
		return fmt.Errorf("netcmd non-zero exit: %v", status)
	}
	return client.RemoveContainer(docker.RemoveContainerOptions{
		ID:    c.ID,
		Force: true,
	})
}

func originalDestinationPort(conn net.Conn) (string, error) {
	f, err := conn.(*net.TCPConn).File()
	if err != nil {
		return "", err
	}
	defer f.Close()
	addr, err := syscall.GetsockoptIPv6Mreq(
		int(f.Fd()), syscall.IPPROTO_IP, 80) // 80 = SO_ORIGINAL_DST
	if err != nil {
		return "", err
	}
	port := uint16(addr.Multiaddr[2])<<8 + uint16(addr.Multiaddr[3])
	return strconv.Itoa(int(port)), nil
}

func inspectBackend(sourceIP, destPort string) (string, error) {
	client, err := docker.NewClient(endpoint)
	if err != nil {
		return "", err
	}
	label := fmt.Sprintf("connect[%s]", destPort)

	// todo: cache, invalidate with container destroy events
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return "", err
	}
	for _, listing := range containers {
		container, err := client.InspectContainer(listing.ID)
		if err != nil {
			return "", err
		}
		if container.NetworkSettings.IPAddress == sourceIP {
			backend, ok := container.Config.Labels[label]
			if !ok {
				return "", fmt.Errorf("connect label '%s' not found: %v", label, container.Config.Labels)
			}
			return backend, nil
		}
	}
	return "", fmt.Errorf("unable to find container with source IP")
}

func lookupBackend(conn net.Conn) string {
	sourceIP, _, _ := net.SplitHostPort(conn.RemoteAddr().String())
	destPort, err := originalDestinationPort(conn)
	if err != nil {
		log.Println("unable to determine destination port")
		return ""
	}

	backend, err := inspectBackend(sourceIP, destPort)
	if err != nil {
		log.Println(err)
		return ""
	}
	return backend
}

func proxyConn(conn net.Conn, addr string) {
	backend, err := net.Dial("tcp", addr)
	defer conn.Close()
	if err != nil {
		log.Println("proxy", err.Error())
		return
	}
	defer backend.Close()

	done := make(chan struct{})
	go func() {
		io.Copy(backend, conn)
		backend.(*net.TCPConn).CloseWrite()
		close(done)
	}()
	io.Copy(conn, backend)
	conn.(*net.TCPConn).CloseWrite()
	<-done
}

func setupContainer(id string) {
	re := regexp.MustCompile("connect\\[(\\d+)\\]")
	client, err := docker.NewClient(endpoint)
	if err != nil {
		log.Println(err)
	}
	container, err := client.InspectContainer(id)
	if err != nil {
		log.Println(err)
	}
	if container.HostConfig.NetworkMode == "bridge" {
		hasBackends := false
		cmds := []string{
			"/sbin/sysctl -w net.ipv4.conf.all.route_localnet=1",
			"iptables -t nat -I POSTROUTING 1 -m addrtype --src-type LOCAL --dst-type UNICAST -j MASQUERADE",
		}
		for k, _ := range container.Config.Labels {
			results := re.FindStringSubmatch(k)
			if len(results) > 1 {
				hasBackends = true
				cmds = append(cmds, fmt.Sprintf(
					"iptables -t nat -I OUTPUT 1 -m addrtype --src-type LOCAL --dst-type LOCAL -p tcp --dport %s -j DNAT --to-destination %s:%s",
					results[1], self.NetworkSettings.IPAddress, results[1]))
			}
		}
		if hasBackends {
			log.Printf("setting iptables on %s \n", container.ID[:12])
			shellCmd := strings.Join(cmds, " && ")
			assert(runNetCmd(container.ID, self.Image, shellCmd))
		}
	}
}

func monitorContainers() {
	client, err := docker.NewClient(endpoint)
	assert(err)
	events := make(chan *docker.APIEvents)
	assert(client.AddEventListener(events))
	list, _ := client.ListContainers(docker.ListContainersOptions{})
	for _, listing := range list {
		go setupContainer(listing.ID)
	}
	for msg := range events {
		switch msg.Status {
		case "create":
			go setupContainer(msg.ID)
		}
	}
}

func main() {
	listener, err := net.Listen("tcp", ":"+port)
	assert(err)

	fmt.Printf("# Connectable %s listening on %s ...\n", Version, port)

	client, err := docker.NewClient(endpoint)
	assert(err)

	list, err := client.ListContainers(docker.ListContainersOptions{})
	assert(err)
	for _, listing := range list {
		c, err := client.InspectContainer(listing.ID)
		assert(err)
		if c.Config.Hostname == os.Getenv("HOSTNAME") {
			self = c
			if c.HostConfig.NetworkMode == "bridge" {
				fmt.Printf("# Setting iptables on connectable... ")
				shellCmd := fmt.Sprintf("iptables -t nat -A PREROUTING -p tcp -j REDIRECT --to-ports %s", port)
				assert(runNetCmd(c.ID, c.Image, shellCmd))
				fmt.Printf("done.\n")
			}
		}
	}

	if self == nil {
		fmt.Println("# unable to find self")
		os.Exit(1)
	}

	go monitorContainers()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Fatal(err)
		}

		backend := lookupBackend(conn)
		if backend == "" {
			conn.Close()
			continue
		}

		backendAddrs, err := lookup.Resolve(backend)
		if err != nil {
			log.Println(err)
			conn.Close()
			continue
		}
		if len(backendAddrs) == 0 {
			log.Println(conn.RemoteAddr(), backend, "no backends")
			conn.Close()
			continue
		}

		log.Println(conn.RemoteAddr(), backend, "->", backendAddrs[0])
		go proxyConn(conn, backendAddrs[0])
	}
}
<commit_msg>Just log an error when failing to setup iptables for a specific container<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	env "github.com/MattAitchison/envconfig"
	"github.com/fsouza/go-dockerclient"
	"github.com/gliderlabs/connectable/pkg/lookup"

	_ "github.com/gliderlabs/connectable/pkg/lookup/dns"
)

var Version string

var (
	endpoint = env.String("docker_host", "unix:///var/run/docker.sock", "docker endpoint")
	port     = env.String("port", "10000", "primary listen port")

	self *docker.Container
)

func assert(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func runNetCmd(container, image string, cmd string) error {
	client, err := docker.NewClient(endpoint)
	if err != nil {
		return err
	}
	c, err := client.CreateContainer(docker.CreateContainerOptions{
		Config: &docker.Config{
			Image:      image,
			Cmd:        []string{cmd},
			Entrypoint: []string{"/bin/sh", "-c"},
		},
		HostConfig: &docker.HostConfig{
			Privileged:  true,
			NetworkMode: fmt.Sprintf("container:%s", container),
		},
	})
	if err != nil {
		return err
	}
	if err := client.StartContainer(c.ID, nil); err != nil {
		return err
	}
	status, err := client.WaitContainer(c.ID)
	if err != nil {
		return err
	}
	if status != 0 {
		return fmt.Errorf("netcmd non-zero exit: %v", status)
	}
	return client.RemoveContainer(docker.RemoveContainerOptions{
		ID:    c.ID,
		Force: true,
	})
}

func originalDestinationPort(conn net.Conn) (string, error) {
	f, err := conn.(*net.TCPConn).File()
	if err != nil {
		return "", err
	}
	defer f.Close()
	addr, err := syscall.GetsockoptIPv6Mreq(
		int(f.Fd()), syscall.IPPROTO_IP, 80) // 80 = SO_ORIGINAL_DST
	if err != nil {
		return "", err
	}
	port := uint16(addr.Multiaddr[2])<<8 + uint16(addr.Multiaddr[3])
	return strconv.Itoa(int(port)), nil
}

func inspectBackend(sourceIP, destPort string) (string, error) {
	client, err := docker.NewClient(endpoint)
	if err != nil {
		return "", err
	}
	label := fmt.Sprintf("connect[%s]", destPort)

	// todo: cache, invalidate with container destroy events
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return "", err
	}
	for _, listing := range containers {
		container, err := client.InspectContainer(listing.ID)
		if err != nil {
			return "", err
		}
		if container.NetworkSettings.IPAddress == sourceIP {
			backend, ok := container.Config.Labels[label]
			if !ok {
				return "", fmt.Errorf("connect label '%s' not found: %v", label, container.Config.Labels)
			}
			return backend, nil
		}
	}
	return "", fmt.Errorf("unable to find container with source IP")
}

func lookupBackend(conn net.Conn) string {
	sourceIP, _, _ := net.SplitHostPort(conn.RemoteAddr().String())
	destPort, err := originalDestinationPort(conn)
	if err != nil {
		log.Println("unable to determine destination port")
		return ""
	}

	backend, err := inspectBackend(sourceIP, destPort)
	if err != nil {
		log.Println(err)
		return ""
	}
	return backend
}

func proxyConn(conn net.Conn, addr string) {
	backend, err := net.Dial("tcp", addr)
	defer conn.Close()
	if err != nil {
		log.Println("proxy", err.Error())
		return
	}
	defer backend.Close()

	done := make(chan struct{})
	go func() {
		io.Copy(backend, conn)
		backend.(*net.TCPConn).CloseWrite()
		close(done)
	}()
	io.Copy(conn, backend)
	conn.(*net.TCPConn).CloseWrite()
	<-done
}

func setupContainer(id string) error {
	re := regexp.MustCompile("connect\\[(\\d+)\\]")
	client, err := docker.NewClient(endpoint)
	if err != nil {
		log.Println(err)
		return err
	}
	container, err := client.InspectContainer(id)
	if err != nil {
		log.Println(err)
		return err
	}
	if container.HostConfig.NetworkMode == "bridge" {
		hasBackends := false
		cmds := []string{
			"/sbin/sysctl -w net.ipv4.conf.all.route_localnet=1",
			"iptables -t nat -I POSTROUTING 1 -m addrtype --src-type LOCAL --dst-type UNICAST -j MASQUERADE",
		}
		for k, _ := range container.Config.Labels {
			results := re.FindStringSubmatch(k)
			if len(results) > 1 {
				hasBackends = true
				cmds = append(cmds, fmt.Sprintf(
					"iptables -t nat -I OUTPUT 1 -m addrtype --src-type LOCAL --dst-type LOCAL -p tcp --dport %s -j DNAT --to-destination %s:%s",
					results[1], self.NetworkSettings.IPAddress, results[1]))
			}
		}
		if hasBackends {
			log.Printf("setting iptables on %s \n", container.ID[:12])
			shellCmd := strings.Join(cmds, " && ")
			err := runNetCmd(container.ID, self.Image, shellCmd)
			log.Printf("error setting iptables on %s: %s \n", container.ID[:12], err)
			return err
		}
	}
	return nil

}

func monitorContainers() {
	client, err := docker.NewClient(endpoint)
	assert(err)
	events := make(chan *docker.APIEvents)
	assert(client.AddEventListener(events))
	list, _ := client.ListContainers(docker.ListContainersOptions{})
	for _, listing := range list {
		go setupContainer(listing.ID)
	}
	for msg := range events {
		switch msg.Status {
		case "create":
			go setupContainer(msg.ID)
		}
	}
}

func main() {
	listener, err := net.Listen("tcp", ":"+port)
	assert(err)

	fmt.Printf("# Connectable %s listening on %s ...\n", Version, port)

	client, err := docker.NewClient(endpoint)
	assert(err)

	list, err := client.ListContainers(docker.ListContainersOptions{})
	assert(err)
	for _, listing := range list {
		c, err := client.InspectContainer(listing.ID)
		assert(err)
		if c.Config.Hostname == os.Getenv("HOSTNAME") {
			self = c
			if c.HostConfig.NetworkMode == "bridge" {
				fmt.Printf("# Setting iptables on connectable... ")
				shellCmd := fmt.Sprintf("iptables -t nat -A PREROUTING -p tcp -j REDIRECT --to-ports %s", port)
				assert(runNetCmd(c.ID, c.Image, shellCmd))
				fmt.Printf("done.\n")
			}
		}
	}

	if self == nil {
		fmt.Println("# unable to find self")
		os.Exit(1)
	}

	go monitorContainers()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Fatal(err)
		}

		backend := lookupBackend(conn)
		if backend == "" {
			conn.Close()
			continue
		}

		backendAddrs, err := lookup.Resolve(backend)
		if err != nil {
			log.Println(err)
			conn.Close()
			continue
		}
		if len(backendAddrs) == 0 {
			log.Println(conn.RemoteAddr(), backend, "no backends")
			conn.Close()
			continue
		}

		log.Println(conn.RemoteAddr(), backend, "->", backendAddrs[0])
		go proxyConn(conn, backendAddrs[0])
	}
}
<|endoftext|>


<commit_before>package libcomfo

import (
	"bytes"
	"fmt"
	"io"
	"reflect"
	"testing"
	"time"
)

// TestConn is a structure that implements io.ReadWriter and can be used
// to mock a connection with the unit. HangStart makes initial Read and Write
// calls hang. HangEOF makes Read hang on the second call when a Scanner
// would expect an EOF error to be returned.
// Not thread-safe.
type TestConn struct {
	Receives []byte
	Emits    []byte

	HangTime  time.Duration
	HangStart bool
	HangEOF   bool

	Limit int

	emitted  int
	received int
}

func (tr *TestConn) Read(p []byte) (n int, err error) {

	// Simulate read delay before returning output or EOF
	if tr.HangTime > 0 && tr.HangStart {
		time.Sleep(tr.HangTime)
	}

	// Return EOF when buffer is fully copied
	if tr.emitted == len(tr.Emits) {

		if tr.HangTime > 0 && tr.HangEOF {
			time.Sleep(tr.HangTime)
		}

		return 0, io.EOF
	}

	// Return the data in the connection's Emits field
	n = copy(p, tr.Emits)

	// Keep track of how many bytes we've been able to copy
	tr.emitted = tr.emitted + n

	// Return mocked read bytes total
	if tr.Limit > 0 {
		n = tr.Limit
	}

	return
}

func (tr *TestConn) Write(p []byte) (n int, err error) {

	// Simulate write delay
	if tr.HangTime > 0 {
		time.Sleep(tr.HangTime)
	}

	// Bounds check for Receives
	if tr.received+len(p) > len(tr.Receives) {
		return n,
			fmt.Errorf("write to TestConn exceeded expected input specified in Receives: %v",
				// This formula determines the offset of the slice
				// of surplus characters over the struct's 'Receives' field.
				p[len(p)-(tr.received+len(p)-len(tr.Receives)):])
	}

	// Expect the written data to correspond to what's in the 'Receives' field.
	// Offset the comparison against the write index (received).
	if want, got := tr.Receives[tr.received:tr.received+len(p)], p; !bytes.Equal(want, got) {
		err = fmt.Errorf(
			"unexpected write to TestConn:\n  - want: %v\n  -  got: %v", want, got)
	}

	// Keep a record of the bytes we successfully received
	// and compared against the 'Receives' buffer.
	tr.received = tr.received + len(p)

	// Return mocked written bytes total
	if tr.Limit > 0 {
		n = tr.Limit
	} else {
		// Return the length of the 'written' slice
		n = len(p)
	}

	return
}

func TestWaitTimeout(t *testing.T) {

	returnChan := make(chan bool)

	// Return timer at 2 milliseconds
	returnTimer := func() {
		time.Sleep(time.Millisecond * 2)
		returnChan <- true
	}

	// Run test 10 times back-to-back
	for i := 0; i < 10; i++ {
		t.Run(t.Name(), func(t *testing.T) {

			// Start a timeout timer with a timeout higher than the return timer
			timeOutTimer := time.NewTimer(time.Millisecond * 5)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err := WaitTimeout(returnChan, timeOutTimer)
			if err != nil {
				t.Fatal("returnChan did not unblock before timeOutTimer")
			}

			// Start a timeout timer with a lower timeout than return timer
			timeOutTimer = time.NewTimer(time.Millisecond * 1)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err = WaitTimeout(returnChan, timeOutTimer)
			if err != errTimeout {
				t.Fatal("timeOutTimer did not expire before returnChan unblock")
			}

			// Wait for the return timer to send on channel
			// This value needs to be read to prevent it from interfering other tests
			<-returnChan
		})
	}
}

func TestReadPacket(t *testing.T) {

	rt := []struct {
		name string
		tc   TestConn
		pkt  Packet
		err  error
	}{
		{
			name: "single and double escaped 0x07s in payload",
			tc: TestConn{
				Emits: []byte{ // Response
					esc, ack,
					0x07, 0xF0, // Frame Start
					0x00, uint8(getFans + 1), // response type
					0x06,       // Length
					0xAA, 0xBB, // in/out percents
					0x11, 0x07, 0x07, // in speed (one seven)
					0x07, 0x07, 0x07, 0x07, // out speed (two sevens)
					0x4A, // Checksum
					esc, end,
					esc, ack,
				},
			},
			pkt: Packet{
				Command: uint8(getFans + 1),
				Data:    []byte{0xAA, 0xBB, 0x11, 0x7, 0x7, 0x7},
			},
		},
		{
			name: "timeout reading first response ACK",
			tc: TestConn{
				HangStart: true,
				HangTime:  2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading start sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading end sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack, esc, start},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "garbage before ACK",
			tc: TestConn{
				Emits: []byte{esc, ack, 0xF0, 0x0B, esc, start},
			},
			err: errScanInput,
		},
		{
			name: "impossible packet length (need at least 4 bytes)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, esc, end},
			},
			err: errTooShort,
		},
		{
			name: "packet unmarshal error (any)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, 0xBA, 0xAA, esc, end},
			},
			err: errPayloadSize,
		},
	}

	for _, tt := range rt {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Start timeout timer
			to := time.NewTimer(time.Millisecond * 1)

			// Attempt to read packet from mock connection
			pkt, err := ReadPacket(&tc, to, true)

			if want, got := tt.err, err; want != got {
				t.Fatalf("unexpected error reading packet:\n- want: %v\n-  got: %v",
					want, got)
			}

			if want, got := tt.pkt, pkt; !reflect.DeepEqual(want, got) {
				t.Fatalf("unexpected packet read from connection:\n- want: %v\n-  got: %v",
					want, got)
			}
		})
	}
}

func TestWritePacket(t *testing.T) {

	wt := []struct {
		name string
		tc   TestConn
		pkt  Packet
		err  error
	}{
		{
			name: "oversized packet (handle MarshalPacket error)",
			pkt: Packet{
				Data: make([]byte, 256),
			},
			err: errTooLong,
		},
	}

	for _, tt := range wt {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Attempt to write packet to mock connection
			_, err := WritePacket(tt.pkt, &tc)

			if want, got := tt.err, err; want != got {
				t.Fatalf("unexpected error writing packet:\n- want: %v\n-  got: %v",
					want, got)
			}
		})
	}
}

func TestWriteAck(t *testing.T) {

	at := []struct {
		name string
		tc   TestConn
		err  bool
	}{
		{
			name: "writer not expecting input",
			tc:   TestConn{},
			err:  true,
		},
		{
			name: "incomplete ACK write",
			tc: TestConn{
				Receives: []byte{esc, ack},
				Limit:    1,
			},
			err: true,
		},
	}

	for _, tt := range at {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Attempt to write ACK to mock connection
			_, err := WriteAck(&tc)

			if tt.err != (err != nil) {
				t.Fatalf("unexpected error condition in ACK write:\n- want error: %v\n-  got error: %v",
					tt.err, err)
			}
		})
	}
}
<commit_msg>libcomfo - TestQueryPacket: harness for end-to-end wire testing<commit_after>package libcomfo

import (
	"bytes"
	"fmt"
	"io"
	"reflect"
	"testing"
	"time"
)

// TestConn is a structure that implements io.ReadWriter and can be used
// to mock a connection with the unit. HangStart makes initial Read and Write
// calls hang. HangEOF makes Read hang on the second call when a Scanner
// would expect an EOF error to be returned.
// Not thread-safe.
type TestConn struct {
	Receives []byte
	Emits    []byte

	HangTime  time.Duration
	HangStart bool
	HangEOF   bool

	Limit int

	emitted  int
	received int
}

func (tr *TestConn) Read(p []byte) (n int, err error) {

	// Simulate read delay before returning output or EOF
	if tr.HangTime > 0 && tr.HangStart {
		time.Sleep(tr.HangTime)
	}

	// Return EOF when buffer is fully copied
	if tr.emitted == len(tr.Emits) {

		if tr.HangTime > 0 && tr.HangEOF {
			time.Sleep(tr.HangTime)
		}

		return 0, io.EOF
	}

	// Return the data in the connection's Emits field
	n = copy(p, tr.Emits)

	// Keep track of how many bytes we've been able to copy
	tr.emitted = tr.emitted + n

	// Return mocked read bytes total
	if tr.Limit > 0 {
		n = tr.Limit
	}

	return
}

func (tr *TestConn) Write(p []byte) (n int, err error) {

	// Simulate write delay
	if tr.HangTime > 0 {
		time.Sleep(tr.HangTime)
	}

	// Bounds check for Receives
	if tr.received+len(p) > len(tr.Receives) {
		return n,
			fmt.Errorf("write to TestConn exceeded expected input specified in Receives: %v",
				// This formula determines the offset of the slice
				// of surplus characters over the struct's 'Receives' field.
				p[len(p)-(tr.received+len(p)-len(tr.Receives)):])
	}

	// Expect the written data to correspond to what's in the 'Receives' field.
	// Offset the comparison against the write index (received).
	if want, got := tr.Receives[tr.received:tr.received+len(p)], p; !bytes.Equal(want, got) {
		err = fmt.Errorf(
			"unexpected write to TestConn:\n  - want: %v\n  -  got: %v", want, got)
	}

	// Keep a record of the bytes we successfully received
	// and compared against the 'Receives' buffer.
	tr.received = tr.received + len(p)

	// Return mocked written bytes total
	if tr.Limit > 0 {
		n = tr.Limit
	} else {
		// Return the length of the 'written' slice
		n = len(p)
	}

	return
}

func TestWaitTimeout(t *testing.T) {

	returnChan := make(chan bool)

	// Return timer at 2 milliseconds
	returnTimer := func() {
		time.Sleep(time.Millisecond * 2)
		returnChan <- true
	}

	// Run test 10 times back-to-back
	for i := 0; i < 10; i++ {
		t.Run(t.Name(), func(t *testing.T) {

			// Start a timeout timer with a timeout higher than the return timer
			timeOutTimer := time.NewTimer(time.Millisecond * 5)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err := WaitTimeout(returnChan, timeOutTimer)
			if err != nil {
				t.Fatal("returnChan did not unblock before timeOutTimer")
			}

			// Start a timeout timer with a lower timeout than return timer
			timeOutTimer = time.NewTimer(time.Millisecond * 1)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err = WaitTimeout(returnChan, timeOutTimer)
			if err != errTimeout {
				t.Fatal("timeOutTimer did not expire before returnChan unblock")
			}

			// Wait for the return timer to send on channel
			// This value needs to be read to prevent it from interfering other tests
			<-returnChan
		})
	}
}

func TestReadPacket(t *testing.T) {

	rt := []struct {
		name string
		tc   TestConn
		pkt  Packet
		err  error
	}{
		{
			name: "single and double escaped 0x07s in payload",
			tc: TestConn{
				Emits: []byte{ // Response
					esc, ack,
					0x07, 0xF0, // Frame Start
					0x00, uint8(getFans + 1), // response type
					0x06,       // Length
					0xAA, 0xBB, // in/out percents
					0x11, 0x07, 0x07, // in speed (one seven)
					0x07, 0x07, 0x07, 0x07, // out speed (two sevens)
					0x4A, // Checksum
					esc, end,
					esc, ack,
				},
			},
			pkt: Packet{
				Command: uint8(getFans + 1),
				Data:    []byte{0xAA, 0xBB, 0x11, 0x7, 0x7, 0x7},
			},
		},
		{
			name: "timeout reading first response ACK",
			tc: TestConn{
				HangStart: true,
				HangTime:  2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading start sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading end sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack, esc, start},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "garbage before ACK",
			tc: TestConn{
				Emits: []byte{esc, ack, 0xF0, 0x0B, esc, start},
			},
			err: errScanInput,
		},
		{
			name: "impossible packet length (need at least 4 bytes)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, esc, end},
			},
			err: errTooShort,
		},
		{
			name: "packet unmarshal error (any)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, 0xBA, 0xAA, esc, end},
			},
			err: errPayloadSize,
		},
	}

	for _, tt := range rt {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Start timeout timer
			to := time.NewTimer(time.Millisecond * 1)

			// Attempt to read packet from mock connection
			pkt, err := ReadPacket(&tc, to, true)

			if want, got := tt.err, err; want != got {
				t.Fatalf("unexpected error reading packet:\n- want: %v\n-  got: %v",
					want, got)
			}

			if want, got := tt.pkt, pkt; !reflect.DeepEqual(want, got) {
				t.Fatalf("unexpected packet read from connection:\n- want: %v\n-  got: %v",
					want, got)
			}
		})
	}
}

func TestWritePacket(t *testing.T) {

	wt := []struct {
		name string
		tc   TestConn
		pkt  Packet
		err  error
	}{
		{
			name: "oversized packet (handle MarshalPacket error)",
			pkt: Packet{
				Data: make([]byte, 256),
			},
			err: errTooLong,
		},
	}

	for _, tt := range wt {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Attempt to write packet to mock connection
			_, err := WritePacket(tt.pkt, &tc)

			if want, got := tt.err, err; want != got {
				t.Fatalf("unexpected error writing packet:\n- want: %v\n-  got: %v",
					want, got)
			}
		})
	}
}

func TestWriteAck(t *testing.T) {

	at := []struct {
		name string
		tc   TestConn
		err  bool
	}{
		{
			name: "writer not expecting input",
			tc:   TestConn{},
			err:  true,
		},
		{
			name: "incomplete ACK write",
			tc: TestConn{
				Receives: []byte{esc, ack},
				Limit:    1,
			},
			err: true,
		},
	}

	for _, tt := range at {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Attempt to write ACK to mock connection
			_, err := WriteAck(&tc)

			if tt.err != (err != nil) {
				t.Fatalf("unexpected error condition in ACK write:\n- want error: %v\n-  got error: %v",
					tt.err, err)
			}
		})
	}
}

func TestQueryPacket(t *testing.T) {
	at := []struct {
		name   string
		pktIn  Packet
		tc     TestConn
		pktOut Packet
		err    error
	}{
		{
			name: "invalid response command type (0x00)",
			pktIn: Packet{
				Command: uint8(getFans),
				Expect:  true,
			},
			tc: TestConn{
				Receives: []byte{esc, start, 0x00, uint8(getFans), 0x00, 0xB8, esc, end},
				Emits:    []byte{esc, ack, esc, start, 0x00, 0x00, 0x00, 0xAD, esc, end},
			},
			err: errInvalidResponse,
			pktOut: Packet{
				Data: []byte{},
			},
		},
	}

	for _, tt := range at {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			pktOut, err := QueryPacket(tt.pktIn, &tc)

			if want, got := tt.err, err; want != got {
				t.Fatalf("unexpected error during packet query:\n- want: %v\n-  got: %v",
					want, got)
			}

			if want, got := tt.pktOut, pktOut; !reflect.DeepEqual(want, got) {
				t.Fatalf("unexpected packet query result:\n- want: %v\n-  got: %v",
					want, got)
			}
		})
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Netstack Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package channel provides the implemention of channel-based data-link layer
// endpoints. Such endpoints allow injection of inbound packets and store
// outbound packets in a channel.
package channel

import (
	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/buffer"
	"github.com/google/netstack/tcpip/stack"
)

// PacketInfo holds all the information about an outbound packet.
type PacketInfo struct {
	Header  buffer.View
	Payload buffer.View
	Proto   tcpip.NetworkProtocolNumber
}

// Endpoint is link layer endpoint that stores outbound packets in a channel
// and allows injection of inbound packets.
type Endpoint struct {
	dispatcher stack.NetworkDispatcher
	mtu        uint32

	// C is where outbound packets are queued.
	C chan PacketInfo
}

// New creates a new channel endpoint.
func New(size int, mtu uint32) (tcpip.LinkEndpointID, *Endpoint) {
	e := &Endpoint{
		C:   make(chan PacketInfo, size),
		mtu: mtu,
	}

	return stack.RegisterLinkEndpoint(e), e
}

// Drain removes all outbound packets from the channel and counts them.
func (e *Endpoint) Drain() int {
	c := 0
	for {
		select {
		case <-e.C:
			c++
		default:
			return c
		}
	}
}

// Inject injects an inbound packet.
func (e *Endpoint) Inject(protocol tcpip.NetworkProtocolNumber, vv *buffer.VectorisedView) {
	uu := vv.Clone(nil)
	e.dispatcher.DeliverNetworkPacket(e, protocol, &uu)
}

// Attach saves the stack network-layer dispatcher for use later when packets
// are injected.
func (e *Endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.dispatcher = dispatcher
}

// MTU implements stack.LinkEndpoint.MTU. It returns the value initialized
// during construction.
func (e *Endpoint) MTU() uint32 {
	return e.mtu
}

// MaxHeaderLength returns the maximum size of the link layer header. Given it
// doesn't have a header, it just returns 0.
func (*Endpoint) MaxHeaderLength() uint16 {
	return 0
}

// WritePacket stores outbound packets into the channel.
func (e *Endpoint) WritePacket(_ *stack.Route, hdr *buffer.Prependable, payload buffer.View, protocol tcpip.NetworkProtocolNumber) error {
	p := PacketInfo{
		Header: hdr.View(),
		Proto:  protocol,
	}

	if payload != nil {
		p.Payload = payload
	}

	select {
	case e.C <- p:
	default:
	}

	return nil
}
<commit_msg>channel: copy payload on WritePacket<commit_after>// Copyright 2016 The Netstack Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package channel provides the implemention of channel-based data-link layer
// endpoints. Such endpoints allow injection of inbound packets and store
// outbound packets in a channel.
package channel

import (
	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/buffer"
	"github.com/google/netstack/tcpip/stack"
)

// PacketInfo holds all the information about an outbound packet.
type PacketInfo struct {
	Header  buffer.View
	Payload buffer.View
	Proto   tcpip.NetworkProtocolNumber
}

// Endpoint is link layer endpoint that stores outbound packets in a channel
// and allows injection of inbound packets.
type Endpoint struct {
	dispatcher stack.NetworkDispatcher
	mtu        uint32

	// C is where outbound packets are queued.
	C chan PacketInfo
}

// New creates a new channel endpoint.
func New(size int, mtu uint32) (tcpip.LinkEndpointID, *Endpoint) {
	e := &Endpoint{
		C:   make(chan PacketInfo, size),
		mtu: mtu,
	}

	return stack.RegisterLinkEndpoint(e), e
}

// Drain removes all outbound packets from the channel and counts them.
func (e *Endpoint) Drain() int {
	c := 0
	for {
		select {
		case <-e.C:
			c++
		default:
			return c
		}
	}
}

// Inject injects an inbound packet.
func (e *Endpoint) Inject(protocol tcpip.NetworkProtocolNumber, vv *buffer.VectorisedView) {
	uu := vv.Clone(nil)
	e.dispatcher.DeliverNetworkPacket(e, protocol, &uu)
}

// Attach saves the stack network-layer dispatcher for use later when packets
// are injected.
func (e *Endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.dispatcher = dispatcher
}

// MTU implements stack.LinkEndpoint.MTU. It returns the value initialized
// during construction.
func (e *Endpoint) MTU() uint32 {
	return e.mtu
}

// MaxHeaderLength returns the maximum size of the link layer header. Given it
// doesn't have a header, it just returns 0.
func (*Endpoint) MaxHeaderLength() uint16 {
	return 0
}

// WritePacket stores outbound packets into the channel.
func (e *Endpoint) WritePacket(_ *stack.Route, hdr *buffer.Prependable, payload buffer.View, protocol tcpip.NetworkProtocolNumber) error {
	p := PacketInfo{
		Header: hdr.View(),
		Proto:  protocol,
	}

	if payload != nil {
		p.Payload = make(buffer.View, len(payload))
		copy(p.Payload, payload)
	}

	select {
	case e.C <- p:
	default:
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Netstack Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package sniffer provides the implementation of data-link layer endpoints that
// wrap another endpoint and logs inbound and outbound packets.
//
// Sniffer endpoints can be used in the networking stack by calling New(eID) to
// create a new endpoint, where eID is the ID of the endpoint being wrapped,
// and then passing it as an argument to Stack.CreateNIC().
package sniffer

import (
	"fmt"
	"sync/atomic"

	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/buffer"
	"github.com/google/netstack/tcpip/header"
	"github.com/google/netstack/tcpip/stack"
	"log"
)

// LogPackets is a flag used to enable or disable packet valid values
// are 0 or 1.
var LogPackets uint32 = 1

type endpoint struct {
	dispatcher stack.NetworkDispatcher
	lower      stack.LinkEndpoint
}

// New creates a new sniffer link-layer endpoint. It wraps around another
// endpoint and logs packets and they traverse the endpoint.
func New(lower tcpip.LinkEndpointID) tcpip.LinkEndpointID {
	return stack.RegisterLinkEndpoint(&endpoint{
		lower: stack.FindLinkEndpoint(lower),
	})
}

// DeliverNetworkPacket implements the stack.NetworkDispatcher interface. It is
// called by the link-layer endpoint being wrapped when a packet arrives, and
// logs the packet before forwarding to the actual dispatcher.
func (e *endpoint) DeliverNetworkPacket(linkEP stack.LinkEndpoint, remoteLinkAddr tcpip.LinkAddress, protocol tcpip.NetworkProtocolNumber, vv *buffer.VectorisedView) {
	if atomic.LoadUint32(&LogPackets) == 1 {
		LogPacket("recv", protocol, vv.First(), nil)
	}
	e.dispatcher.DeliverNetworkPacket(e, remoteLinkAddr, protocol, vv)
}

// Attach implements the stack.LinkEndpoint interface. It saves the dispatcher
// and registers with the lower endpoint as its dispatcher so that "e" is called
// for inbound packets.
func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.dispatcher = dispatcher
	e.lower.Attach(e)
}

// MTU implements stack.LinkEndpoint.MTU. It just forwards the request to the
// lower endpoint.
func (e *endpoint) MTU() uint32 {
	return e.lower.MTU()
}

// MaxHeaderLength implements the stack.LinkEndpoint interface. It just forwards
// the request to the lower endpoint.
func (e *endpoint) MaxHeaderLength() uint16 {
	return e.lower.MaxHeaderLength()
}

func (e *endpoint) LinkAddress() tcpip.LinkAddress {
	return e.lower.LinkAddress()
}

// WritePacket implements the stack.LinkEndpoint interface. It is called by
// higher-level protocols to write packets; it just logs the packet and forwards
// the request to the lower endpoint.
func (e *endpoint) WritePacket(r *stack.Route, hdr *buffer.Prependable, payload buffer.View, protocol tcpip.NetworkProtocolNumber) *tcpip.Error {
	if atomic.LoadUint32(&LogPackets) == 1 {
		LogPacket("send", protocol, hdr.UsedBytes(), payload)
	}
	return e.lower.WritePacket(r, hdr, payload, protocol)
}

// LogPacket logs the given packet.
func LogPacket(prefix string, protocol tcpip.NetworkProtocolNumber, b, plb []byte) {
	// Figure out the network layer info.
	var transProto uint8
	src := tcpip.Address("unknown")
	dst := tcpip.Address("unknown")
	id := 0
	size := uint16(0)
	switch protocol {
	case header.IPv4ProtocolNumber:
		ipv4 := header.IPv4(b)
		src = ipv4.SourceAddress()
		dst = ipv4.DestinationAddress()
		transProto = ipv4.Protocol()
		size = ipv4.TotalLength() - uint16(ipv4.HeaderLength())
		b = b[ipv4.HeaderLength():]
		id = int(ipv4.ID())

	case header.IPv6ProtocolNumber:
		ipv6 := header.IPv6(b)
		src = ipv6.SourceAddress()
		dst = ipv6.DestinationAddress()
		transProto = ipv6.NextHeader()
		size = ipv6.PayloadLength()
		b = b[header.IPv6MinimumSize:]

	case header.ARPProtocolNumber:
		arp := header.ARP(b)
		log.Printf(
			"%s arp %v (%v) -> %v (%v) valid:%v",
			prefix,
			tcpip.Address(arp.ProtocolAddressSender()), tcpip.LinkAddress(arp.HardwareAddressSender()),
			tcpip.Address(arp.ProtocolAddressTarget()), tcpip.LinkAddress(arp.HardwareAddressTarget()),
			arp.IsValid(),
		)
		return
	default:
		log.Printf("%s unknown network protocol", prefix)
		return
	}

	// Figure out the transport layer info.
	transName := "unknown"
	srcPort := uint16(0)
	dstPort := uint16(0)
	details := ""
	switch tcpip.TransportProtocolNumber(transProto) {
	case header.ICMPv4ProtocolNumber:
		transName = "icmp"
		icmp := header.ICMPv4(b)
		icmpType := "unknown"
		switch icmp.Type() {
		case header.ICMPv4EchoReply:
			icmpType = "echo reply"
		case header.ICMPv4DstUnreachable:
			icmpType = "destination unreachable"
		case header.ICMPv4SrcQuench:
			icmpType = "source quench"
		case header.ICMPv4Redirect:
			icmpType = "redirect"
		case header.ICMPv4Echo:
			icmpType = "echo"
		case header.ICMPv4TimeExceeded:
			icmpType = "time exceeded"
		case header.ICMPv4ParamProblem:
			icmpType = "param problem"
		case header.ICMPv4Timestamp:
			icmpType = "timestamp"
		case header.ICMPv4TimestampReply:
			icmpType = "timestamp reply"
		case header.ICMPv4InfoRequest:
			icmpType = "info request"
		case header.ICMPv4InfoReply:
			icmpType = "info reply"
		}
		log.Printf("%s %s %v -> %v %s len:%d id:%04x code:%d", prefix, transName, src, dst, icmpType, size, id, icmp.Code())
		return

	case header.UDPProtocolNumber:
		transName = "udp"
		udp := header.UDP(b)
		srcPort = udp.SourcePort()
		dstPort = udp.DestinationPort()
		size -= header.UDPMinimumSize

		details = fmt.Sprintf("xsum: 0x%x", udp.Checksum())

	case header.TCPProtocolNumber:
		transName = "tcp"
		tcp := header.TCP(b)
		srcPort = tcp.SourcePort()
		dstPort = tcp.DestinationPort()
		size -= uint16(tcp.DataOffset())

		// Initialize the TCP flags.
		flags := tcp.Flags()
		flagsStr := []byte("FSRPAU")
		for i := range flagsStr {
			if flags&(1<<uint(i)) == 0 {
				flagsStr[i] = ' '
			}
		}

		details = fmt.Sprintf("flags:0x%02x (%v) seqnum: %v ack: %v win: %v xsum:0x%x", flags, string(flagsStr), tcp.SequenceNumber(), tcp.AckNumber(), tcp.WindowSize(), tcp.Checksum())

	default:
		log.Printf("%s %v -> %v unknown transport protocol: %d", prefix, src, dst, transProto)
		return
	}

	log.Printf("%s %s %v:%v -> %v:%v len:%d id:%04x %s", prefix, transName, src, srcPort, dst, dstPort, size, id, details)
}
<commit_msg>Adds support to log TCP options when logging packets.<commit_after>// Copyright 2016 The Netstack Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package sniffer provides the implementation of data-link layer endpoints that
// wrap another endpoint and logs inbound and outbound packets.
//
// Sniffer endpoints can be used in the networking stack by calling New(eID) to
// create a new endpoint, where eID is the ID of the endpoint being wrapped,
// and then passing it as an argument to Stack.CreateNIC().
package sniffer

import (
	"fmt"
	"sync/atomic"

	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/buffer"
	"github.com/google/netstack/tcpip/header"
	"github.com/google/netstack/tcpip/stack"
	"log"
)

// LogPackets is a flag used to enable or disable packet valid values
// are 0 or 1.
var LogPackets uint32 = 1

type endpoint struct {
	dispatcher stack.NetworkDispatcher
	lower      stack.LinkEndpoint
}

// New creates a new sniffer link-layer endpoint. It wraps around another
// endpoint and logs packets and they traverse the endpoint.
func New(lower tcpip.LinkEndpointID) tcpip.LinkEndpointID {
	return stack.RegisterLinkEndpoint(&endpoint{
		lower: stack.FindLinkEndpoint(lower),
	})
}

// DeliverNetworkPacket implements the stack.NetworkDispatcher interface. It is
// called by the link-layer endpoint being wrapped when a packet arrives, and
// logs the packet before forwarding to the actual dispatcher.
func (e *endpoint) DeliverNetworkPacket(linkEP stack.LinkEndpoint, remoteLinkAddr tcpip.LinkAddress, protocol tcpip.NetworkProtocolNumber, vv *buffer.VectorisedView) {
	if atomic.LoadUint32(&LogPackets) == 1 {
		LogPacket("recv", protocol, vv.First(), nil)
	}
	e.dispatcher.DeliverNetworkPacket(e, remoteLinkAddr, protocol, vv)
}

// Attach implements the stack.LinkEndpoint interface. It saves the dispatcher
// and registers with the lower endpoint as its dispatcher so that "e" is called
// for inbound packets.
func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.dispatcher = dispatcher
	e.lower.Attach(e)
}

// MTU implements stack.LinkEndpoint.MTU. It just forwards the request to the
// lower endpoint.
func (e *endpoint) MTU() uint32 {
	return e.lower.MTU()
}

// MaxHeaderLength implements the stack.LinkEndpoint interface. It just forwards
// the request to the lower endpoint.
func (e *endpoint) MaxHeaderLength() uint16 {
	return e.lower.MaxHeaderLength()
}

func (e *endpoint) LinkAddress() tcpip.LinkAddress {
	return e.lower.LinkAddress()
}

// WritePacket implements the stack.LinkEndpoint interface. It is called by
// higher-level protocols to write packets; it just logs the packet and forwards
// the request to the lower endpoint.
func (e *endpoint) WritePacket(r *stack.Route, hdr *buffer.Prependable, payload buffer.View, protocol tcpip.NetworkProtocolNumber) *tcpip.Error {
	if atomic.LoadUint32(&LogPackets) == 1 {
		LogPacket("send", protocol, hdr.UsedBytes(), payload)
	}
	return e.lower.WritePacket(r, hdr, payload, protocol)
}

// LogPacket logs the given packet.
func LogPacket(prefix string, protocol tcpip.NetworkProtocolNumber, b, plb []byte) {
	// Figure out the network layer info.
	var transProto uint8
	src := tcpip.Address("unknown")
	dst := tcpip.Address("unknown")
	id := 0
	size := uint16(0)
	switch protocol {
	case header.IPv4ProtocolNumber:
		ipv4 := header.IPv4(b)
		src = ipv4.SourceAddress()
		dst = ipv4.DestinationAddress()
		transProto = ipv4.Protocol()
		size = ipv4.TotalLength() - uint16(ipv4.HeaderLength())
		b = b[ipv4.HeaderLength():]
		id = int(ipv4.ID())

	case header.IPv6ProtocolNumber:
		ipv6 := header.IPv6(b)
		src = ipv6.SourceAddress()
		dst = ipv6.DestinationAddress()
		transProto = ipv6.NextHeader()
		size = ipv6.PayloadLength()
		b = b[header.IPv6MinimumSize:]

	case header.ARPProtocolNumber:
		arp := header.ARP(b)
		log.Printf(
			"%s arp %v (%v) -> %v (%v) valid:%v",
			prefix,
			tcpip.Address(arp.ProtocolAddressSender()), tcpip.LinkAddress(arp.HardwareAddressSender()),
			tcpip.Address(arp.ProtocolAddressTarget()), tcpip.LinkAddress(arp.HardwareAddressTarget()),
			arp.IsValid(),
		)
		return
	default:
		log.Printf("%s unknown network protocol", prefix)
		return
	}

	// Figure out the transport layer info.
	transName := "unknown"
	srcPort := uint16(0)
	dstPort := uint16(0)
	details := ""
	switch tcpip.TransportProtocolNumber(transProto) {
	case header.ICMPv4ProtocolNumber:
		transName = "icmp"
		icmp := header.ICMPv4(b)
		icmpType := "unknown"
		switch icmp.Type() {
		case header.ICMPv4EchoReply:
			icmpType = "echo reply"
		case header.ICMPv4DstUnreachable:
			icmpType = "destination unreachable"
		case header.ICMPv4SrcQuench:
			icmpType = "source quench"
		case header.ICMPv4Redirect:
			icmpType = "redirect"
		case header.ICMPv4Echo:
			icmpType = "echo"
		case header.ICMPv4TimeExceeded:
			icmpType = "time exceeded"
		case header.ICMPv4ParamProblem:
			icmpType = "param problem"
		case header.ICMPv4Timestamp:
			icmpType = "timestamp"
		case header.ICMPv4TimestampReply:
			icmpType = "timestamp reply"
		case header.ICMPv4InfoRequest:
			icmpType = "info request"
		case header.ICMPv4InfoReply:
			icmpType = "info reply"
		}
		log.Printf("%s %s %v -> %v %s len:%d id:%04x code:%d", prefix, transName, src, dst, icmpType, size, id, icmp.Code())
		return

	case header.UDPProtocolNumber:
		transName = "udp"
		udp := header.UDP(b)
		srcPort = udp.SourcePort()
		dstPort = udp.DestinationPort()
		size -= header.UDPMinimumSize

		details = fmt.Sprintf("xsum: 0x%x", udp.Checksum())

	case header.TCPProtocolNumber:
		transName = "tcp"
		tcp := header.TCP(b)
		srcPort = tcp.SourcePort()
		dstPort = tcp.DestinationPort()
		size -= uint16(tcp.DataOffset())

		// Initialize the TCP flags.
		flags := tcp.Flags()
		flagsStr := []byte("FSRPAU")
		for i := range flagsStr {
			if flags&(1<<uint(i)) == 0 {
				flagsStr[i] = ' '
			}
		}
		details = fmt.Sprintf("flags:0x%02x (%v) seqnum: %v ack: %v win: %v xsum:0x%x", flags, string(flagsStr), tcp.SequenceNumber(), tcp.AckNumber(), tcp.WindowSize(), tcp.Checksum())
		if flags&header.TCPFlagSyn != 0 {
			details += fmt.Sprintf(" options: %+v", header.ParseSynOptions(tcp.Options(), flags&header.TCPFlagAck != 0))
		} else {
			details += fmt.Sprintf(" options: %+v", tcp.ParsedOptions())
		}
	default:
		log.Printf("%s %v -> %v unknown transport protocol: %d", prefix, src, dst, transProto)
		return
	}

	log.Printf("%s %s %v:%v -> %v:%v len:%d id:%04x %s", prefix, transName, src, srcPort, dst, dstPort, size, id, details)
}
<|endoftext|>


<commit_before>package curler

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
)

// New creates a new curler
func New(orig http.Handler, out io.Writer) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		parts := []string{"curl", "-i", "-X", r.Method}

		for k, v := range r.Header {
			ck := http.CanonicalHeaderKey(k)
			if ck != "Host" && ck != "User-Agent" {
				parts = append(parts, "-H", fmt.Sprintf("'%s: %s'", k, strings.Join(v, ",")))
			}
		}

		if r.Method == "POST" || r.Method == "PATCH" || r.Method == "PUT" {
			b, err := ioutil.ReadAll(r.Body)

			if err != nil {
				fmt.Fprintln(out, "curler:", err)
			}
			if err == nil && len(b) > 0 {
				r.Body = ioutil.NopCloser(bytes.NewBuffer(b))
				parts = append(parts, "-d '", string(b)+"'")
			}
		}

		scheme := "http"
		if r.TLS != nil {
			scheme += "s"
		}

		parts = append(parts, scheme+"://"+r.Host+r.URL.String())
		fmt.Fprintln(out, strings.Join(parts, " "))
		orig.ServeHTTP(rw, r)
	})
}
<commit_msg>fix repeated headers<commit_after>package curler

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
)

// New creates a new curler
func New(orig http.Handler, out io.Writer) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		parts := []string{"curl", "-i", "-X", r.Method}

		for k, v := range r.Header {
			ck := http.CanonicalHeaderKey(k)
			if ck != "Host" && ck != "User-Agent" {
				for _, vv := range v {
					parts = append(parts, "-H", fmt.Sprintf("'%s: %s'", k, vv))
				}

			}
		}

		if r.Method == "POST" || r.Method == "PATCH" || r.Method == "PUT" {
			b, err := ioutil.ReadAll(r.Body)

			if err != nil {
				fmt.Fprintln(out, "curler:", err)
			}
			if err == nil && len(b) > 0 {
				r.Body = ioutil.NopCloser(bytes.NewBuffer(b))
				parts = append(parts, "-d '", string(b)+"'")
			}
		}

		scheme := "http"
		if r.TLS != nil {
			scheme += "s"
		}

		parts = append(parts, scheme+"://"+r.Host+r.URL.String())
		fmt.Fprintln(out, strings.Join(parts, " "))
		orig.ServeHTTP(rw, r)
	})
}
<|endoftext|>


<commit_before>package immortal

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"sync/atomic"
	"syscall"
	"time"
)

// Daemon struct
type Daemon struct {
	cfg            *Config
	count          uint64
	lock, lockOnce uint32
	process        *process
	quit           chan struct{}
	sTime          time.Time
	supDir         string
}

// Run returns a process instance
func (d *Daemon) Run(p Process) (*process, error) {
	if atomic.SwapUint32(&d.lock, uint32(1)) != 0 {
		return nil, fmt.Errorf("lock: %d lock once: %d", d.lock, d.lockOnce)
	}

	var err error

	// increment count by 1
	atomic.AddUint64(&d.count, 1)

	time.Sleep(time.Duration(d.cfg.Wait) * time.Second)

	d.process, err = p.Start()
	if err != nil {
		atomic.StoreUint32(&d.lock, d.lockOnce)
		return nil, err
	}

	// write parent pid
	if d.cfg.Pid.Parent != "" {
		if err := d.WritePid(d.cfg.Pid.Parent, os.Getpid()); err != nil {
			log.Println(err)
		}
	}

	// write child pid
	if d.cfg.Pid.Child != "" {
		if err := d.WritePid(d.cfg.Pid.Child, p.Pid()); err != nil {
			log.Println(err)
		}
	}

	return d.process, nil
}

// WritePid write pid to file
func (d *Daemon) WritePid(file string, pid int) error {
	if err := ioutil.WriteFile(file, []byte(fmt.Sprintf("%d", pid)), 0644); err != nil {
		return err
	}
	return nil
}

// New creates a new daemon
func New(cfg *Config) (*Daemon, error) {
	var supDir string

	if cfg.ctl {
		d, err := os.Getwd()
		if err != nil {
			return nil, err
		}
		supDir = filepath.Join(d, "supervise")
	} else {
		// create an .immotal dir on HOME user when calling immortal directly
		// and not using immortal-dir, this helps to run immortal-ctl and
		// check status of all daemons
		usr, err := user.Current()
		if err != nil {
			return nil, err
		}
		supDir = filepath.Join(usr.HomeDir,
			".immortal",
			fmt.Sprintf("%d", os.Getpid()),
			"supervise")
	}

	// create supervise dir
	if err := os.MkdirAll(supDir, os.ModePerm); err != nil {
		return nil, err
	}

	// lock
	if lock, err := os.Create(filepath.Join(supDir, "lock")); err != nil {
		return nil, err
	} else if err = syscall.Flock(int(lock.Fd()), syscall.LOCK_EX+syscall.LOCK_NB); err != nil {
		return nil, err
	}

	// remove previous socket in case exists
	os.Remove(filepath.Join(supDir, "immortal.sock"))

	return &Daemon{
		cfg:    cfg,
		supDir: supDir,
		quit:   make(chan struct{}),
		sTime:  time.Now(),
	}, nil
}
<commit_msg>	modified:   daemon.go<commit_after>package immortal

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"sync/atomic"
	"syscall"
	"time"
)

// Daemon struct
type Daemon struct {
	cfg            *Config
	count          uint64
	lock, lockOnce uint32
	process        *process
	quit           chan struct{}
	sTime          time.Time
	supDir         string
}

// Run returns a process instance
func (d *Daemon) Run(p Process) (*process, error) {
	if atomic.SwapUint32(&d.lock, uint32(1)) != 0 {
		return nil, fmt.Errorf("lock: %d lock once: %d", d.lock, d.lockOnce)
	}

	var err error

	// increment count by 1
	atomic.AddUint64(&d.count, 1)

	time.Sleep(time.Duration(d.cfg.Wait) * time.Second)

	d.process, err = p.Start()
	if err != nil {
		atomic.StoreUint32(&d.lock, d.lockOnce)
		return nil, err
	}

	// write parent pid
	if d.cfg.Pid.Parent != "" {
		if err := d.WritePid(d.cfg.Pid.Parent, os.Getpid()); err != nil {
			log.Println(err)
		}
	}

	// write child pid
	if d.cfg.Pid.Child != "" {
		if err := d.WritePid(d.cfg.Pid.Child, p.Pid()); err != nil {
			log.Println(err)
		}
	}

	return d.process, nil
}

// WritePid write pid to file
func (d *Daemon) WritePid(file string, pid int) error {
	if err := ioutil.WriteFile(file, []byte(fmt.Sprintf("%d", pid)), 0644); err != nil {
		return err
	}
	return nil
}

// New creates a new daemon
func New(cfg *Config) (*Daemon, error) {
	var supDir string

	// create supervise directory in current directory
	if cfg.ctl {
		d, err := os.Getwd()
		if err != nil {
			return nil, err
		}
		supDir = filepath.Join(d, "supervise")
	} else {
		// create an .immotal dir on HOME user when calling immortal directly
		// and not using immortal-dir, this helps to run immortal-ctl and
		// check status of all daemons
		usr, err := user.Current()
		if err != nil {
			return nil, err
		}
		supDir = filepath.Join(usr.HomeDir,
			".immortal",
			fmt.Sprintf("%d", os.Getpid()),
			"supervise")
	}

	// create supervise dir
	if err := os.MkdirAll(supDir, os.ModePerm); err != nil {
		return nil, err
	}

	// lock
	if lock, err := os.Create(filepath.Join(supDir, "lock")); err != nil {
		return nil, err
	} else if err = syscall.Flock(int(lock.Fd()), syscall.LOCK_EX+syscall.LOCK_NB); err != nil {
		return nil, err
	}

	// remove previous socket in case exists
	os.Remove(filepath.Join(supDir, "immortal.sock"))

	return &Daemon{
		cfg:    cfg,
		supDir: supDir,
		quit:   make(chan struct{}),
		sTime:  time.Now(),
	}, nil
}
<|endoftext|>


<commit_before>package db

import (
	"fmt"

	"github.com/jinzhu/gorm"
	"github.com/knqyf263/gost/models"
	// Required MySQL.  See http://jinzhu.me/gorm/database.html#connecting-to-a-database
	_ "github.com/jinzhu/gorm/dialects/mysql"
	_ "github.com/jinzhu/gorm/dialects/postgres"
	// Required SQLite3.
	_ "github.com/jinzhu/gorm/dialects/sqlite"
)

// Supported DB dialects.
const (
	dialectSqlite3    = "sqlite3"
	dialectMysql      = "mysql"
	dialectPostgreSQL = "postgres"
)

// RDBDriver is Driver for RDB
type RDBDriver struct {
	name string
	conn *gorm.DB
}

// Name return db name
func (r *RDBDriver) Name() string {
	return r.name
}

// OpenDB opens Database
func (r *RDBDriver) OpenDB(dbType, dbPath string, debugSQL bool) (err error) {
	r.conn, err = gorm.Open(dbType, dbPath)
	if err != nil {
		err = fmt.Errorf("Failed to open DB. dbtype: %s, dbpath: %s, err: %s", dbType, dbPath, err)
		return
	}
	r.conn.LogMode(debugSQL)
	if r.name == dialectSqlite3 {
		r.conn.Exec("PRAGMA foreign_keys = ON")
		r.conn.Exec("PRAGMA journal_mode=WAL;")
	}
	return
}

// MigrateDB migrates Database
func (r *RDBDriver) MigrateDB() error {
	if err := r.conn.AutoMigrate(
		&models.RedhatCVE{},
		&models.RedhatDetail{},
		&models.RedhatReference{},
		&models.RedhatBugzilla{},
		&models.RedhatCvss{},
		&models.RedhatCvss3{},
		&models.RedhatAffectedRelease{},
		&models.RedhatPackageState{},

		&models.DebianCVE{},
		&models.DebianPackage{},
		&models.DebianRelease{},
	).Error; err != nil {
		return fmt.Errorf("Failed to migrate. err: %s", err)
	}
	return nil
}
<commit_msg>fix(db): no wal mode with SQLite3 backend to avoid `database is locked`<commit_after>package db

import (
	"fmt"

	"github.com/jinzhu/gorm"
	"github.com/knqyf263/gost/models"
	// Required MySQL.  See http://jinzhu.me/gorm/database.html#connecting-to-a-database
	_ "github.com/jinzhu/gorm/dialects/mysql"
	_ "github.com/jinzhu/gorm/dialects/postgres"
	// Required SQLite3.
	_ "github.com/jinzhu/gorm/dialects/sqlite"
)

// Supported DB dialects.
const (
	dialectSqlite3    = "sqlite3"
	dialectMysql      = "mysql"
	dialectPostgreSQL = "postgres"
)

// RDBDriver is Driver for RDB
type RDBDriver struct {
	name string
	conn *gorm.DB
}

// Name return db name
func (r *RDBDriver) Name() string {
	return r.name
}

// OpenDB opens Database
func (r *RDBDriver) OpenDB(dbType, dbPath string, debugSQL bool) (err error) {
	r.conn, err = gorm.Open(dbType, dbPath)
	if err != nil {
		err = fmt.Errorf("Failed to open DB. dbtype: %s, dbpath: %s, err: %s", dbType, dbPath, err)
		return
	}
	r.conn.LogMode(debugSQL)
	if r.name == dialectSqlite3 {
		r.conn.Exec("PRAGMA foreign_keys = ON")
	}
	return
}

// MigrateDB migrates Database
func (r *RDBDriver) MigrateDB() error {
	if err := r.conn.AutoMigrate(
		&models.RedhatCVE{},
		&models.RedhatDetail{},
		&models.RedhatReference{},
		&models.RedhatBugzilla{},
		&models.RedhatCvss{},
		&models.RedhatCvss3{},
		&models.RedhatAffectedRelease{},
		&models.RedhatPackageState{},

		&models.DebianCVE{},
		&models.DebianPackage{},
		&models.DebianRelease{},
	).Error; err != nil {
		return fmt.Errorf("Failed to migrate. err: %s", err)
	}
	return nil
}
<|endoftext|>


<commit_before>package pcap

/*
#include <pcap.h>
*/
import "C"

import (
	"fmt"
	"net"
	"strings"
	"syscall"
)

// Type constants.
const (
	TypeIP    = syscall.ETH_P_IP
	TypeARP   = syscall.ETH_P_ARP
	TypeIP6   = syscall.ETH_P_IPV6
	TypeEAPOL = syscall.ETH_P_PAE

	IPICMP   = syscall.IPPROTO_ICMP
	IPInIP   = syscall.IPPROTO_IPIP
	IPTCP    = syscall.IPPROTO_TCP
	IPUDP    = syscall.IPPROTO_UDP
	IPICMPv6 = syscall.IPPROTO_ICMPV6
)

// Port from sf-pcap.c file.
const (
	TcpdumpMagic          = 0xa1b2c3d4
	KuznetsovTcpdumpMagic = 0xa1b2cd34
	FmesqiutaTcpdumpMagic = 0xa1b234cd
	NavtelTcpdumpMagic    = 0xa12b3c4d
	NsecTcpdumpMagic      = 0xa1b23c4d
)

// DLT* are the data link types. They can differ per platform.
// See http://www.tcpdump.org/linktypes.html
const (
	DLTNULL     = C.DLT_NULL
	DLTRAW      = C.DLT_RAW
	DLTEN10MB   = C.DLT_EN10MB    // Ethernet (10Mb)
	DLTEN3MB    = C.DLT_EN3MB     // Experimental Ethernet (3Mb)
	DLTAX25     = C.DLT_AX25      // Amateur Radio AX.25
	DLTPRONET   = C.DLT_PRONET    // Proteon ProNET Token Ring
	DLTCHAOS    = C.DLT_CHAOS     // Chaos
	DLTIEEE802  = C.DLT_IEEE802   // 802.5 Token Ring
	DLTARCNET   = C.DLT_ARCNET    // ARCNET, with BSD-style header
	DLTSLIP     = C.DLT_SLIP      // Serial Line IP
	DLTPPP      = C.DLT_PPP       // Point-to-point Protocol
	DLTFDDI     = C.DLT_FDDI      // FDDI
	DLTLINUXSSL = C.DLT_LINUX_SLL // Linux cooked
)

const errbufSize = 256

// LinkType constants according to pcap-linktype(7).
const (
	LinkTypeNull      = DLTNULL
	LinkTypeEthernet  = DLTEN10MB
	LinkTypeTokenRing = DLTIEEE802

	LinkTypeExpEthernet = DLTEN3MB /* 3Mb experimental Ethernet */
	LinkTypeAX25        = DLTAX25
	LinkTypePRONET      = DLTPRONET
	LinkTypeCHAOS       = DLTCHAOS
	LinkTypeARCNETBSD   = DLTARCNET /* BSD-style headers */
	LinkTypeSLIP        = DLTSLIP
	LinkTypePPP         = DLTPPP
	LinkTypeFDDI        = DLTFDDI

	LinkTypeARCNET         = 7
	LinkTypeATMRFC1483     = 100
	LinkTypeRAW            = 101
	LinkTypePPPHDLC        = 50
	LinkTypePPPETHER       = 51
	LinkTypeCHDLC          = 104
	LinkTypeIEEE80211      = 105
	LinkTypeFRELAY         = 107
	LinkTypeLOOP           = 108
	LinkTypeLINUXSLL       = 113
	LinkTypeLTALK          = 104
	LinkTypePFLOG          = 117
	LinkTypePRISMHeader    = 119
	LinkTypeIPOverFC       = 122
	LinkTypeSUNATM         = 123
	LinkTypeIEEE80211Radio = 127
	LinkTypeARCNETLinux    = 129
	LinkTypeLinuxIRDA      = 144
	LinkTypeLinuxLAPD      = 177
)

type addrHdr interface {
	SrcAddr() string
	DestAddr() string
	Len() int
}

type addrStringer interface {
	String(addr addrHdr) string
}

func decodemac(pkt []byte) uint64 {
	mac := uint64(0)
	for i := 0; i < 6; i++ {
		mac = (mac << 8) + uint64(pkt[i])
	}
	return mac
}

// ARPHdr is a ARP packet header.
type ARPHdr struct {
	Addrtype          uint16
	Protocol          uint16
	HwAddressSize     uint8
	ProtAddressSize   uint8
	Operation         uint16
	SourceHwAddress   []byte
	SourceProtAddress []byte
	DestHwAddress     []byte
	DestProtAddress   []byte
}

func (arp *ARPHdr) String() (s string) {
	switch arp.Operation {
	case 1:
		s = "ARP request"
	case 2:
		s = "ARP Reply"
	}
	if arp.Addrtype == LinkTypeEthernet && arp.Protocol == TypeIP {
		s = fmt.Sprintf("%012x (%s) > %012x (%s)",
			decodemac(arp.SourceHwAddress), arp.SourceProtAddress,
			decodemac(arp.DestHwAddress), arp.DestProtAddress)
	} else {
		s = fmt.Sprintf("addrtype = %d protocol = %d", arp.Addrtype, arp.Protocol)
	}
	return
}

// Len returns the (calculated) ARP payload length.
func (arp *ARPHdr) Len() int {
	return int(8 + 2*arp.HwAddressSize + 2*arp.ProtAddressSize)
}

// IPHdr is the header of an IP packet.
type IPHdr struct {
	Version    uint8
	Ihl        uint8
	Tos        uint8
	Length     uint16
	ID         uint16
	Flags      uint8
	FragOffset uint16
	Ttl        uint8
	Protocol   uint8
	Checksum   uint16
	SrcIP      []byte
	DestIP     []byte
}

// SrcAddr returns the string version of the source IP.
func (ip *IPHdr) SrcAddr() string { return net.IP(ip.SrcIP).String() }

// DestAddr returns the string version of the destination IP.
func (ip *IPHdr) DestAddr() string { return net.IP(ip.DestIP).String() }

// Len returns the ip.Length.
func (ip *IPHdr) Len() int { return int(ip.Length) }

// TCPHdr is the header of a TCP packet.
type TCPHdr struct {
	SrcPort    uint16
	DestPort   uint16
	Seq        uint32
	Ack        uint32
	DataOffset uint8
	Flags      uint16
	Window     uint16
	Checksum   uint16
	Urgent     uint16
	Data       []byte
}

// TCP flags.
const (
	TCPFIN = 1 << iota
	TCPSYN
	TCPRST
	TCPPSH
	TCPACK
	TCPURG
	TCPECE
	TCPCWR
	TCPNS
)

// String TODO
func (tcp *TCPHdr) String(hdr addrHdr) string {
	return fmt.Sprintf("TCP %s:%d > %s:%d %s SEQ=%d ACK=%d LEN=%d",
		hdr.SrcAddr(), int(tcp.SrcPort), hdr.DestAddr(), int(tcp.DestPort),
		tcp.FlagsString(), int64(tcp.Seq), int64(tcp.Ack), hdr.Len())
}

// FlagsString TODO
func (tcp *TCPHdr) FlagsString() string {
	var sflags []string
	if 0 != (tcp.Flags & TCPSYN) {
		sflags = append(sflags, "syn")
	}
	if 0 != (tcp.Flags & TCPFIN) {
		sflags = append(sflags, "fin")
	}
	if 0 != (tcp.Flags & TCPACK) {
		sflags = append(sflags, "ack")
	}
	if 0 != (tcp.Flags & TCPPSH) {
		sflags = append(sflags, "psh")
	}
	if 0 != (tcp.Flags & TCPRST) {
		sflags = append(sflags, "rst")
	}
	if 0 != (tcp.Flags & TCPURG) {
		sflags = append(sflags, "urg")
	}
	if 0 != (tcp.Flags & TCPNS) {
		sflags = append(sflags, "ns")
	}
	if 0 != (tcp.Flags & TCPCWR) {
		sflags = append(sflags, "cwr")
	}
	if 0 != (tcp.Flags & TCPECE) {
		sflags = append(sflags, "ece")
	}
	return fmt.Sprintf("[%s]", strings.Join(sflags, " "))
}

// UDPHdr is the header of a UDP packet.
type UDPHdr struct {
	SrcPort  uint16
	DestPort uint16
	Length   uint16
	Checksum uint16
}

func (udp *UDPHdr) String(hdr addrHdr) string {
	return fmt.Sprintf("UDP %s:%d > %s:%d LEN=%d CHKSUM=%d",
		hdr.SrcAddr(), int(udp.SrcPort), hdr.DestAddr(), int(udp.DestPort),
		int(udp.Length), int(udp.Checksum))
}

// ICMPHdr is the header of an ICMP packet.
type ICMPHdr struct {
	Type     uint8
	Code     uint8
	Checksum uint16
	ID       uint16
	Seq      uint16
	Data     []byte
}

// String TODO
func (icmp *ICMPHdr) String(hdr addrHdr) string {
	return fmt.Sprintf("ICMP %s > %s Type = %d Code = %d ",
		hdr.SrcAddr(), hdr.DestAddr(), icmp.Type, icmp.Code)
}

// TypeString TODO
func (icmp *ICMPHdr) TypeString() (result string) {
	switch icmp.Type {
	case 0:
		result = fmt.Sprintf("Echo reply seq=%d", icmp.Seq)
	case 3:
		switch icmp.Code {
		case 0:
			result = "Network unreachable"
		case 1:
			result = "Host unreachable"
		case 2:
			result = "Protocol unreachable"
		case 3:
			result = "Port unreachable"
		default:
			result = "Destination unreachable"
		}
	case 8:
		result = fmt.Sprintf("Echo request seq=%d", icmp.Seq)
	case 30:
		result = "Traceroute"
	}
	return
}

// IP6Hdr is the header of an IPv6 packet.
type IP6Hdr struct {
	// http://www.networksorcery.com/enp/protocol/ipv6.htm
	Version      uint8  // 4 bits
	TrafficClass uint8  // 8 bits
	FlowLabel    uint32 // 20 bits
	Length       uint16 // 16 bits
	NextHeader   uint8  // 8 bits, same as Protocol in IPHdr
	HopLimit     uint8  // 8 bits
	SrcIP        []byte // 16 bytes
	DestIP       []byte // 16 bytes
}

// SrcAddr returns the string version of the source IP.
func (ip6 *IP6Hdr) SrcAddr() string { return net.IP(ip6.SrcIP).String() }

// DestAddr returns the string version of the destination IP.
func (ip6 *IP6Hdr) DestAddr() string { return net.IP(ip6.DestIP).String() }

// Len returns the ip6.Length.
func (ip6 *IP6Hdr) Len() int { return int(ip6.Length) }

// ICMPv6Hdr is for ICMPv6
type ICMPv6Hdr struct {
	Type     uint8
	Code     uint8
	Checksum uint16
}
<commit_msg>Same same<commit_after>package pcap

/*
#include <pcap.h>
*/
import "C"

import (
	"fmt"
	"net"
	"strings"
	"syscall"
)

// Type constants.
const (
	TypeIP    = syscall.ETH_P_IP
	TypeARP   = syscall.ETH_P_ARP
	TypeIP6   = syscall.ETH_P_IPV6
	TypeEAPOL = syscall.ETH_P_PAE

	IPICMP   = syscall.IPPROTO_ICMP
	IPInIP   = syscall.IPPROTO_IPIP
	IPTCP    = syscall.IPPROTO_TCP
	IPUDP    = syscall.IPPROTO_UDP
	IPICMPv6 = syscall.IPPROTO_ICMPV6
)

// Port from sf-pcap.c file.
const (
	TcpdumpMagic          = 0xa1b2c3d4
	KuznetsovTcpdumpMagic = 0xa1b2cd34
	FmesqiutaTcpdumpMagic = 0xa1b234cd
	NavtelTcpdumpMagic    = 0xa12b3c4d
	NsecTcpdumpMagic      = 0xa1b23c4d
)

// DLT* are the data link types. They can differ per platform.
// See http://www.tcpdump.org/linktypes.html
const (
	DLTNULL     = C.DLT_NULL
	DLTRAW      = C.DLT_RAW
	DLTEN10MB   = C.DLT_EN10MB    // Ethernet (10Mb)
	DLTEN3MB    = C.DLT_EN3MB     // Experimental Ethernet (3Mb)
	DLTAX25     = C.DLT_AX25      // Amateur Radio AX.25
	DLTPRONET   = C.DLT_PRONET    // Proteon ProNET Token Ring
	DLTCHAOS    = C.DLT_CHAOS     // Chaos
	DLTIEEE802  = C.DLT_IEEE802   // 802.5 Token Ring
	DLTARCNET   = C.DLT_ARCNET    // ARCNET, with BSD-style header
	DLTSLIP     = C.DLT_SLIP      // Serial Line IP
	DLTPPP      = C.DLT_PPP       // Point-to-point Protocol
	DLTFDDI     = C.DLT_FDDI      // FDDI
	DLTLINUXSSL = C.DLT_LINUX_SLL // Linux cooked
)

const errbufSize = C.PCAP_ERRBUF_SIZE

// LinkType constants according to pcap-linktype(7).
const (
	LinkTypeNull      = DLTNULL
	LinkTypeEthernet  = DLTEN10MB
	LinkTypeTokenRing = DLTIEEE802

	LinkTypeExpEthernet = DLTEN3MB /* 3Mb experimental Ethernet */
	LinkTypeAX25        = DLTAX25
	LinkTypePRONET      = DLTPRONET
	LinkTypeCHAOS       = DLTCHAOS
	LinkTypeARCNETBSD   = DLTARCNET /* BSD-style headers */
	LinkTypeSLIP        = DLTSLIP
	LinkTypePPP         = DLTPPP
	LinkTypeFDDI        = DLTFDDI

	LinkTypeARCNET         = 7
	LinkTypeATMRFC1483     = 100
	LinkTypeRAW            = 101
	LinkTypePPPHDLC        = 50
	LinkTypePPPETHER       = 51
	LinkTypeCHDLC          = 104
	LinkTypeIEEE80211      = 105
	LinkTypeFRELAY         = 107
	LinkTypeLOOP           = 108
	LinkTypeLINUXSLL       = 113
	LinkTypeLTALK          = 104
	LinkTypePFLOG          = 117
	LinkTypePRISMHeader    = 119
	LinkTypeIPOverFC       = 122
	LinkTypeSUNATM         = 123
	LinkTypeIEEE80211Radio = 127
	LinkTypeARCNETLinux    = 129
	LinkTypeLinuxIRDA      = 144
	LinkTypeLinuxLAPD      = 177
)

type addrHdr interface {
	SrcAddr() string
	DestAddr() string
	Len() int
}

type addrStringer interface {
	String(addr addrHdr) string
}

func decodemac(pkt []byte) uint64 {
	mac := uint64(0)
	for i := 0; i < 6; i++ {
		mac = (mac << 8) + uint64(pkt[i])
	}
	return mac
}

// ARPHdr is a ARP packet header.
type ARPHdr struct {
	Addrtype          uint16
	Protocol          uint16
	HwAddressSize     uint8
	ProtAddressSize   uint8
	Operation         uint16
	SourceHwAddress   []byte
	SourceProtAddress []byte
	DestHwAddress     []byte
	DestProtAddress   []byte
}

func (arp *ARPHdr) String() (s string) {
	switch arp.Operation {
	case 1:
		s = "ARP request"
	case 2:
		s = "ARP Reply"
	}
	if arp.Addrtype == LinkTypeEthernet && arp.Protocol == TypeIP {
		s = fmt.Sprintf("%012x (%s) > %012x (%s)",
			decodemac(arp.SourceHwAddress), arp.SourceProtAddress,
			decodemac(arp.DestHwAddress), arp.DestProtAddress)
	} else {
		s = fmt.Sprintf("addrtype = %d protocol = %d", arp.Addrtype, arp.Protocol)
	}
	return
}

// Len returns the (calculated) ARP payload length.
func (arp *ARPHdr) Len() int {
	return int(8 + 2*arp.HwAddressSize + 2*arp.ProtAddressSize)
}

// IPHdr is the header of an IP packet.
type IPHdr struct {
	Version    uint8
	Ihl        uint8
	Tos        uint8
	Length     uint16
	ID         uint16
	Flags      uint8
	FragOffset uint16
	Ttl        uint8
	Protocol   uint8
	Checksum   uint16
	SrcIP      []byte
	DestIP     []byte
}

// SrcAddr returns the string version of the source IP.
func (ip *IPHdr) SrcAddr() string { return net.IP(ip.SrcIP).String() }

// DestAddr returns the string version of the destination IP.
func (ip *IPHdr) DestAddr() string { return net.IP(ip.DestIP).String() }

// Len returns the ip.Length.
func (ip *IPHdr) Len() int { return int(ip.Length) }

// TCPHdr is the header of a TCP packet.
type TCPHdr struct {
	SrcPort    uint16
	DestPort   uint16
	Seq        uint32
	Ack        uint32
	DataOffset uint8
	Flags      uint16
	Window     uint16
	Checksum   uint16
	Urgent     uint16
	Data       []byte
}

// TCP flags.
const (
	TCPFIN = 1 << iota
	TCPSYN
	TCPRST
	TCPPSH
	TCPACK
	TCPURG
	TCPECE
	TCPCWR
	TCPNS
)

// String TODO
func (tcp *TCPHdr) String(hdr addrHdr) string {
	return fmt.Sprintf("TCP %s:%d > %s:%d %s SEQ=%d ACK=%d LEN=%d",
		hdr.SrcAddr(), int(tcp.SrcPort), hdr.DestAddr(), int(tcp.DestPort),
		tcp.FlagsString(), int64(tcp.Seq), int64(tcp.Ack), hdr.Len())
}

// FlagsString TODO
func (tcp *TCPHdr) FlagsString() string {
	var sflags []string
	if 0 != (tcp.Flags & TCPSYN) {
		sflags = append(sflags, "syn")
	}
	if 0 != (tcp.Flags & TCPFIN) {
		sflags = append(sflags, "fin")
	}
	if 0 != (tcp.Flags & TCPACK) {
		sflags = append(sflags, "ack")
	}
	if 0 != (tcp.Flags & TCPPSH) {
		sflags = append(sflags, "psh")
	}
	if 0 != (tcp.Flags & TCPRST) {
		sflags = append(sflags, "rst")
	}
	if 0 != (tcp.Flags & TCPURG) {
		sflags = append(sflags, "urg")
	}
	if 0 != (tcp.Flags & TCPNS) {
		sflags = append(sflags, "ns")
	}
	if 0 != (tcp.Flags & TCPCWR) {
		sflags = append(sflags, "cwr")
	}
	if 0 != (tcp.Flags & TCPECE) {
		sflags = append(sflags, "ece")
	}
	return fmt.Sprintf("[%s]", strings.Join(sflags, " "))
}

// UDPHdr is the header of a UDP packet.
type UDPHdr struct {
	SrcPort  uint16
	DestPort uint16
	Length   uint16
	Checksum uint16
}

func (udp *UDPHdr) String(hdr addrHdr) string {
	return fmt.Sprintf("UDP %s:%d > %s:%d LEN=%d CHKSUM=%d",
		hdr.SrcAddr(), int(udp.SrcPort), hdr.DestAddr(), int(udp.DestPort),
		int(udp.Length), int(udp.Checksum))
}

// ICMPHdr is the header of an ICMP packet.
type ICMPHdr struct {
	Type     uint8
	Code     uint8
	Checksum uint16
	ID       uint16
	Seq      uint16
	Data     []byte
}

// String TODO
func (icmp *ICMPHdr) String(hdr addrHdr) string {
	return fmt.Sprintf("ICMP %s > %s Type = %d Code = %d ",
		hdr.SrcAddr(), hdr.DestAddr(), icmp.Type, icmp.Code)
}

// TypeString TODO
func (icmp *ICMPHdr) TypeString() (result string) {
	switch icmp.Type {
	case 0:
		result = fmt.Sprintf("Echo reply seq=%d", icmp.Seq)
	case 3:
		switch icmp.Code {
		case 0:
			result = "Network unreachable"
		case 1:
			result = "Host unreachable"
		case 2:
			result = "Protocol unreachable"
		case 3:
			result = "Port unreachable"
		default:
			result = "Destination unreachable"
		}
	case 8:
		result = fmt.Sprintf("Echo request seq=%d", icmp.Seq)
	case 30:
		result = "Traceroute"
	}
	return
}

// IP6Hdr is the header of an IPv6 packet.
type IP6Hdr struct {
	// http://www.networksorcery.com/enp/protocol/ipv6.htm
	Version      uint8  // 4 bits
	TrafficClass uint8  // 8 bits
	FlowLabel    uint32 // 20 bits
	Length       uint16 // 16 bits
	NextHeader   uint8  // 8 bits, same as Protocol in IPHdr
	HopLimit     uint8  // 8 bits
	SrcIP        []byte // 16 bytes
	DestIP       []byte // 16 bytes
}

// SrcAddr returns the string version of the source IP.
func (ip6 *IP6Hdr) SrcAddr() string { return net.IP(ip6.SrcIP).String() }

// DestAddr returns the string version of the destination IP.
func (ip6 *IP6Hdr) DestAddr() string { return net.IP(ip6.DestIP).String() }

// Len returns the ip6.Length.
func (ip6 *IP6Hdr) Len() int { return int(ip6.Length) }

// ICMPv6Hdr is for ICMPv6
type ICMPv6Hdr struct {
	Type     uint8
	Code     uint8
	Checksum uint16
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"net/url"
	"strconv"
)

type TaskListController struct{}

// # GET /tasks
// ============
//
// Fetchs all **Task** objects.
//
func (t *TaskListController) Get(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	return http.StatusOK, GetTaskList(db)
}

type TaskController struct{}

// GET /task
// =========
//
// Fetch a **Task** by id.
//
// ### Params:
//
// - **id**: The object id hex string
//
func (t *TaskController) Get(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task, err := GetTaskById(db, params.Get("id"))
	if err != nil {
		return http.StatusNotFound, nil
	}

	return http.StatusOK, task
}

// POST /task
// ==========

// Create a new **Task** object.

// ### Params:

// - **ok**: The status as boolean string
// - **desc**: The description string

// #### Boolean string:

// ```
// 1, t, T, TRUE, true, True
// 0, f, F, FALSE, false, False
// ````
//
func (t *TaskController) Post(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task := new(Task)
	task.Completed = false
	task.Description = params.Get("desc")

	if err := task.Insert(db); err != nil {
		return http.StatusBadRequest, nil
	}

	return http.StatusOK, task
}

// PUT /task
// =========

// Update a given *Task* by id.

// ### Params:

// - **id**: The object id hex string
// - **ok**: The status as boolean string
// - **desc**: The description string

// #### Boolean string:

// ```
// 1, t, T, TRUE, true, True
// 0, f, F, FALSE, false, False
// ````
//
func (t *TaskController) Put(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task, err := GetTaskById(db, params.Get("id"))
	if err != nil {
		return http.StatusNotFound, nil
	}

	if ok := params.Get("ok"); ok != "" {
		task.Completed, _ = strconv.ParseBool(ok)
	}
	if desc := params.Get("desc"); desc != "" {
		task.Description = desc
	}
	// Do update
	if err := task.Update(db); err != nil {
		return http.StatusBadRequest, nil
	}

	return http.StatusOK, task
}

// DELETE /task
// ============

// Remove a **Task** by id.

// ### Params:

// - **id**: The object id hex string
//
func (t *TaskController) Delete(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task, err := GetTaskById(db, params.Get("id"))
	if err != nil {
		return http.StatusNotFound, nil
	}
	// Do remove
	if err := task.Remove(db); err != nil {
		log.Println(err)
		return http.StatusBadRequest, nil
	}

	return http.StatusAccepted, nil
}
<commit_msg>Remove log message<commit_after>package main

import (
	"net/http"
	"net/url"
	"strconv"
)

type TaskListController struct{}

// # GET /tasks
// ============
//
// Fetchs all **Task** objects.
//
func (t *TaskListController) Get(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	return http.StatusOK, GetTaskList(db)
}

type TaskController struct{}

// GET /task
// =========
//
// Fetch a **Task** by id.
//
// ### Params:
//
// - **id**: The object id hex string
//
func (t *TaskController) Get(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task, err := GetTaskById(db, params.Get("id"))
	if err != nil {
		return http.StatusNotFound, nil
	}

	return http.StatusOK, task
}

// POST /task
// ==========

// Create a new **Task** object.

// ### Params:

// - **ok**: The status as boolean string
// - **desc**: The description string

// #### Boolean string:

// ```
// 1, t, T, TRUE, true, True
// 0, f, F, FALSE, false, False
// ````
//
func (t *TaskController) Post(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task := new(Task)
	task.Completed = false
	task.Description = params.Get("desc")

	if err := task.Insert(db); err != nil {
		return http.StatusBadRequest, nil
	}

	return http.StatusOK, task
}

// PUT /task
// =========

// Update a given *Task* by id.

// ### Params:

// - **id**: The object id hex string
// - **ok**: The status as boolean string
// - **desc**: The description string

// #### Boolean string:

// ```
// 1, t, T, TRUE, true, True
// 0, f, F, FALSE, false, False
// ````
//
func (t *TaskController) Put(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task, err := GetTaskById(db, params.Get("id"))
	if err != nil {
		return http.StatusNotFound, nil
	}

	if ok := params.Get("ok"); ok != "" {
		task.Completed, _ = strconv.ParseBool(ok)
	}
	if desc := params.Get("desc"); desc != "" {
		task.Description = desc
	}
	// Do update
	if err := task.Update(db); err != nil {
		return http.StatusBadRequest, nil
	}

	return http.StatusOK, task
}

// DELETE /task
// ============

// Remove a **Task** by id.

// ### Params:

// - **id**: The object id hex string
//
func (t *TaskController) Delete(params url.Values) (int, interface{}) {
	db, conn := GetDatabase()
	defer conn.Close()

	task, err := GetTaskById(db, params.Get("id"))
	if err != nil {
		return http.StatusNotFound, nil
	}
	// Do remove
	if err := task.Remove(db); err != nil {
		return http.StatusBadRequest, nil
	}

	return http.StatusAccepted, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Jonathan Pincas

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

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

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

package core

import (
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "ecosystem [command] [arguments]",
	Short: "EcoSystem command line tool",
	Long: `Use to initialise or launch the EcoSystem server or create new users or bundles.
	Use the bare command 'ecosystem' to create a new config.json or verify an existing one.`,
	RunE: createConfigIfNotExists,
}

func init() {

	RootCmd.PersistentFlags().StringP("pgpw", "p", "", "Postgres superuser password")
	RootCmd.PersistentFlags().StringP("configfile", "c", "config", "Name of config file (without extension)")
	viper.BindPFlags(RootCmd.PersistentFlags())

}

// initConfig reads in config file and ENV variables if set.
func createConfigIfNotExists(cmd *cobra.Command, args []string) error {

	viper.SetConfigName(viper.GetString("configfile"))

	if err := viper.ReadInConfig(); err == nil {
		LogFatal(LogEntry{"CORE.CONFIG", true, "Config file already exists:" + viper.ConfigFileUsed()})
	} else {
		if err := createDefaultConfigFile(viper.GetString("configfile")); err != nil {
			LogFatal(LogEntry{"CORE.CONFIG", false, "Error creating config file: " + err.Error()})
		} else {
			//Otherwise create one
			Log(LogEntry{"CORE.CONFIG", true, "Config file created"})
		}
	}

	return nil
}
<commit_msg>Added ping test<commit_after>// Copyright 2017 Jonathan Pincas

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

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

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

package core

import (
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "ecosystem [command] [arguments]",
	Short: "EcoSystem command line tool",
	Long: `Use to initialise or launch the EcoSystem server or create new users or bundles.
	Use the bare command 'ecosystem' to create a new config.json or verify an existing one.`,
	RunE: createConfigIfNotExists,
}

// serveCmd represents the serve command
var pingCmd = &cobra.Command{
	Use:   "ping",
	Short: "Ping test",
	Long:  `Pings the DB and returns OK if the connection is ready.`,
	RunE:  ping,
}

func init() {

	RootCmd.AddCommand(pingCmd)
	RootCmd.PersistentFlags().StringP("pgpw", "p", "", "Postgres superuser password")
	RootCmd.PersistentFlags().StringP("configfile", "c", "config", "Name of config file (without extension)")
	viper.BindPFlags(RootCmd.PersistentFlags())

}

// initConfig reads in config file and ENV variables if set.
func createConfigIfNotExists(cmd *cobra.Command, args []string) error {

	viper.SetConfigName(viper.GetString("configfile"))

	if err := viper.ReadInConfig(); err == nil {
		LogFatal(LogEntry{"CORE.CONFIG", true, "Config file already exists:" + viper.ConfigFileUsed()})
	} else {
		if err := createDefaultConfigFile(viper.GetString("configfile")); err != nil {
			LogFatal(LogEntry{"CORE.CONFIG", false, "Error creating config file: " + err.Error()})
		} else {
			//Otherwise create one
			Log(LogEntry{"CORE.CONFIG", true, "Config file created"})
		}
	}

	return nil
}

func ping(cmd *cobra.Command, args []string) error {

	readConfig()

	//Attempt to open a db connection
	db := SuperUserDBConfig.ReturnDBConnection("")
	defer db.Close()
	//IF we get this far, just exit with success
	Log(LogEntry{"PING", true, "Ping test passed"})
	os.Exit(0)

	return nil

}
<|endoftext|>


<commit_before>package core

type Bagger struct {
	Bag      *Bag
	Profile  *BagItProfile
	Sources  []string
	Defaults map[string]string
}

// sources is a list of source files/directories to put into the data dir.
// defaults is a map of default tag values.
func NewBagger(bag *Bag, profile *BagItProfile, sources []string, defaults map[string]string) *Bagger {
	return &Bagger{
		Bag:     bag,
		Profile: profile,
		Sources: sources,
	}
}

/*

- Validate Profile
- Validate that all defaults are present
- Include progress callback or io.Writer for writing progress messages
- Add all payload files to bag in working dir
- Create manifests

- Tar bag if serialization is required and serlialization format is tar?
- Provide a way of specifying tag files and where to put them?
- Parse existing tag files, and skip defaults if they're already defined in existing tag files?

*/
<commit_msg>Working on bagger.<commit_after>package core

import (
	"fmt"
	"github.com/APTrust/bagit/util/fileutil"
	"os"
)

type Bagger struct {
	Bag            *Bag
	Profile        *BagItProfile
	PayloadDir     string
	TagValues      map[string]string
	ForceOverwrite bool
	errors         []string
}

// NewBagger creates a new Bagger.
//
// Param bagPath is the path to the location where the bag should
// be built. This should be a directory name. The bagger will
// create the directory, if it does not already exist. If the directory
// exists, the bagger will not overwrite it, unless the forceOverwrite
// (see below) is true.
//
// Param payloadDir is the name of the directory containing the files to be
// bagged.
//
// Param profile is the BagIt profile that describes the requirements
// for the bag.
//
// Param tagValues is a map of tag values that will be used
// to create tag files.
//
// If param forceOverwrite is true, the bagger will delete and then
// overwrite the directory at bagPath.
func NewBagger(bagPath, payloadDir string, profile *BagItProfile, tagValues map[string]string, forceOverwrite bool) *Bagger {
	return &Bagger{
		Bag:            NewBag(bagPath),
		Profile:        profile,
		PayloadDir:     payloadDir,
		TagValues:      tagValues,
		ForceOverwrite: forceOverwrite,
		errors:         make([]string, 0),
	}
}

/*

- Validate Profile
- Validate that all defaults are present
- Include progress callback or io.Writer for writing progress messages
- Add all payload files to bag in working dir
- Build tag files
- Create manifests

- Tar bag if serialization is required and serlialization format is tar?
- Provide a way of specifying tag files and where to put them?
- Parse existing tag files, and skip defaults if they're already defined in existing tag files?

*/

func (bagger *Bagger) BuildBag() bool {
	ok := bagger.makeDir()
	if !ok {
		return false
	}
	return true
}

// makeDir creates the directory in which we'll assemble the bag,
// performing some safety checks along the way. Returns true on
// success, false otherwise.
func (bagger *Bagger) makeDir() bool {
	if fileutil.FileExists(bagger.Bag.Path) {
		if bagger.ForceOverwrite {
			if fileutil.LooksSafeToDelete(bagger.Bag.Path, 12, 3) {
				err := os.RemoveAll(bagger.Bag.Path)
				if err != nil {
					bagger.addError("Error removing existing %s", bagger.Bag.Path, err.Error())
					return false
				}
			} else {
				bagger.addError("%s already exists and does not look safe to delete. "+
					"Choose a different path for this bag.",
					bagger.Bag.Path)
				return false
			}
		} else {
			bagger.addError("%s already exists. Use forceOverwrite=true to replace it.",
				bagger.Bag.Path)
			return false
		}
	}
	err := os.MkdirAll(bagger.Bag.Path, 0755)
	if err != nil {
		bagger.addError(err.Error())
		return false
	}
	return true
}

// addError adds a message to the list of bagging errors.
func (bagger *Bagger) addError(format string, a ...interface{}) {
	bagger.errors = append(bagger.errors, fmt.Sprintf(format, a...))
}
<|endoftext|>


<commit_before>package core

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
)

const (
	DefaultAPIBaseURL = "http://demo6970933.mockable.io"
	DefaultAPIVersion = 1
)

type APIClient struct {
	Server    string
	UUID      string
	AuthToken string
	Client    *http.Client
}

func NewAPIClient(server string, uuid string, authToken string) *APIClient {
	return &APIClient{
		Client:    &http.Client{},
		UUID:      uuid,
		AuthToken: authToken,
		Server:    server,
	}
}

type ConfigResponse struct {
	Signed string `json:"signed"`
	Raw    string `json:"raw"`
}

func (api *APIClient) GetFormattedURL(prefix ...string) string {
	return fmt.Sprintf("%s/%s/%s", api.Server, strconv.Itoa(DefaultAPIVersion),
		strings.Join(prefix, "/"))
}

func (api *APIClient) GetConfig() (*ConfigResponse, error) {
	request, err := http.NewRequest("GET",
		api.GetFormattedURL(api.UUID), nil)

	if err != nil {
		return nil, err
	}

	if api.AuthToken != "" {
		request.Header.Add("Auth-Token", api.AuthToken)
	}

	response, err := api.Client.Do(request)
	if err != nil {
		return nil, err
	}

	if response.StatusCode != 200 {
		return nil, fmt.Errorf("Invalid server response: %s", response.Status)
	}

	body, err := ioutil.ReadAll(response.Body)

	if err != nil {
		return nil, err
	}

	config := new(ConfigResponse)
	err = json.Unmarshal(body, &config)

	if err != nil {
		return nil, err
	}

	return config, nil
}

type Client struct {
	Hostname    string
	ReportsPath string
	APIClient   *APIClient
}

func NewClient(server string, uuid string, authToken string) (*Client, error) {
	reportsPath, err := GetDefaultReportsDirectory()
	if err != nil {
		return nil, err
	}

	return &Client{
		ReportsPath: reportsPath,
		APIClient:   NewAPIClient(server, uuid, authToken),
	}, nil
}

func (client *Client) PromptPGPConfirmation(config *Config) bool {
	var answer string

	for _, key := range config.Signature.Keys {
		fmt.Printf("Configuration file Signed-off by PGP Key: %s\n", key.PublicKey.KeyIdShortString())
		for _, identity := range key.Entity.Identities {
			fmt.Printf(" - %s\n", identity.UserId.Id)
		}
	}

	fmt.Printf("Proceed (y/n)? ")
	fmt.Scanf("%s", &answer)

	return answer == "y"
}

func (client *Client) Run(pgp bool, upload bool) error {
	reportPath, err := GetTempReportDirectory()

	if err != nil {
		return err
	}

	apiConfig, err := client.APIClient.GetConfig()
	if err != nil {
		return fmt.Errorf("Error getting configuration from server: %s", err)
	}

	raw, err := base64.StdEncoding.DecodeString(apiConfig.Raw)
	if err != nil {
		return err
	}

	signed, err := base64.StdEncoding.DecodeString(apiConfig.Signed)

	if err != nil {
		return err
	}

	config, err := NewConfig(string(raw), string(signed))

	if pgp {
		err := config.CheckPGPSignature()
		if err != nil {
			return err
		}
		answer := client.PromptPGPConfirmation(config)
		if answer != true {
			return fmt.Errorf("PGP key has not been accepted")
		}
	}

	wg := new(sync.WaitGroup)
	log.Printf("Starting a new report on: %s", reportPath)

	for _, command := range config.Commands {
		wg.Add(1)
		go client.RunCommand(reportPath, command, wg)
	}

	for _, file := range config.Files {
		finfo, err := os.Stat(file.Path)
		if err != nil {
			log.Printf("Cannot stat file:%s", file.Path)
		} else {
			log.Printf("Archiving file:%s", file.Path)
			if finfo.IsDir() {
				CopyDir(file.Path, path.Join(reportPath, file.Path))
			} else {
				CopyFile(file.Path, reportPath)
			}
		}
	}

	wg.Wait()

	return nil
}

func (m *Client) RunCommand(reportPath string, command Command, wg *sync.WaitGroup) {
	//TODO: Refactor this to handle properly all the error cases
	log.Printf("Running %s", command.Executable)
	ran, _ := exec.Command("/bin/bash", "-c", command.Executable).Output()
	outfile, _ := os.Create(command.GetFileName(reportPath))
	outfile.WriteString(string(ran))
	defer outfile.Close()
	defer wg.Done()
}

func (c *Command) GetFileName(Base string) string {
	return path.Join(Base, MangleCommand(c.Executable))
}
<commit_msg>[core] moved api client to core/api.go<commit_after>package core

import (
	"encoding/base64"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path"
	"sync"
)

type Client struct {
	Hostname    string
	ReportsPath string
	APIClient   *APIClient
}

func NewClient(server string, uuid string, authToken string) (*Client, error) {
	reportsPath, err := GetDefaultReportsDirectory()
	if err != nil {
		return nil, err
	}

	return &Client{
		ReportsPath: reportsPath,
		APIClient:   NewAPIClient(server, uuid, authToken),
	}, nil
}

func (client *Client) Run(pgp bool, upload bool) error {
	reportPath, err := GetTempReportDirectory()

	if err != nil {
		return err
	}

	apiConfig, err := client.APIClient.GetConfig()
	if err != nil {
		return fmt.Errorf("Error getting configuration from server: %s", err)
	}

	raw, err := base64.StdEncoding.DecodeString(apiConfig.Raw)
	if err != nil {
		return err
	}

	signed, err := base64.StdEncoding.DecodeString(apiConfig.Signed)

	if err != nil {
		return err
	}

	config, err := NewConfig(string(raw), string(signed))

	if pgp {
		err := config.CheckPGPSignature()
		if err != nil {
			return err
		}
		answer := PromptPGPConfirmation(config)
		if answer != true {
			return fmt.Errorf("PGP key has not been accepted")
		}
	}

	wg := new(sync.WaitGroup)
	log.Printf("Starting a new report on: %s", reportPath)

	for _, command := range config.Commands {
		wg.Add(1)
		go client.RunCommand(reportPath, command, wg)
	}

	for _, file := range config.Files {
		finfo, err := os.Stat(file.Path)
		if err != nil {
			log.Printf("Cannot stat file:%s", file.Path)
		} else {
			log.Printf("Archiving file:%s", file.Path)
			if finfo.IsDir() {
				CopyDir(file.Path, path.Join(reportPath, file.Path))
			} else {
				CopyFile(file.Path, reportPath)
			}
		}
	}

	wg.Wait()
	return nil
}

func (m *Client) RunCommand(reportPath string, command Command, wg *sync.WaitGroup) {
	//TODO: Refactor this to handle properly all the error cases
	log.Printf("Running %s", command.Executable)
	ran, _ := exec.Command("/bin/bash", "-c", command.Executable).Output()
	outfile, _ := os.Create(command.GetFileName(reportPath))
	outfile.WriteString(string(ran))
	defer outfile.Close()
	defer wg.Done()
}

func (c *Command) GetFileName(Base string) string {
	return path.Join(Base, MangleCommand(c.Executable))
}
<|endoftext|>