<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package storage

import (
	"fmt"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/api/core/v1"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/test/e2e/framework"
)

type testBody func(c clientset.Interface, f *framework.Framework, clientPod *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume)
type disruptiveTest struct {
	testItStmt string
	runTest    testBody
}
type kubeletOpt string

const (
	MinNodes                    = 2
	NodeStateTimeout            = 1 * time.Minute
	kStart           kubeletOpt = "start"
	kStop            kubeletOpt = "stop"
	kRestart         kubeletOpt = "restart"
)

var _ = SIGDescribe("PersistentVolumes[Disruptive][Flaky]", func() {

	f := framework.NewDefaultFramework("disruptive-pv")
	var (
		c                         clientset.Interface
		ns                        string
		nfsServerPod              *v1.Pod
		nfsPVconfig               framework.PersistentVolumeConfig
		pvcConfig                 framework.PersistentVolumeClaimConfig
		nfsServerIP, clientNodeIP string
		clientNode                *v1.Node
		volLabel                  labels.Set
		selector                  *metav1.LabelSelector
	)

	BeforeEach(func() {
		// To protect the NFS volume pod from the kubelet restart, we isolate it on its own node.
		framework.SkipUnlessNodeCountIsAtLeast(MinNodes)
		framework.SkipIfProviderIs("local")

		c = f.ClientSet
		ns = f.Namespace.Name
		volLabel = labels.Set{framework.VolumeSelectorKey: ns}
		selector = metav1.SetAsLabelSelector(volLabel)
		// Start the NFS server pod.
		_, nfsServerPod, nfsServerIP = framework.NewNFSServer(c, ns, []string{"-G", "777", "/exports"})
		nfsPVconfig = framework.PersistentVolumeConfig{
			NamePrefix: "nfs-",
			Labels:     volLabel,
			PVSource: v1.PersistentVolumeSource{
				NFS: &v1.NFSVolumeSource{
					Server:   nfsServerIP,
					Path:     "/exports",
					ReadOnly: false,
				},
			},
		}
		pvcConfig = framework.PersistentVolumeClaimConfig{
			Annotations: map[string]string{
				v1.BetaStorageClassAnnotation: "",
			},
			Selector: selector,
		}
		// Get the first ready node IP that is not hosting the NFS pod.
		if clientNodeIP == "" {
			framework.Logf("Designating test node")
			nodes := framework.GetReadySchedulableNodesOrDie(c)
			for _, node := range nodes.Items {
				if node.Name != nfsServerPod.Spec.NodeName {
					clientNode = &node
					clientNodeIP = framework.GetNodeExternalIP(clientNode)
					break
				}
			}
			Expect(clientNodeIP).NotTo(BeEmpty())
		}
	})

	AfterEach(func() {
		framework.DeletePodWithWait(f, c, nfsServerPod)
	})

	Context("when kubelet restarts", func() {
		var (
			clientPod *v1.Pod
			pv        *v1.PersistentVolume
			pvc       *v1.PersistentVolumeClaim
		)

		BeforeEach(func() {
			framework.Logf("Initializing test spec")
			clientPod, pv, pvc = initTestCase(f, c, nfsPVconfig, pvcConfig, ns, clientNode.Name)
		})

		AfterEach(func() {
			framework.Logf("Tearing down test spec")
			tearDownTestCase(c, f, ns, clientPod, pvc, pv)
			pv, pvc, clientPod = nil, nil, nil
		})

		// Test table housing the It() title string and test spec.  runTest is type testBody, defined at
		// the start of this file.  To add tests, define a function mirroring the testBody signature and assign
		// to runTest.
		disruptiveTestTable := []disruptiveTest{
			{
				testItStmt: "Should test that a file written to the mount before kubelet restart is readable after restart.",
				runTest:    testKubeletRestartsAndRestoresMount,
			},
			{
				testItStmt: "Should test that a volume mounted to a pod that is deleted while the kubelet is down unmounts when the kubelet returns.",
				runTest:    testVolumeUnmountsFromDeletedPod,
			},
		}

		// Test loop executes each disruptiveTest iteratively.
		for _, test := range disruptiveTestTable {
			func(t disruptiveTest) {
				It(t.testItStmt, func() {
					By("Executing Spec")
					t.runTest(c, f, clientPod, pvc, pv)
				})
			}(test)
		}
	})
})

// testKubeletRestartsAndRestoresMount tests that a volume mounted to a pod remains mounted after a kubelet restarts
func testKubeletRestartsAndRestoresMount(c clientset.Interface, f *framework.Framework, clientPod *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume) {
	By("Writing to the volume.")
	file := "/mnt/_SUCCESS"
	out, err := podExec(clientPod, fmt.Sprintf("touch %s", file))
	framework.Logf(out)
	Expect(err).NotTo(HaveOccurred())

	By("Restarting kubelet")
	kubeletCommand(kRestart, c, clientPod)

	By("Testing that written file is accessible.")
	out, err = podExec(clientPod, fmt.Sprintf("cat %s", file))
	framework.Logf(out)
	Expect(err).NotTo(HaveOccurred())
	framework.Logf("Volume mount detected on pod %s and written file %s is readable post-restart.", clientPod.Name, file)
}

// testVolumeUnmountsFromDeletedPod tests that a volume unmounts if the client pod was deleted while the kubelet was down.
func testVolumeUnmountsFromDeletedPod(c clientset.Interface, f *framework.Framework, clientPod *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume) {
	nodeIP, err := framework.GetHostExternalAddress(c, clientPod)
	Expect(err).NotTo(HaveOccurred())
	nodeIP = nodeIP + ":22"

	By("Expecting the volume mount to be found.")
	result, err := framework.SSH(fmt.Sprintf("mount | grep %s", clientPod.UID), nodeIP, framework.TestContext.Provider)
	framework.LogSSHResult(result)
	Expect(err).NotTo(HaveOccurred(), "Encountered SSH error.")
	Expect(result.Code).To(BeZero(), fmt.Sprintf("Expected grep exit code of 0, got %d", result.Code))

	By("Stopping the kubelet.")
	kubeletCommand(kStop, c, clientPod)
	defer func() {
		if err != nil {
			kubeletCommand(kStart, c, clientPod)
		}
	}()
	By(fmt.Sprintf("Deleting Pod %q", clientPod.Name))
	err = c.Core().Pods(clientPod.Namespace).Delete(clientPod.Name, &metav1.DeleteOptions{})
	Expect(err).NotTo(HaveOccurred())
	By("Starting the kubelet and waiting for pod to delete.")
	kubeletCommand(kStart, c, clientPod)
	err = f.WaitForPodTerminated(clientPod.Name, "")
	if !apierrs.IsNotFound(err) && err != nil {
		Expect(err).NotTo(HaveOccurred(), "Expected pod to terminate.")
	}

	By("Expecting the volume mount not to be found.")
	result, err = framework.SSH(fmt.Sprintf("mount | grep %s", clientPod.UID), nodeIP, framework.TestContext.Provider)
	framework.LogSSHResult(result)
	Expect(err).NotTo(HaveOccurred(), "Encountered SSH error.")
	Expect(result.Stdout).To(BeEmpty(), "Expected grep stdout to be empty (i.e. no mount found).")
	framework.Logf("Volume unmounted on node %s", clientPod.Spec.NodeName)
}

// initTestCase initializes spec resources (pv, pvc, and pod) and returns pointers to be consumed
// by the test.
func initTestCase(f *framework.Framework, c clientset.Interface, pvConfig framework.PersistentVolumeConfig, pvcConfig framework.PersistentVolumeClaimConfig, ns, nodeName string) (*v1.Pod, *v1.PersistentVolume, *v1.PersistentVolumeClaim) {
	pv, pvc, err := framework.CreatePVPVC(c, pvConfig, pvcConfig, ns, false)
	defer func() {
		if err != nil {
			framework.DeletePersistentVolumeClaim(c, pvc.Name, ns)
			framework.DeletePersistentVolume(c, pv.Name)
		}
	}()
	Expect(err).NotTo(HaveOccurred())
	pod := framework.MakePod(ns, []*v1.PersistentVolumeClaim{pvc}, true, "")
	pod.Spec.NodeName = nodeName
	framework.Logf("Creating NFS client pod.")
	pod, err = c.CoreV1().Pods(ns).Create(pod)
	framework.Logf("NFS client Pod %q created on Node %q", pod.Name, nodeName)
	Expect(err).NotTo(HaveOccurred())
	defer func() {
		if err != nil {
			framework.DeletePodWithWait(f, c, pod)
		}
	}()
	err = framework.WaitForPodRunningInNamespace(c, pod)
	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("Pod %q timed out waiting for phase: Running", pod.Name))
	// Return created api objects
	pod, err = c.CoreV1().Pods(ns).Get(pod.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	pvc, err = c.CoreV1().PersistentVolumeClaims(ns).Get(pvc.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	pv, err = c.CoreV1().PersistentVolumes().Get(pv.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	return pod, pv, pvc
}

// tearDownTestCase destroy resources created by initTestCase.
func tearDownTestCase(c clientset.Interface, f *framework.Framework, ns string, client *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume) {
	// Ignore deletion errors.  Failing on them will interrupt test cleanup.
	framework.DeletePodWithWait(f, c, client)
	framework.DeletePersistentVolumeClaim(c, pvc.Name, ns)
	framework.DeletePersistentVolume(c, pv.Name)
}

// kubeletCommand performs `start`, `restart`, or `stop` on the kubelet running on the node of the target pod and waits
// for the desired statues..
// - First issues the command via `systemctl`
// - If `systemctl` returns stderr "command not found, issues the command via `service`
// - If `service` also returns stderr "command not found", the test is aborted.
// Allowed kubeletOps are `kStart`, `kStop`, and `kRestart`
func kubeletCommand(kOp kubeletOpt, c clientset.Interface, pod *v1.Pod) {
	nodeIP, err := framework.GetHostExternalAddress(c, pod)
	Expect(err).NotTo(HaveOccurred())
	nodeIP = nodeIP + ":22"
	systemctlCmd := fmt.Sprintf("sudo systemctl %s kubelet", string(kOp))
	framework.Logf("Attempting `%s`", systemctlCmd)
	sshResult, err := framework.SSH(systemctlCmd, nodeIP, framework.TestContext.Provider)
	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("SSH to Node %q errored.", pod.Spec.NodeName))
	framework.LogSSHResult(sshResult)
	if strings.Contains(sshResult.Stderr, "command not found") {
		serviceCmd := fmt.Sprintf("sudo service kubelet %s", string(kOp))
		framework.Logf("Attempting `%s`", serviceCmd)
		sshResult, err = framework.SSH(serviceCmd, nodeIP, framework.TestContext.Provider)
		Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("SSH to Node %q errored.", pod.Spec.NodeName))
		framework.LogSSHResult(sshResult)
	}
	Expect(sshResult.Code).To(BeZero(), "Failed to [%s] kubelet:\n%#v", string(kOp), sshResult)
	// On restart, waiting for node NotReady prevents a race condition where the node takes a few moments to leave the
	// Ready state which in turn short circuits WaitForNodeToBeReady()
	if kOp == kStop || kOp == kRestart {
		if ok := framework.WaitForNodeToBeNotReady(c, pod.Spec.NodeName, NodeStateTimeout); !ok {
			framework.Failf("Node %s failed to enter NotReady state", pod.Spec.NodeName)
		}
	}
	if kOp == kStart || kOp == kRestart {
		if ok := framework.WaitForNodeToBeReady(c, pod.Spec.NodeName, NodeStateTimeout); !ok {
			framework.Failf("Node %s failed to enter Ready state", pod.Spec.NodeName)
		}
	}
}

// podExec wraps RunKubectl to execute a bash cmd in target pod
func podExec(pod *v1.Pod, bashExec string) (string, error) {
	return framework.RunKubectl("exec", fmt.Sprintf("--namespace=%s", pod.Namespace), pod.Name, "--", "/bin/sh", "-c", bashExec)
}
<commit_msg>Fixing kubelet restart<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package storage

import (
	"fmt"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/api/core/v1"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/test/e2e/framework"
)

type testBody func(c clientset.Interface, f *framework.Framework, clientPod *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume)
type disruptiveTest struct {
	testItStmt string
	runTest    testBody
}
type kubeletOpt string

const (
	MinNodes                    = 2
	NodeStateTimeout            = 1 * time.Minute
	kStart           kubeletOpt = "start"
	kStop            kubeletOpt = "stop"
	kRestart         kubeletOpt = "restart"
)

var _ = SIGDescribe("PersistentVolumes[Disruptive][Flaky]", func() {

	f := framework.NewDefaultFramework("disruptive-pv")
	var (
		c                         clientset.Interface
		ns                        string
		nfsServerPod              *v1.Pod
		nfsPVconfig               framework.PersistentVolumeConfig
		pvcConfig                 framework.PersistentVolumeClaimConfig
		nfsServerIP, clientNodeIP string
		clientNode                *v1.Node
		volLabel                  labels.Set
		selector                  *metav1.LabelSelector
	)

	BeforeEach(func() {
		// To protect the NFS volume pod from the kubelet restart, we isolate it on its own node.
		framework.SkipUnlessNodeCountIsAtLeast(MinNodes)
		framework.SkipIfProviderIs("local")

		c = f.ClientSet
		ns = f.Namespace.Name
		volLabel = labels.Set{framework.VolumeSelectorKey: ns}
		selector = metav1.SetAsLabelSelector(volLabel)
		// Start the NFS server pod.
		_, nfsServerPod, nfsServerIP = framework.NewNFSServer(c, ns, []string{"-G", "777", "/exports"})
		nfsPVconfig = framework.PersistentVolumeConfig{
			NamePrefix: "nfs-",
			Labels:     volLabel,
			PVSource: v1.PersistentVolumeSource{
				NFS: &v1.NFSVolumeSource{
					Server:   nfsServerIP,
					Path:     "/exports",
					ReadOnly: false,
				},
			},
		}
		pvcConfig = framework.PersistentVolumeClaimConfig{
			Annotations: map[string]string{
				v1.BetaStorageClassAnnotation: "",
			},
			Selector: selector,
		}
		// Get the first ready node IP that is not hosting the NFS pod.
		if clientNodeIP == "" {
			framework.Logf("Designating test node")
			nodes := framework.GetReadySchedulableNodesOrDie(c)
			for _, node := range nodes.Items {
				if node.Name != nfsServerPod.Spec.NodeName {
					clientNode = &node
					clientNodeIP = framework.GetNodeExternalIP(clientNode)
					break
				}
			}
			Expect(clientNodeIP).NotTo(BeEmpty())
		}
	})

	AfterEach(func() {
		framework.DeletePodWithWait(f, c, nfsServerPod)
	})

	Context("when kubelet restarts", func() {
		var (
			clientPod *v1.Pod
			pv        *v1.PersistentVolume
			pvc       *v1.PersistentVolumeClaim
		)

		BeforeEach(func() {
			framework.Logf("Initializing test spec")
			clientPod, pv, pvc = initTestCase(f, c, nfsPVconfig, pvcConfig, ns, clientNode.Name)
		})

		AfterEach(func() {
			framework.Logf("Tearing down test spec")
			tearDownTestCase(c, f, ns, clientPod, pvc, pv)
			pv, pvc, clientPod = nil, nil, nil
		})

		// Test table housing the It() title string and test spec.  runTest is type testBody, defined at
		// the start of this file.  To add tests, define a function mirroring the testBody signature and assign
		// to runTest.
		disruptiveTestTable := []disruptiveTest{
			{
				testItStmt: "Should test that a file written to the mount before kubelet restart is readable after restart.",
				runTest:    testKubeletRestartsAndRestoresMount,
			},
			{
				testItStmt: "Should test that a volume mounted to a pod that is deleted while the kubelet is down unmounts when the kubelet returns.",
				runTest:    testVolumeUnmountsFromDeletedPod,
			},
		}

		// Test loop executes each disruptiveTest iteratively.
		for _, test := range disruptiveTestTable {
			func(t disruptiveTest) {
				It(t.testItStmt, func() {
					By("Executing Spec")
					t.runTest(c, f, clientPod, pvc, pv)
				})
			}(test)
		}
	})
})

// testKubeletRestartsAndRestoresMount tests that a volume mounted to a pod remains mounted after a kubelet restarts
func testKubeletRestartsAndRestoresMount(c clientset.Interface, f *framework.Framework, clientPod *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume) {
	By("Writing to the volume.")
	file := "/mnt/_SUCCESS"
	out, err := podExec(clientPod, fmt.Sprintf("touch %s", file))
	framework.Logf(out)
	Expect(err).NotTo(HaveOccurred())

	By("Restarting kubelet")
	kubeletCommand(kRestart, c, clientPod)

	By("Testing that written file is accessible.")
	out, err = podExec(clientPod, fmt.Sprintf("cat %s", file))
	framework.Logf(out)
	Expect(err).NotTo(HaveOccurred())
	framework.Logf("Volume mount detected on pod %s and written file %s is readable post-restart.", clientPod.Name, file)
}

// testVolumeUnmountsFromDeletedPod tests that a volume unmounts if the client pod was deleted while the kubelet was down.
func testVolumeUnmountsFromDeletedPod(c clientset.Interface, f *framework.Framework, clientPod *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume) {
	nodeIP, err := framework.GetHostExternalAddress(c, clientPod)
	Expect(err).NotTo(HaveOccurred())
	nodeIP = nodeIP + ":22"

	By("Expecting the volume mount to be found.")
	result, err := framework.SSH(fmt.Sprintf("mount | grep %s", clientPod.UID), nodeIP, framework.TestContext.Provider)
	framework.LogSSHResult(result)
	Expect(err).NotTo(HaveOccurred(), "Encountered SSH error.")
	Expect(result.Code).To(BeZero(), fmt.Sprintf("Expected grep exit code of 0, got %d", result.Code))

	By("Stopping the kubelet.")
	kubeletCommand(kStop, c, clientPod)
	defer func() {
		if err != nil {
			kubeletCommand(kStart, c, clientPod)
		}
	}()
	By(fmt.Sprintf("Deleting Pod %q", clientPod.Name))
	err = c.Core().Pods(clientPod.Namespace).Delete(clientPod.Name, &metav1.DeleteOptions{})
	Expect(err).NotTo(HaveOccurred())
	By("Starting the kubelet and waiting for pod to delete.")
	kubeletCommand(kStart, c, clientPod)
	err = f.WaitForPodTerminated(clientPod.Name, "")
	if !apierrs.IsNotFound(err) && err != nil {
		Expect(err).NotTo(HaveOccurred(), "Expected pod to terminate.")
	}

	By("Expecting the volume mount not to be found.")
	result, err = framework.SSH(fmt.Sprintf("mount | grep %s", clientPod.UID), nodeIP, framework.TestContext.Provider)
	framework.LogSSHResult(result)
	Expect(err).NotTo(HaveOccurred(), "Encountered SSH error.")
	Expect(result.Stdout).To(BeEmpty(), "Expected grep stdout to be empty (i.e. no mount found).")
	framework.Logf("Volume unmounted on node %s", clientPod.Spec.NodeName)
}

// initTestCase initializes spec resources (pv, pvc, and pod) and returns pointers to be consumed
// by the test.
func initTestCase(f *framework.Framework, c clientset.Interface, pvConfig framework.PersistentVolumeConfig, pvcConfig framework.PersistentVolumeClaimConfig, ns, nodeName string) (*v1.Pod, *v1.PersistentVolume, *v1.PersistentVolumeClaim) {
	pv, pvc, err := framework.CreatePVPVC(c, pvConfig, pvcConfig, ns, false)
	defer func() {
		if err != nil {
			framework.DeletePersistentVolumeClaim(c, pvc.Name, ns)
			framework.DeletePersistentVolume(c, pv.Name)
		}
	}()
	Expect(err).NotTo(HaveOccurred())
	pod := framework.MakePod(ns, []*v1.PersistentVolumeClaim{pvc}, true, "")
	pod.Spec.NodeName = nodeName
	framework.Logf("Creating NFS client pod.")
	pod, err = c.CoreV1().Pods(ns).Create(pod)
	framework.Logf("NFS client Pod %q created on Node %q", pod.Name, nodeName)
	Expect(err).NotTo(HaveOccurred())
	defer func() {
		if err != nil {
			framework.DeletePodWithWait(f, c, pod)
		}
	}()
	err = framework.WaitForPodRunningInNamespace(c, pod)
	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("Pod %q timed out waiting for phase: Running", pod.Name))
	// Return created api objects
	pod, err = c.CoreV1().Pods(ns).Get(pod.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	pvc, err = c.CoreV1().PersistentVolumeClaims(ns).Get(pvc.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	pv, err = c.CoreV1().PersistentVolumes().Get(pv.Name, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	return pod, pv, pvc
}

// tearDownTestCase destroy resources created by initTestCase.
func tearDownTestCase(c clientset.Interface, f *framework.Framework, ns string, client *v1.Pod, pvc *v1.PersistentVolumeClaim, pv *v1.PersistentVolume) {
	// Ignore deletion errors.  Failing on them will interrupt test cleanup.
	framework.DeletePodWithWait(f, c, client)
	framework.DeletePersistentVolumeClaim(c, pvc.Name, ns)
	framework.DeletePersistentVolume(c, pv.Name)
}

// kubeletCommand performs `start`, `restart`, or `stop` on the kubelet running on the node of the target pod and waits
// for the desired statues..
// - First issues the command via `systemctl`
// - If `systemctl` returns stderr "command not found, issues the command via `service`
// - If `service` also returns stderr "command not found", the test is aborted.
// Allowed kubeletOps are `kStart`, `kStop`, and `kRestart`
func kubeletCommand(kOp kubeletOpt, c clientset.Interface, pod *v1.Pod) {
	command := ""
	sudoPresent := false
	systemctlPresent := false
	kubeletPid := ""

	nodeIP, err := framework.GetHostExternalAddress(c, pod)
	Expect(err).NotTo(HaveOccurred())
	nodeIP = nodeIP + ":22"

	framework.Logf("Checking if sudo command is present")
	sshResult, err := framework.SSH("sudo --version", nodeIP, framework.TestContext.Provider)
	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("SSH to Node %q errored.", pod.Spec.NodeName))
	if !strings.Contains(sshResult.Stderr, "command not found") {
		sudoPresent = true
	}

	framework.Logf("Checking if systemctl command is present")
	sshResult, err = framework.SSH("systemctl --version", nodeIP, framework.TestContext.Provider)
	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("SSH to Node %q errored.", pod.Spec.NodeName))
	if !strings.Contains(sshResult.Stderr, "command not found") {
		command = fmt.Sprintf("systemctl %s kubelet", string(kOp))
		systemctlPresent = true
	} else {
		command = fmt.Sprintf("service kubelet %s", string(kOp))
	}
	if sudoPresent {
		command = fmt.Sprintf("sudo %s", command)
	}

	if kOp == kRestart {
		kubeletPid = getKubeletMainPid(nodeIP, sudoPresent, systemctlPresent)
	}

	framework.Logf("Attempting `%s`", command)
	sshResult, err = framework.SSH(command, nodeIP, framework.TestContext.Provider)
	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("SSH to Node %q errored.", pod.Spec.NodeName))
	framework.LogSSHResult(sshResult)
	Expect(sshResult.Code).To(BeZero(), "Failed to [%s] kubelet:\n%#v", string(kOp), sshResult)

	if kOp == kStop {
		if ok := framework.WaitForNodeToBeNotReady(c, pod.Spec.NodeName, NodeStateTimeout); !ok {
			framework.Failf("Node %s failed to enter NotReady state", pod.Spec.NodeName)
		}
	}
	if kOp == kRestart {
		// Wait for a minute to check if kubelet Pid is getting changed
		isPidChanged := false
		for start := time.Now(); time.Since(start) < 1*time.Minute; time.Sleep(2 * time.Second) {
			kubeletPidAfterRestart := getKubeletMainPid(nodeIP, sudoPresent, systemctlPresent)
			if kubeletPid != kubeletPidAfterRestart {
				isPidChanged = true
				break
			}
		}
		Expect(isPidChanged).To(BeTrue(), "Kubelet PID remained unchanged after restarting Kubelet")
		framework.Logf("Noticed that kubelet PID is changed. Waiting for 30 Seconds for Kubelet to come back")
		time.Sleep(30 * time.Second)
	}
	if kOp == kStart || kOp == kRestart {
		// For kubelet start and restart operations, Wait until Node becomes Ready
		if ok := framework.WaitForNodeToBeReady(c, pod.Spec.NodeName, NodeStateTimeout); !ok {
			framework.Failf("Node %s failed to enter Ready state", pod.Spec.NodeName)
		}
	}
}

// return the Main PID of the Kubelet Process
func getKubeletMainPid(nodeIP string, sudoPresent bool, systemctlPresent bool) string {
	command := ""
	if systemctlPresent {
		command = "systemctl status kubelet | grep 'Main PID'"
	} else {
		command = "service kubelet status | grep 'Main PID'"
	}
	if sudoPresent {
		command = fmt.Sprintf("sudo %s", command)
	}
	framework.Logf("Attempting `%s`", command)
	sshResult, err := framework.SSH(command, nodeIP, framework.TestContext.Provider)
	Expect(err).NotTo(HaveOccurred(), fmt.Sprintf("SSH to Node %q errored.", nodeIP))
	framework.LogSSHResult(sshResult)
	Expect(sshResult.Code).To(BeZero(), "Failed to get kubelet PID")
	Expect(sshResult.Stdout).NotTo(BeEmpty(), "Kubelet Main PID should not be Empty")
	return sshResult.Stdout
}

// podExec wraps RunKubectl to execute a bash cmd in target pod
func podExec(pod *v1.Pod, bashExec string) (string, error) {
	return framework.RunKubectl("exec", fmt.Sprintf("--namespace=%s", pod.Namespace), pod.Name, "--", "/bin/sh", "-c", bashExec)
}
<|endoftext|>


<commit_before>// Package routeros provides a programmatic interface to to the Mikrotik RouterOS API
package routeros

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/jcelliott/lumber"
	"io"
	"net"
	"strconv"
	"strings"
)

// Client is a RouterOS API client.
type Client struct {
	// Network Address.
	// E.g. "10.0.0.1:8728" or "router.example.com:8728"
	address  string
	user     string
	password string
	logger   *lumber.Logger
	debug    bool     // debug logging enabled
	ready    bool     // Ready for work (login ok and connection not terminated)
	conn     net.Conn // Connection to pass around
}

// Pair is a Key-Value pair for RouterOS Attribute, Query, and Reply words
// use slices of pairs instead of map because we care about order
type Pair struct {
	Key   string
	Value string
	// Op is used for Query words to signify logical operations
	// valid operators are -, =, <, >
	// see http://wiki.mikrotik.com/wiki/Manual:API#Queries for details.
	Op    string
}

type Query struct {
	Pairs    []Pair
	Op       string
	Proplist []string
}

func NewPair(key string, value string) *Pair {
	p := new(Pair)
	p.Key = key
	p.Value = value
	return p
}

// Get value for a specific key
func GetPairValue(p []Pair, key string) (string, error) {
	for _, v := range p {
		if v.Key == key {
			return v.Value, nil
		}
	}
	return "", errors.New("key not found")
}

// Create a new instance of the RouterOS API client
func New(address string) (*Client, error) {
	// basic validation of host address
	_, _, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}

	var c Client
	c.address = address

	return &c, nil
}

func (c *Client) Close() {
	c.conn.Close()
}

func (c *Client) Connect(user string, password string) error {
	conn, err := net.Dial("tcp", c.address)
	if err != nil {
		return err
	}

	// stash conn in instance
	c.conn = conn

	// try to log in
	res, err := c.Call("/login", nil)
	if err != nil {
		return err
	}

	// handle challenge/response
	challengeEnc, err := GetPairValue(res, "ret")
	if err != nil {
		return errors.New("Didn't get challenge from ROS")
	}
	challenge, err := hex.DecodeString(challengeEnc)
	if err != nil {
		return err
	}
	h := md5.New()
	io.WriteString(h, "\000")
	io.WriteString(h, password)
	h.Write(challenge)
	resp := fmt.Sprintf("00%x", h.Sum(nil))
	var loginParams []Pair
	loginParams = append(loginParams, *NewPair("name", password))
	loginParams = append(loginParams, *NewPair("response", resp))

	// try to log in again with challenge/response
	res, err = c.Call("/login", loginParams)
	if err != nil {
		return err
	}
	fmt.Println(res)

	return nil
}

// Encode and send a single line
func (c *Client) send(word string) error {
	bword := []byte(word)
	prefix := prefixlen(len(bword))

	_, err := c.conn.Write(prefix.Bytes())
	if err != nil {
		return err
	}

	_, err = c.conn.Write(bword)
	if err != nil {
		return err
	}

	return nil
}

// Listen for reply
func (c *Client) receive() ([]Pair, error) {
	var pairs []Pair

	for {
		length := c.getlen()
		if length == 0 {
			break
		}
		inbuf := make([]byte, length)
		c.conn.Read(inbuf)
		word := string(inbuf)

		if word == "!done" {
			continue
		}

		if strings.Contains(word, "=") {
			var p Pair
			parts := strings.SplitN(word, "=", 3)
			p.Key = parts[1]
			p.Value = parts[2]
			pairs = append(pairs, p)
		}

	}

	return pairs, nil
}

func (c *Client) Query(command string, q Query) ([]Pair, error) {
	err := c.send(command)
	if err != nil {
		return nil, err
	}

	// Set property list if present
	if len(q.Proplist) > 0 {
		proplist := fmt.Sprintf("=.proplist=%s", strings.Join(q.Proplist, ","))
		err = c.send(proplist)
		if err != nil {
			return nil, err
		}
	}

	// send params if we got them
	if len(q.Pairs) > 0 {
		for _, v := range q.Pairs {
			word := fmt.Sprintf("?%s%s=%s", v.Op, v.Key, v.Value)
			c.send(word)
		}

		if q.Op != "" {
			word := fmt.Sprintf("?#%s", q.Op)
			c.send(word)
		}
	}

	// send terminator
	err = c.send("")
	if err != nil {
		return nil, err
	}

	res, err := c.receive()
	if err != nil {
		return nil, err
	}

	return res, nil
}

func (c *Client) Call(command string, params []Pair) ([]Pair, error) {
	err := c.send(command)
	if err != nil {
		return nil, err
	}

	// send params if we got them
	if len(params) > 0 {
		for _, v := range params {
			word := fmt.Sprintf("=%s=%s", v.Key, v.Value)
			c.send(word)
		}
	}

	// send terminator
	err = c.send("")
	if err != nil {
		return nil, err
	}

	res, err := c.receive()
	if err != nil {
		return nil, err
	}

	return res, nil
}

// Enabled logging by providing a Logger instance from Lumber.
func (c *Client) Logging(l *lumber.Logger, level string) {
	c.logger = l
	if level == "DEBUG" {
		c.debug = true
	}
}

// Get just one byte because MT's size prefix is overoptimized
func (c *Client) getone() int {
	charlet := make([]byte, 1)
	c.conn.Read(charlet)
	numlet := int(charlet[0])
	return numlet
}

// Decode RouterOS API Word Size Prefix / Figure out how much to read
// TODO: based on MT Docs. Look for way to make this cleaner later
func (client *Client) getlen() int64 {
	c := int64(client.getone())

	if (c & 0x80) == 0x00 {

	} else if (c & 0xC0) == 0x80 {
		c &= ^0xC0
		c <<= 8
		c += int64(client.getone())
	} else if (c & 0xE0) == 0xC0 {
		c &= ^0xE0
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
	} else if (c & 0xF0) == 0xE0 {
		c &= ^0xF0
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
	} else if (c & 0xF8) == 0xF0 {
		c = int64(client.getone())
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
	}

	return c
}

// Calculate RouterOS API Word Size Prefix
// TODO: based on MT Docs. Look for way to make this cleaner later
func prefixlen(l int) bytes.Buffer {
	var b bytes.Buffer

	if l < 0x80 {
		b.Write([]byte(string(l)))
	} else if l < 0x4000 {
		l |= 0x8000
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	} else if l < 0x200000 {
		l |= 0xC00000
		b.Write([]byte(strconv.Itoa((l >> 16) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	} else if l < 0x10000000 {
		l |= 0xE0000000
		b.Write([]byte(strconv.Itoa((l >> 24) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 16) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	} else {
		b.Write([]byte(strconv.Itoa(0xF0)))
		b.Write([]byte(strconv.Itoa((l >> 24) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 16) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	}

	return b
}
<commit_msg>improved error handling in login<commit_after>// Package routeros provides a programmatic interface to to the Mikrotik RouterOS API
package routeros

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/jcelliott/lumber"
	"io"
	"net"
	"strconv"
	"strings"
)

// Client is a RouterOS API client.
type Client struct {
	// Network Address.
	// E.g. "10.0.0.1:8728" or "router.example.com:8728"
	address  string
	user     string
	password string
	logger   *lumber.Logger
	debug    bool     // debug logging enabled
	ready    bool     // Ready for work (login ok and connection not terminated)
	conn     net.Conn // Connection to pass around
}

// Pair is a Key-Value pair for RouterOS Attribute, Query, and Reply words
// use slices of pairs instead of map because we care about order
type Pair struct {
	Key   string
	Value string
	// Op is used for Query words to signify logical operations
	// valid operators are -, =, <, >
	// see http://wiki.mikrotik.com/wiki/Manual:API#Queries for details.
	Op    string
}

type Query struct {
	Pairs    []Pair
	Op       string
	Proplist []string
}

func NewPair(key string, value string) *Pair {
	p := new(Pair)
	p.Key = key
	p.Value = value
	return p
}

// Get value for a specific key
func GetPairValue(p []Pair, key string) (string, error) {
	for _, v := range p {
		if v.Key == key {
			return v.Value, nil
		}
	}
	return "", errors.New("key not found")
}

// Create a new instance of the RouterOS API client
func New(address string) (*Client, error) {
	// basic validation of host address
	_, _, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}

	var c Client
	c.address = address

	return &c, nil
}

func (c *Client) Close() {
	c.conn.Close()
}

func (c *Client) Connect(user string, password string) error {
	conn, err := net.Dial("tcp", c.address)
	if err != nil {
		return err
	}

	// stash conn in instance
	c.conn = conn

	// try to log in
	res, err := c.Call("/login", nil)
	if err != nil {
		return err
	}

	// handle challenge/response
	challengeEnc, err := GetPairValue(res, "ret")
	if err != nil {
		return errors.New("Didn't get challenge from ROS")
	}
	challenge, err := hex.DecodeString(challengeEnc)
	if err != nil {
		return err
	}
	h := md5.New()
	io.WriteString(h, "\000")
	io.WriteString(h, password)
	h.Write(challenge)
	resp := fmt.Sprintf("00%x", h.Sum(nil))
	var loginParams []Pair
	loginParams = append(loginParams, *NewPair("name", password))
	loginParams = append(loginParams, *NewPair("response", resp))

	// try to log in again with challenge/response
	res, err = c.Call("/login", loginParams)
	if err != nil {
		return err
	}

	if len(res) > 0 {
		return fmt.Errorf("Unexpected result on login: %+v", res)
	}

	return nil
}

// Encode and send a single line
func (c *Client) send(word string) error {
	bword := []byte(word)
	prefix := prefixlen(len(bword))

	_, err := c.conn.Write(prefix.Bytes())
	if err != nil {
		return err
	}

	_, err = c.conn.Write(bword)
	if err != nil {
		return err
	}

	return nil
}

// Listen for reply
func (c *Client) receive() ([]Pair, error) {
	var pairs []Pair

	for {
		length := c.getlen()
		if length == 0 {
			break
		}
		inbuf := make([]byte, length)
		c.conn.Read(inbuf)
		word := string(inbuf)

		if word == "!done" {
			continue
		}

		if strings.Contains(word, "=") {
			var p Pair
			parts := strings.SplitN(word, "=", 3)
			p.Key = parts[1]
			p.Value = parts[2]
			pairs = append(pairs, p)
		}

	}

	return pairs, nil
}

func (c *Client) Query(command string, q Query) ([]Pair, error) {
	err := c.send(command)
	if err != nil {
		return nil, err
	}

	// Set property list if present
	if len(q.Proplist) > 0 {
		proplist := fmt.Sprintf("=.proplist=%s", strings.Join(q.Proplist, ","))
		err = c.send(proplist)
		if err != nil {
			return nil, err
		}
	}

	// send params if we got them
	if len(q.Pairs) > 0 {
		for _, v := range q.Pairs {
			word := fmt.Sprintf("?%s%s=%s", v.Op, v.Key, v.Value)
			c.send(word)
		}

		if q.Op != "" {
			word := fmt.Sprintf("?#%s", q.Op)
			c.send(word)
		}
	}

	// send terminator
	err = c.send("")
	if err != nil {
		return nil, err
	}

	res, err := c.receive()
	if err != nil {
		return nil, err
	}

	return res, nil
}

func (c *Client) Call(command string, params []Pair) ([]Pair, error) {
	err := c.send(command)
	if err != nil {
		return nil, err
	}

	// send params if we got them
	if len(params) > 0 {
		for _, v := range params {
			word := fmt.Sprintf("=%s=%s", v.Key, v.Value)
			c.send(word)
		}
	}

	// send terminator
	err = c.send("")
	if err != nil {
		return nil, err
	}

	res, err := c.receive()
	if err != nil {
		return nil, err
	}

	return res, nil
}

// Enabled logging by providing a Logger instance from Lumber.
func (c *Client) Logging(l *lumber.Logger, level string) {
	c.logger = l
	if level == "DEBUG" {
		c.debug = true
	}
}

// Get just one byte because MT's size prefix is overoptimized
func (c *Client) getone() int {
	charlet := make([]byte, 1)
	c.conn.Read(charlet)
	numlet := int(charlet[0])
	return numlet
}

// Decode RouterOS API Word Size Prefix / Figure out how much to read
// TODO: based on MT Docs. Look for way to make this cleaner later
func (client *Client) getlen() int64 {
	c := int64(client.getone())

	if (c & 0x80) == 0x00 {

	} else if (c & 0xC0) == 0x80 {
		c &= ^0xC0
		c <<= 8
		c += int64(client.getone())
	} else if (c & 0xE0) == 0xC0 {
		c &= ^0xE0
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
	} else if (c & 0xF0) == 0xE0 {
		c &= ^0xF0
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
	} else if (c & 0xF8) == 0xF0 {
		c = int64(client.getone())
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
		c <<= 8
		c += int64(client.getone())
	}

	return c
}

// Calculate RouterOS API Word Size Prefix
// TODO: based on MT Docs. Look for way to make this cleaner later
func prefixlen(l int) bytes.Buffer {
	var b bytes.Buffer

	if l < 0x80 {
		b.Write([]byte(string(l)))
	} else if l < 0x4000 {
		l |= 0x8000
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	} else if l < 0x200000 {
		l |= 0xC00000
		b.Write([]byte(strconv.Itoa((l >> 16) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	} else if l < 0x10000000 {
		l |= 0xE0000000
		b.Write([]byte(strconv.Itoa((l >> 24) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 16) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	} else {
		b.Write([]byte(strconv.Itoa(0xF0)))
		b.Write([]byte(strconv.Itoa((l >> 24) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 16) & 0xFF)))
		b.Write([]byte(strconv.Itoa((l >> 8) & 0xFF)))
		b.Write([]byte(strconv.Itoa(l & 0xFF)))
	}

	return b
}
<|endoftext|>


<commit_before>package sdk

import (
	"net/http"
	"fmt"
)

func NewHTTPStatusError(statusCode int, content []byte) *HTTPStatusError {
	return &HTTPStatusError{
		statusCode: statusCode,
		content:    content,
	}
}

// HTTPStatusError stands in for the error type but also provides convenience methods
// for accessing the status code and content of the request that caused the error.
// Instances of this type are returned by sdk.HTTPSender.Send(). When nil, the methods
// of this type behave as if called on a non-nil instance instantiated with http.StatusOK (200).
type HTTPStatusError struct {
	statusCode int
	content    []byte
}

func (e *HTTPStatusError) Error() string {
	if e == nil {
		return statusText(http.StatusOK)
	}
	return statusText(e.statusCode)
}

func statusText(code int) string {
	return fmt.Sprintf("HTTP %d %s", code, http.StatusText(code))
}

func (e *HTTPStatusError) StatusCode() int {
	if e == nil {
		return http.StatusOK
	}
	return e.statusCode
}

func (e *HTTPStatusError) Content() []byte {
	if e == nil {
		return nil
	}
	return e.content
}
<commit_msg>goimports.<commit_after>package sdk

import (
	"fmt"
	"net/http"
)

func NewHTTPStatusError(statusCode int, content []byte) *HTTPStatusError {
	return &HTTPStatusError{
		statusCode: statusCode,
		content:    content,
	}
}

// HTTPStatusError stands in for the error type but also provides convenience methods
// for accessing the status code and content of the request that caused the error.
// Instances of this type are returned by sdk.HTTPSender.Send(). When nil, the methods
// of this type behave as if called on a non-nil instance instantiated with http.StatusOK (200).
type HTTPStatusError struct {
	statusCode int
	content    []byte
}

func (e *HTTPStatusError) Error() string {
	if e == nil {
		return statusText(http.StatusOK)
	}
	return statusText(e.statusCode)
}

func statusText(code int) string {
	return fmt.Sprintf("HTTP %d %s", code, http.StatusText(code))
}

func (e *HTTPStatusError) StatusCode() int {
	if e == nil {
		return http.StatusOK
	}
	return e.statusCode
}

func (e *HTTPStatusError) Content() []byte {
	if e == nil {
		return nil
	}
	return e.content
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/vladimirvivien/automi/collectors"
	"github.com/vladimirvivien/automi/emitters"
	"github.com/vladimirvivien/automi/stream"
)

func main() {
	// stream := stream.New([]string{
	// 	"Hello World",
	// 	"Hello Milkyway",
	// 	"Hello Universe",
	// })

	f, err := os.Open("./twotw.txt")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	stream := stream.New(emitters.Reader(f, bufio.ScanLines))

	regSpaces := regexp.MustCompile(`\s+`)

	// coerce incoming data to string
	// then remove extranuous spaces.
	stream.Process(func(line interface{}) string {
		return regSpaces.ReplaceAllLiteralString(fmt.Sprintf("%s", line), " ")
	})

	// split lines into slice of words which
	// are then streamed individually
	stream.FlatMap(func(line string) []string {
		return strings.Split(line, " ")
	})

	// map word to an array where array[0]=word
	// array[1]=1 to mark the occurence of word
	stream.Map(func(word string) [2]interface{} {
		return [2]interface{}{word, 1}
	})

	// Next:
	// 1) batch the stream of arrays from last step
	// 2) group by position 0 which has the word which returns a map[word]count
	// 3) Sum by the keys of the group
	stream.Batch().GroupByPos(0).SumAllKeys()

	// sink resunt to collector function
	stream.SinkTo(collectors.Func(func(items interface{}) error {
		words := items.([]map[interface{}]float64)
		for _, wmap := range words {
			for word, count := range wmap {
				fmt.Printf("%s:%0.0f\n", word, count)
			}
		}
		return nil
	}))

	// stream.SinkTo(collectors.Writer(os.Stdout))

	// open the stream
	if err := <-stream.Open(); err != nil {
		fmt.Println(err)
		return
	}
}
<commit_msg>fix wordcount example<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/vladimirvivien/automi/collectors"
	"github.com/vladimirvivien/automi/emitters"
	"github.com/vladimirvivien/automi/stream"
)

func main() {
	// stream := stream.New([]string{
	// 	"Hello World",
	// 	"Hello Milkyway",
	// 	"Hello Universe",
	// })

	f, err := os.Open("./twotw.txt")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	stream := stream.New(emitters.Scanner(f, bufio.ScanLines))

	regSpaces := regexp.MustCompile(`\s+`)

	// coerce incoming data to string
	// then remove extranuous spaces.
	stream.Process(func(line interface{}) string {
		return regSpaces.ReplaceAllLiteralString(fmt.Sprintf("%s", line), " ")
	})

	// split lines into slice of words which
	// are then streamed individually
	stream.FlatMap(func(line string) []string {
		return strings.Split(line, " ")
	})

	// map word to an array where array[0]=word
	// array[1]=1 to mark the occurence of word
	stream.Map(func(word string) [2]interface{} {
		return [2]interface{}{word, 1}
	})

	// Next:
	// 1) batch the stream of arrays from last step
	// 2) group by position 0 which has the word which returns a map[word]count
	// 3) Sum by the keys of the group
	stream.Batch().GroupByPos(0).SumAllKeys()

	// sink resunt to collector function
	stream.SinkTo(collectors.Func(func(items interface{}) error {
		words := items.([]map[interface{}]float64)
		for _, wmap := range words {
			for word, count := range wmap {
				fmt.Printf("%s:%0.0f\n", word, count)
			}
		}
		return nil
	}))

	// stream.SinkTo(collectors.Writer(os.Stdout))

	// open the stream
	if err := <-stream.Open(); err != nil {
		fmt.Println(err)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/twinj/uuid"
	"net/http"
	"strings"
	"testing"
	"time"
)

// TODO: factor out magic numbers and strings

func TestRoutes(t *testing.T) {
	go StartServer()
	time.Sleep(100 * time.Millisecond)

	// Test message api
	testMessage(t)

	// Test user api
	testUser(t)

	// Test auth api
	testAuth(t)
}

func testMessage(t *testing.T) {
	// GET /api/message
	// TODO: Decide on proper response and functionality for GET /api/message
	r := Get("/api/message", t)
	if r != 200 {
		t.Errorf("GET /api/message did not return 200, instead returned %d\n", r)
	}

	// POST /api/message
	// Creates a new message with the given body parameter.
	// Should return 201 on successful creation or 5** for server error.
	r = Post("/api/message", "{\"body\": \"This is the body of the first message.\"}", t)
	if r != 201 {
		t.Errorf("POST /api/message did not return 201, instead returned %d\n", r)
	}
}

func testUser(t *testing.T) {
	id := uuid.NewV1()
	// GET /api/user
	// TODO: Decide on use of this call. Is it needed?
	r := Get("/api/user", t)
	if r != 200 {
		t.Errorf("GET /api/user did not return 200, instead returned %d\n", r)
	}

	// PUT /api/user/{id}
	// Create a new user or update a user record.
	// Should return 201 on successful creation, 200 on update, or 5** for server error.
	r = Put("/api/user/"+id.String(), "{\"id\": \""+id.String()+"\", \"username\":\"fdsa\", \"password\": \"fdsa\"}", t)
	if r != 201 {
		t.Errorf("PUT /api/user/%s did not return 201, instead returned %d\n", id.String(), r)
	}

	// GET /api/user/{id}
	// Returns the given user resource
	// Should return 200 if found or 404 if resource isn't found.
	r = Get("/api/user/"+id.String(), t)
	if r != 200 {
		t.Errorf("GET /api/user/%s did not return 200, instead returned %d\n", id.String(), r)
	}

	id2 := uuid.NewV1()
	// POST /api/user
	// Create a new user
	// TODO: Remove id from this and generate a unique one automatically.
	// Should return 201 for created or 409 for conflict.
	r = Post("/api/user", "{\"username\":\"asdf\", \"password\": \"fdsa\"}", t)
	if r != 201 && r != 409 {
		t.Errorf("POST /api/user did not return 201, instead returned %d\n", r)
	}

	// Get newly created user.
	r = Get("/api/user/"+id2.String(), t)
	if r != 200 {
		t.Errorf("GET /api/user/%s did not return 200, instead returned %d\n", id2.String(), r)
	}

	// Update the user's record.
	r = Put("/api/user/"+id2.String(), "{\"id\": \""+id2.String()+"\", \"username\":\"fdsa\", \"password\": \"fdsa\"}", t)
	if r != 200 {
		t.Errorf("PUT /api/user/%s did not return 200, instead returned %d\n", id2.String(), r)
	}

	// DELETE /api/user/{id}
	// Delete's the user resource.
	// Should return 204 for deletion, 404 if not found, or 5** for server error
	r = Delete("/api/user/"+id.String(), t)
	if r != 204 {
		t.Errorf("DELETE did not return 204, instead returned %d\n", r)
	}

	// Delete user record.
	r = Delete("/api/user/"+id2.String(), t)
	if r != 204 {
		t.Errorf("DELETE did not return 204, instead returned %d\n", r)
	}

	// Test delete on non existant user
	r = Delete("/api/user/1337", t)
	if r != 404 {
		t.Errorf("DELETE on user not found\n")
	}
}

func testAuth(t *testing.T) {
	// GET /api/auth
	// TODO: Will return an auth token or something eventually...
	// Should return 405 unauthorized until implemented.
	r := Get("/api/auth", t)
	if r != 405 {
		t.Errorf("GET /api/auth did not return 405, instead returned %d\n", r)
	}
}

// Make a GET request to the specified path.
func Get(path string, t *testing.T) int {
	get_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	res, err := http.Get(get_url)
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}
	defer res.Body.Close()

	fmt.Printf("GET %s: %d\n", path, res.StatusCode)
	return res.StatusCode
}

// Make a POST request to the specified path and with the given JSON body.
func Post(path, body string, t *testing.T) int {
	post_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	bodyType := "application/json"
	res, err := http.Post(post_url, bodyType, strings.NewReader(body))
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}
	defer res.Body.Close()

	fmt.Printf("POST %s: %d\n", path, res.StatusCode)
	return res.StatusCode
}

// Make a PUT request to the specified path with the given JSON body
func Put(path, body string, t *testing.T) int {
	put_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	bodyType := "application/json"
	req, err := http.NewRequest("PUT", put_url, strings.NewReader(body))
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}

	// Add a header for JSON content type
	req.Header.Add("Content-Type", bodyType)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}
	defer res.Body.Close()

	fmt.Printf("PUT %s: %d\n", path, res.StatusCode)
	return res.StatusCode
}

// Make a DELETE request to the given path.
func Delete(path string, t *testing.T) int {
	del_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	req, err := http.NewRequest("DELETE", del_url, strings.NewReader(""))
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}
	defer res.Body.Close()

	fmt.Printf("DELETE %s: %d\n", path, res.StatusCode)
	return res.StatusCode
}
<commit_msg>Fix failing user get test. (#5)<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/twinj/uuid"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"
)

// TODO: factor out magic numbers and strings

func TestRoutes(t *testing.T) {
	go StartServer()
	time.Sleep(100 * time.Millisecond)

	// Test message api
	testMessage(t)

	// Test user api
	testUser(t)

	// Test auth api
	testAuth(t)
}

func testMessage(t *testing.T) {
	// GET /api/message
	// TODO: Decide on proper response and functionality for GET /api/message
	r := Get("/api/message", t)
	if r != 200 {
		t.Errorf("GET /api/message did not return 200, instead returned %d\n", r)
	}

	// POST /api/message
	// Creates a new message with the given body parameter.
	// Should return 201 on successful creation or 5** for server error.
	_, r = Post("/api/message", "{\"body\": \"This is the body of the first message.\"}", t)
	if r != 201 {
		t.Errorf("POST /api/message did not return 201, instead returned %d\n", r)
	}
}

func testUser(t *testing.T) {
	id := uuid.NewV1()
	// GET /api/user
	// TODO: Decide on use of this call. Is it needed?
	r := Get("/api/user", t)
	if r != 200 {
		t.Errorf("GET /api/user did not return 200, instead returned %d\n", r)
	}

	// PUT /api/user/{id}
	// Create a new user or update a user record.
	// Should return 201 on successful creation, 200 on update, or 5** for server error.
	r = Put("/api/user/"+id.String(), "{\"id\": \""+id.String()+"\", \"username\":\"fdsa\", \"password\": \"fdsa\"}", t)
	if r != 201 {
		t.Errorf("PUT /api/user/%s did not return 201, instead returned %d\n", id.String(), r)
	}

	// GET /api/user/{id}
	// Returns the given user resource
	// Should return 200 if found or 404 if resource isn't found.
	r = Get("/api/user/"+id.String(), t)
	if r != 200 {
		t.Errorf("GET /api/user/%s did not return 200, instead returned %d\n", id.String(), r)
	}

	// POST /api/user
	// Create a new user
	// Should return 201 for created or 409 for conflict.
	res, r := Post("/api/user", "{\"username\":\"asdf\", \"password\": \"fdsa\"}", t)
	if r != 201 && r != 409 {
		t.Errorf("POST /api/user did not return 201, instead returned %d\n", r)
	}

	var user User
	err := json.Unmarshal(res, &user)
	if err != nil {
		t.Errorf(err.Error())
	}

	// Get newly created user.
	r = Get("/api/user/"+user.Id, t)
	if r != 200 {
		t.Errorf("GET /api/user/%s did not return 200, instead returned %d\n", user.Id, r)
	}

	// Update the user's record.
	r = Put("/api/user/"+user.Id, "{\"id\": \""+user.Id+"\", \"username\":\"fdsa\", \"password\": \"fdsa\"}", t)
	if r != 200 {
		t.Errorf("PUT /api/user/%s did not return 200, instead returned %d\n", user.Id, r)
	}

	// DELETE /api/user/{id}
	// Delete's the user resource.
	// Should return 204 for deletion, 404 if not found, or 5** for server error
	r = Delete("/api/user/"+id.String(), t)
	if r != 204 {
		t.Errorf("DELETE did not return 204, instead returned %d\n", r)
	}

	// Delete user record.
	r = Delete("/api/user/"+user.Id, t)
	if r != 204 {
		t.Errorf("DELETE did not return 204, instead returned %d\n", r)
	}

	// Test delete on non existant user
	r = Delete("/api/user/1337", t)
	if r != 404 {
		t.Errorf("DELETE on user not found\n")
	}
}

func testAuth(t *testing.T) {
	// GET /api/auth
	// TODO: Will return an auth token or something eventually...
	// Should return 405 unauthorized until implemented.
	r := Get("/api/auth", t)
	if r != 405 {
		t.Errorf("GET /api/auth did not return 405, instead returned %d\n", r)
	}
}

// Make a GET request to the specified path.
func Get(path string, t *testing.T) int {
	get_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	res, err := http.Get(get_url)
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}
	defer res.Body.Close()

	fmt.Printf("GET %s: %d\n", path, res.StatusCode)
	return res.StatusCode
}

// Make a POST request to the specified path and with the given JSON body.
func Post(path, body string, t *testing.T) ([]byte, int) {
	post_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	bodyType := "application/json"
	res, err := http.Post(post_url, bodyType, strings.NewReader(body))
	if err != nil {
		t.Errorf(err.Error())
		return nil, 0
	}
	defer res.Body.Close()
	resbody, err := ioutil.ReadAll(res.Body)
	if err != nil {
		t.Errorf(err.Error())
		return nil, 0
	}

	fmt.Printf("POST %s: %d\n", path, res.StatusCode)
	return resbody, res.StatusCode
}

// Make a PUT request to the specified path with the given JSON body
func Put(path, body string, t *testing.T) int {
	put_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	bodyType := "application/json"
	req, err := http.NewRequest("PUT", put_url, strings.NewReader(body))
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}

	// Add a header for JSON content type
	req.Header.Add("Content-Type", bodyType)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}
	defer res.Body.Close()

	fmt.Printf("PUT %s: %d\n", path, res.StatusCode)
	return res.StatusCode
}

// Make a DELETE request to the given path.
func Delete(path string, t *testing.T) int {
	del_url := fmt.Sprintf("http://127.0.0.1:8080%s", path)
	req, err := http.NewRequest("DELETE", del_url, strings.NewReader(""))
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Errorf(err.Error())
		return 0
	}
	defer res.Body.Close()

	fmt.Printf("DELETE %s: %d\n", path, res.StatusCode)
	return res.StatusCode
}
<|endoftext|>


<commit_before>package rpcc

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
)

var (
	// ErrStreamClosing indicates that the operation is illegal because
	// the stream is closing and there are no pending messages.
	ErrStreamClosing = errors.New("rpcc: the stream is closing")
)

type streamMsg struct {
	method string
	data   []byte
}

type messageBuffer struct {
	ch    chan *streamMsg
	mu    sync.Mutex
	queue []*streamMsg
}

func newMessageBuffer() *messageBuffer {
	return &messageBuffer{
		ch: make(chan *streamMsg, 1),
	}
}

func (b *messageBuffer) store(m *streamMsg) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) == 0 {
		select {
		case b.ch <- m:
			return
		default:
		}
	}
	b.queue = append(b.queue, m)
}

func (b *messageBuffer) load() {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) > 0 {
		select {
		case b.ch <- b.queue[0]:
			b.queue[0] = nil // Remove reference from underlying array.
			b.queue = b.queue[1:]
		default:
		}
	}
}

func (b *messageBuffer) get() <-chan *streamMsg {
	return b.ch
}

// Stream represents a stream of notifications for a certain method.
type Stream interface {
	// RecvMsg unmarshals pending messages onto m. Blocks until the
	// next message is received, context is canceled or stream is
	// closed.
	RecvMsg(m interface{}) error
	// Close closes the stream and no new messages will be received.
	// RecvMsg will return ErrStreamClosing once all pending messages
	// have been received.
	Close() error
}

// NewStream creates a new stream that listens to notifications from the
// RPC server. This function is called by generated code.
func NewStream(ctx context.Context, method string, conn *Conn) (Stream, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	s := &streamClient{userCtx: ctx, done: make(chan struct{})}
	s.msgBuf = newMessageBuffer()
	s.ctx, s.cancel = context.WithCancel(context.Background())

	var err error
	s.remove, err = conn.listen(method, s)
	if err != nil {
		return nil, err
	}

	go func() {
		select {
		case <-s.ctx.Done():
		case <-conn.ctx.Done():
			s.close(ErrConnClosing)
		case <-ctx.Done():
			s.close(ctx.Err())
		}
	}()

	return s, nil
}

type streamClient struct {
	userCtx context.Context
	ctx     context.Context
	cancel  context.CancelFunc

	// msgBuf stores all incoming messages
	// until they are ready to be received.
	msgBuf *messageBuffer

	mu     sync.Mutex // Protects following.
	remove func()     // Unsubscribes from messages.

	done chan struct{} // Protects err.
	err  error
}

func (s *streamClient) RecvMsg(m interface{}) (err error) {
	msg, err := s.recv()
	if err != nil {
		return err
	}

	return json.Unmarshal(msg.data, m)
}

func (s *streamClient) recv() (m *streamMsg, err error) {
	userCancelled := func() bool {
		select {
		case <-s.userCtx.Done():
			return true
		default:
			return false
		}
	}

	select {
	case <-s.userCtx.Done():
		return m, s.userCtx.Err()
	case <-s.ctx.Done():
		// Give precedence for user cancellation.
		if userCancelled() {
			return m, s.userCtx.Err()
		}

		// Send all messages before returning error.
		select {
		case m = <-s.msgBuf.get():
		default:
			<-s.done
			return m, s.err
		}
	case m = <-s.msgBuf.get():
		// Give precedence for user cancellation.
		if userCancelled() {
			return m, s.userCtx.Err()
		}
	}

	// Preload the next message.
	s.msgBuf.load()

	return m, nil
}

// Close closes the stream client.
func (s *streamClient) close(err error) error {
	s.mu.Lock()
	remove := s.remove
	s.remove = nil
	s.mu.Unlock()

	if remove == nil {
		return errors.New("rpcc: the stream is already closed")
	}

	if err == nil {
		err = ErrStreamClosing
	}

	// Unsubscribe first to prevent incoming messages.
	remove()
	s.cancel()
	s.err = err
	close(s.done)

	return nil
}

// Close closes the stream client.
func (s *streamClient) Close() error {
	return s.close(nil)
}

// streamClients handles multiple instances of streamClient and
// enables sending of the same message to multiple clients.
type streamClients struct {
	mu      sync.Mutex
	seq     uint64
	clients map[uint64]*streamClient
}

func newStreamService() *streamClients {
	return &streamClients{
		clients: make(map[uint64]*streamClient),
	}
}

func (s *streamClients) add(client *streamClient) (seq uint64) {
	s.mu.Lock()
	seq = s.seq
	s.seq++
	s.clients[seq] = client
	s.mu.Unlock()
	return seq
}

func (s *streamClients) remove(seq uint64) {
	s.mu.Lock()
	delete(s.clients, seq)
	s.mu.Unlock()
}

func (s *streamClients) send(method string, args []byte) {
	s.mu.Lock()
	for _, client := range s.clients {
		client.msgBuf.store(&streamMsg{method: method, data: args})
	}
	s.mu.Unlock()
}
<commit_msg>rpcc: Copy raw bytes in RecvMsg when m is *[]byte<commit_after>package rpcc

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
)

var (
	// ErrStreamClosing indicates that the operation is illegal because
	// the stream is closing and there are no pending messages.
	ErrStreamClosing = errors.New("rpcc: the stream is closing")
)

type streamMsg struct {
	method string
	data   []byte
}

type messageBuffer struct {
	ch    chan *streamMsg
	mu    sync.Mutex
	queue []*streamMsg
}

func newMessageBuffer() *messageBuffer {
	return &messageBuffer{
		ch: make(chan *streamMsg, 1),
	}
}

func (b *messageBuffer) store(m *streamMsg) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) == 0 {
		select {
		case b.ch <- m:
			return
		default:
		}
	}
	b.queue = append(b.queue, m)
}

func (b *messageBuffer) load() {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) > 0 {
		select {
		case b.ch <- b.queue[0]:
			b.queue[0] = nil // Remove reference from underlying array.
			b.queue = b.queue[1:]
		default:
		}
	}
}

func (b *messageBuffer) get() <-chan *streamMsg {
	return b.ch
}

// Stream represents a stream of notifications for a certain method.
type Stream interface {
	// RecvMsg unmarshals pending messages onto m. Blocks until the
	// next message is received, context is canceled or stream is
	// closed.
	//
	// When m is a *[]byte the message will not be decoded and the
	// raw bytes are copied into m.
	RecvMsg(m interface{}) error
	// Close closes the stream and no new messages will be received.
	// RecvMsg will return ErrStreamClosing once all pending messages
	// have been received.
	Close() error
}

// NewStream creates a new stream that listens to notifications from the
// RPC server. This function is called by generated code.
func NewStream(ctx context.Context, method string, conn *Conn) (Stream, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	s := &streamClient{userCtx: ctx, done: make(chan struct{})}
	s.msgBuf = newMessageBuffer()
	s.ctx, s.cancel = context.WithCancel(context.Background())

	var err error
	s.remove, err = conn.listen(method, s)
	if err != nil {
		return nil, err
	}

	go func() {
		select {
		case <-s.ctx.Done():
		case <-conn.ctx.Done():
			s.close(ErrConnClosing)
		case <-ctx.Done():
			s.close(ctx.Err())
		}
	}()

	return s, nil
}

type streamClient struct {
	userCtx context.Context
	ctx     context.Context
	cancel  context.CancelFunc

	// msgBuf stores all incoming messages
	// until they are ready to be received.
	msgBuf *messageBuffer

	mu     sync.Mutex // Protects following.
	remove func()     // Unsubscribes from messages.

	done chan struct{} // Protects err.
	err  error
}

func (s *streamClient) RecvMsg(m interface{}) (err error) {
	msg, err := s.recv()
	if err != nil {
		return err
	}

	if m, ok := m.(*[]byte); ok {
		*m = append(*m, msg.data...)
		return nil
	}

	return json.Unmarshal(msg.data, m)
}

func (s *streamClient) recv() (m *streamMsg, err error) {
	userCancelled := func() bool {
		select {
		case <-s.userCtx.Done():
			return true
		default:
			return false
		}
	}

	select {
	case <-s.userCtx.Done():
		return m, s.userCtx.Err()
	case <-s.ctx.Done():
		// Give precedence for user cancellation.
		if userCancelled() {
			return m, s.userCtx.Err()
		}

		// Send all messages before returning error.
		select {
		case m = <-s.msgBuf.get():
		default:
			<-s.done
			return m, s.err
		}
	case m = <-s.msgBuf.get():
		// Give precedence for user cancellation.
		if userCancelled() {
			return m, s.userCtx.Err()
		}
	}

	// Preload the next message.
	s.msgBuf.load()

	return m, nil
}

// Close closes the stream client.
func (s *streamClient) close(err error) error {
	s.mu.Lock()
	remove := s.remove
	s.remove = nil
	s.mu.Unlock()

	if remove == nil {
		return errors.New("rpcc: the stream is already closed")
	}

	if err == nil {
		err = ErrStreamClosing
	}

	// Unsubscribe first to prevent incoming messages.
	remove()
	s.cancel()
	s.err = err
	close(s.done)

	return nil
}

// Close closes the stream client.
func (s *streamClient) Close() error {
	return s.close(nil)
}

// streamClients handles multiple instances of streamClient and
// enables sending of the same message to multiple clients.
type streamClients struct {
	mu      sync.Mutex
	seq     uint64
	clients map[uint64]*streamClient
}

func newStreamService() *streamClients {
	return &streamClients{
		clients: make(map[uint64]*streamClient),
	}
}

func (s *streamClients) add(client *streamClient) (seq uint64) {
	s.mu.Lock()
	seq = s.seq
	s.seq++
	s.clients[seq] = client
	s.mu.Unlock()
	return seq
}

func (s *streamClients) remove(seq uint64) {
	s.mu.Lock()
	delete(s.clients, seq)
	s.mu.Unlock()
}

func (s *streamClients) send(method string, args []byte) {
	s.mu.Lock()
	for _, client := range s.clients {
		client.msgBuf.store(&streamMsg{method: method, data: args})
	}
	s.mu.Unlock()
}
<|endoftext|>


<commit_before>package yum

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"time"
)

// global registry of known backends
var g_backends map[string]func(repo *Repository) (Backend, error)

// NewBackend returns a new backend of type "backend"
func NewBackend(backend string, repo *Repository) (Backend, error) {
	factory, ok := g_backends[backend]
	if !ok {
		return nil, fmt.Errorf("yum: no such backend [%s]", backend)
	}
	return factory(repo)
}

// Backend queries a YUM DB repository
type Backend interface {

	// YumDataType returns the ID for the data type as used in the repomd.xml file
	YumDataType() string

	// Download the DB from server
	GetLatestDB(url string) error

	// Check whether the DB is there
	HasDB() bool

	// Load loads the DB
	LoadDB() error

	// FindLatestMatchingName locats a package by name, returns the latest available version.
	FindLatestMatchingName(name, version, release string) (*Package, error)

	// FindLatestMatchingRequire locates a package providing a given functionality.
	FindLatestMatchingRequire(requirement string) (*Package, error)

	// GetPackages returns all the packages known by a YUM repository
	GetPackages() []*Package
}

// Repository represents a YUM repository with all associated metadata.
type Repository struct {
	Name           string
	RepoUrl        string
	RepoMdUrl      string
	LocalRepoMdXml string
	CacheDir       string
	Backends       []string
	Backend        Backend
}

// NewRepository create a new Repository with name and from url.
func NewRepository(name, url, cachedir string, backends []string, setupBackend, checkForUpdates bool) (*Repository, error) {
	repo := Repository{
		Name:           name,
		RepoUrl:        url,
		RepoMdUrl:      url + "/repodata/repomd.xml",
		LocalRepoMdXml: filepath.Join(cachedir, "repomd.xml"),
		CacheDir:       cachedir,
		Backends:       make([]string, len(backends)),
	}
	copy(repo.Backends, backends)

	err := os.MkdirAll(cachedir, 0644)
	if err != nil {
		return nil, err
	}

	// load appropriate backend if requested
	if setupBackend {
		if checkForUpdates {
			err = repo.setupBackendFromRemote()
			if err != nil {
				return nil, err
			}
		} else {
			err = repo.setupBackendFromLocal()
			if err != nil {
				return nil, err
			}
		}
	}
	return &repo, err
}

// FindLatestMatchingName locats a package by name, returns the latest available version.
func (repo *Repository) FindLatestMatchingName(name, version, release string) (*Package, error) {
	return repo.Backend.FindLatestMatchingName(name, version, release)
}

// FindLatestMatchingRequire locates a package providing a given functionality.
func (repo *Repository) FindLatestMatchingRequire(requirement string) (*Package, error) {
	return repo.Backend.FindLatestMatchingRequire(requirement)
}

// GetPackages returns all the packages known by a YUM repository
func (repo *Repository) GetPackages() []*Package {
	return repo.Backend.GetPackages()
}

// setupBackendFromRemote checks which backend should be used and updates the DB files.
func (repo *Repository) setupBackendFromRemote() error {
	var err error
	return err
}

func (repo *Repository) setupBackendFromLocal() error {
	var err error
	return err
}

// remoteMetadata retrieves the repo metadata file content
func (repo *Repository) remoteMetadata() ([]byte, error) {
	resp, err := http.Get(repo.RepoMdUrl)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	buf := new(bytes.Buffer)
	_, err = io.Copy(buf, resp.Body)
	if err != nil && err != io.EOF {
		return nil, err
	}
	return buf.Bytes(), err
}

// localMetadata retrieves the repo metadata from the repomd file
func (repo *Repository) localMetadata() ([]byte, error) {
	f, err := os.Open(repo.LocalRepoMdXml)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	buf := new(bytes.Buffer)
	_, err = io.Copy(buf, f)
	if err != nil && err != io.EOF {
		return nil, err
	}
	return buf.Bytes(), err
}

// checkRepoMD parses the Repository metadata XML content
func (repo *Repository) checkRepoMD(data []byte) (map[string]RepoMD, error) {
	db := make(map[string]RepoMD)
	var err error

	return db, err
}

type RepoMD struct {
	Checksum  string
	Timestamp time.Time
	Location  string
}

// EOF
<commit_msg>yum: initialize global registry of backends<commit_after>package yum

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"time"
)

// global registry of known backends
var g_backends = make(map[string]func(repo *Repository) (Backend, error))

// NewBackend returns a new backend of type "backend"
func NewBackend(backend string, repo *Repository) (Backend, error) {
	factory, ok := g_backends[backend]
	if !ok {
		return nil, fmt.Errorf("yum: no such backend [%s]", backend)
	}
	return factory(repo)
}

// Backend queries a YUM DB repository
type Backend interface {

	// YumDataType returns the ID for the data type as used in the repomd.xml file
	YumDataType() string

	// Download the DB from server
	GetLatestDB(url string) error

	// Check whether the DB is there
	HasDB() bool

	// Load loads the DB
	LoadDB() error

	// FindLatestMatchingName locats a package by name, returns the latest available version.
	FindLatestMatchingName(name, version, release string) (*Package, error)

	// FindLatestMatchingRequire locates a package providing a given functionality.
	FindLatestMatchingRequire(requirement string) (*Package, error)

	// GetPackages returns all the packages known by a YUM repository
	GetPackages() []*Package
}

// Repository represents a YUM repository with all associated metadata.
type Repository struct {
	Name           string
	RepoUrl        string
	RepoMdUrl      string
	LocalRepoMdXml string
	CacheDir       string
	Backends       []string
	Backend        Backend
}

// NewRepository create a new Repository with name and from url.
func NewRepository(name, url, cachedir string, backends []string, setupBackend, checkForUpdates bool) (*Repository, error) {
	repo := Repository{
		Name:           name,
		RepoUrl:        url,
		RepoMdUrl:      url + "/repodata/repomd.xml",
		LocalRepoMdXml: filepath.Join(cachedir, "repomd.xml"),
		CacheDir:       cachedir,
		Backends:       make([]string, len(backends)),
	}
	copy(repo.Backends, backends)

	err := os.MkdirAll(cachedir, 0644)
	if err != nil {
		return nil, err
	}

	// load appropriate backend if requested
	if setupBackend {
		if checkForUpdates {
			err = repo.setupBackendFromRemote()
			if err != nil {
				return nil, err
			}
		} else {
			err = repo.setupBackendFromLocal()
			if err != nil {
				return nil, err
			}
		}
	}
	return &repo, err
}

// FindLatestMatchingName locats a package by name, returns the latest available version.
func (repo *Repository) FindLatestMatchingName(name, version, release string) (*Package, error) {
	return repo.Backend.FindLatestMatchingName(name, version, release)
}

// FindLatestMatchingRequire locates a package providing a given functionality.
func (repo *Repository) FindLatestMatchingRequire(requirement string) (*Package, error) {
	return repo.Backend.FindLatestMatchingRequire(requirement)
}

// GetPackages returns all the packages known by a YUM repository
func (repo *Repository) GetPackages() []*Package {
	return repo.Backend.GetPackages()
}

// setupBackendFromRemote checks which backend should be used and updates the DB files.
func (repo *Repository) setupBackendFromRemote() error {
	var err error
	return err
}

func (repo *Repository) setupBackendFromLocal() error {
	var err error
	return err
}

// remoteMetadata retrieves the repo metadata file content
func (repo *Repository) remoteMetadata() ([]byte, error) {
	resp, err := http.Get(repo.RepoMdUrl)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	buf := new(bytes.Buffer)
	_, err = io.Copy(buf, resp.Body)
	if err != nil && err != io.EOF {
		return nil, err
	}
	return buf.Bytes(), err
}

// localMetadata retrieves the repo metadata from the repomd file
func (repo *Repository) localMetadata() ([]byte, error) {
	f, err := os.Open(repo.LocalRepoMdXml)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	buf := new(bytes.Buffer)
	_, err = io.Copy(buf, f)
	if err != nil && err != io.EOF {
		return nil, err
	}
	return buf.Bytes(), err
}

// checkRepoMD parses the Repository metadata XML content
func (repo *Repository) checkRepoMD(data []byte) (map[string]RepoMD, error) {
	db := make(map[string]RepoMD)
	var err error

	return db, err
}

type RepoMD struct {
	Checksum  string
	Timestamp time.Time
	Location  string
}

// EOF
<|endoftext|>


<commit_before>// +build freebsd

package mem

import (
	"os/exec"
	"strconv"
	"strings"

	common "github.com/shirou/gopsutil/common"
)

func VirtualMemory() (*VirtualMemoryStat, error) {
	pageSize, err := common.DoSysctrl("vm.stats.vm.v_page_size")
	if err != nil {
		return nil, err
	}
	p, err := strconv.ParseUint(pageSize[0], 10, 64)
	if err != nil {
		return nil, err
	}

	pageCount, err := common.DoSysctrl("vm.stats.vm.v_page_count")
	if err != nil {
		return nil, err
	}
	free, err := common.DoSysctrl("vm.stats.vm.v_free_count")
	if err != nil {
		return nil, err
	}
	active, err := common.DoSysctrl("vm.stats.vm.v_active_count")
	if err != nil {
		return nil, err
	}
	inactive, err := common.DoSysctrl("vm.stats.vm.v_inactive_count")
	if err != nil {
		return nil, err
	}
	cache, err := common.DoSysctrl("vm.stats.vm.v_cache_count")
	if err != nil {
		return nil, err
	}
	buffer, err := common.DoSysctrl("vfs.bufspace")
	if err != nil {
		return nil, err
	}
	wired, err := common.DoSysctrl("vm.stats.vm.v_wire_count")
	if err != nil {
		return nil, err
	}

	parsed := make([]uint64, 0, 7)
	vv := []string{
		pageCount[0],
		free[0],
		active[0],
		inactive[0],
		cache[0],
		buffer[0],
		wired[0],
	}
	for _, target := range vv {
		t, err := strconv.ParseUint(target, 10, 64)
		if err != nil {
			return nil, err
		}
		parsed = append(parsed, t)
	}

	ret := &VirtualMemoryStat{
		Total:    parsed[0] * p,
		Free:     parsed[1] * p,
		Active:   parsed[2] * p,
		Inactive: parsed[3] * p,
		Cached:   parsed[4] * p,
		Buffers:  parsed[5],
		Wired:    parsed[6] * p,
	}

	// TODO: platform independent (worked freebsd?)
	ret.Available = ret.Free + ret.Buffers + ret.Cached

	ret.Used = ret.Total - ret.Free
	ret.UsedPercent = float64(ret.Total-ret.Available) / float64(ret.Total) * 100.0

	return ret, nil
}

// Return swapinfo
// FreeBSD can have multiple swap devices. but use only first device
func SwapMemory() (*SwapMemoryStat, error) {
	out, err := exec.Command("swapinfo").Output()
	if err != nil {
		return nil, err
	}
	var ret *SwapMemoryStat
	for _, line := range strings.Split(string(out), "\n") {
		values := strings.Fields(line)
		// skip title line
		if len(values) == 0 || values[0] == "Device" {
			continue
		}

		u := strings.Replace(values[4], "%", "", 1)
		total_v, err := strconv.ParseUint(values[1], 10, 64)
		if err != nil {
			return nil, err
		}
		used_v, err := strconv.ParseUint(values[2], 10, 64)
		if err != nil {
			return nil, err
		}
		free_v, err := strconv.ParseUint(values[3], 10, 64)
		if err != nil {
			return nil, err
		}
		up_v, err := strconv.ParseFloat(u, 64)
		if err != nil {
			return nil, err
		}

		ret = &SwapMemoryStat{
			Total:       total_v,
			Used:        used_v,
			Free:        free_v,
			UsedPercent: up_v,
		}
	}

	return ret, nil
}
<commit_msg>mem[freebsd]: fix virtual memory available and used.<commit_after>// +build freebsd

package mem

import (
	"os/exec"
	"strconv"
	"strings"

	common "github.com/shirou/gopsutil/common"
)

func VirtualMemory() (*VirtualMemoryStat, error) {
	pageSize, err := common.DoSysctrl("vm.stats.vm.v_page_size")
	if err != nil {
		return nil, err
	}
	p, err := strconv.ParseUint(pageSize[0], 10, 64)
	if err != nil {
		return nil, err
	}

	pageCount, err := common.DoSysctrl("vm.stats.vm.v_page_count")
	if err != nil {
		return nil, err
	}
	free, err := common.DoSysctrl("vm.stats.vm.v_free_count")
	if err != nil {
		return nil, err
	}
	active, err := common.DoSysctrl("vm.stats.vm.v_active_count")
	if err != nil {
		return nil, err
	}
	inactive, err := common.DoSysctrl("vm.stats.vm.v_inactive_count")
	if err != nil {
		return nil, err
	}
	cache, err := common.DoSysctrl("vm.stats.vm.v_cache_count")
	if err != nil {
		return nil, err
	}
	buffer, err := common.DoSysctrl("vfs.bufspace")
	if err != nil {
		return nil, err
	}
	wired, err := common.DoSysctrl("vm.stats.vm.v_wire_count")
	if err != nil {
		return nil, err
	}

	parsed := make([]uint64, 0, 7)
	vv := []string{
		pageCount[0],
		free[0],
		active[0],
		inactive[0],
		cache[0],
		buffer[0],
		wired[0],
	}
	for _, target := range vv {
		t, err := strconv.ParseUint(target, 10, 64)
		if err != nil {
			return nil, err
		}
		parsed = append(parsed, t)
	}

	ret := &VirtualMemoryStat{
		Total:    parsed[0] * p,
		Free:     parsed[1] * p,
		Active:   parsed[2] * p,
		Inactive: parsed[3] * p,
		Cached:   parsed[4] * p,
		Buffers:  parsed[5],
		Wired:    parsed[6] * p,
	}

	ret.Available = ret.Inactive + ret.Cached + ret.Free
	ret.Used = ret.Active + ret.Wired + ret.Cached
	ret.UsedPercent = float64(ret.Total-ret.Available) / float64(ret.Total) * 100.0

	return ret, nil
}

// Return swapinfo
// FreeBSD can have multiple swap devices. but use only first device
func SwapMemory() (*SwapMemoryStat, error) {
	out, err := exec.Command("swapinfo").Output()
	if err != nil {
		return nil, err
	}
	var ret *SwapMemoryStat
	for _, line := range strings.Split(string(out), "\n") {
		values := strings.Fields(line)
		// skip title line
		if len(values) == 0 || values[0] == "Device" {
			continue
		}

		u := strings.Replace(values[4], "%", "", 1)
		total_v, err := strconv.ParseUint(values[1], 10, 64)
		if err != nil {
			return nil, err
		}
		used_v, err := strconv.ParseUint(values[2], 10, 64)
		if err != nil {
			return nil, err
		}
		free_v, err := strconv.ParseUint(values[3], 10, 64)
		if err != nil {
			return nil, err
		}
		up_v, err := strconv.ParseFloat(u, 64)
		if err != nil {
			return nil, err
		}

		ret = &SwapMemoryStat{
			Total:       total_v,
			Used:        used_v,
			Free:        free_v,
			UsedPercent: up_v,
		}
	}

	return ret, nil
}
<|endoftext|>


<commit_before>package main

import (
	mqttc "./mqttc"
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/grindhold/gominatim"
	geoipc "github.com/rubiojr/freegeoip-client"
	"gopkg.in/alecthomas/kingpin.v1"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Host struct {
	IP          string  `json:"ip"`
	Name        string  `json:"hostname"`
	Hop         int     `json:"hop-number"`
	Sent        int     `json:"sent"`
	LostPercent float64 `json:"lost-percent"`
	Last        float64 `json:"mean"`
	Avg         float64 `json:"mean"`
	Best        float64 `json:"best"`
	Worst       float64 `json:"worst"`
	StDev       float64 `json:"standard-dev"`
}

type Report struct {
	Time        time.Time       `json:"time"`
	Hosts       []*Host         `json:"hosts"`
	Hops        int             `json:"hops"`
	ElapsedTime time.Duration   `json:"elapsed_time"`
	Location    *ReportLocation `json:"location"`
}

// slightly simpler struct than the one provided by geoipc
type ReportLocation struct {
	IP          string  `json:"ip"`
	CountryCode string  `json:"country_code"`
	CountryName string  `json:"country_name"`
	City        string  `json:"city"`
	Latitude    float64 `json:"latitude"`
	Longitude   float64 `json:"longitude"`
}

func NewReport(reportCycles int, host string, args ...string) *Report {
	loc, err := geoipc.GetLocation()
	if err != nil {
		log.Panicf("Error getting location from geoip server: %s", err)
	}

	l := ReportLocation{
		CountryCode: loc.CountryCode,
		CountryName: loc.CountryName,
		Latitude:    loc.Latitude,
		Longitude:   loc.Longitude,
		IP:          loc.IP,
	}
	return NewReportWithLoc(reportCycles, host, &l, args...)
}

func NewReportWithLoc(reportCycles int, host string, loc *ReportLocation, args ...string) *Report {
	report := &Report{}
	report.Time = time.Now()
	args = append([]string{"--report", "-n", "-c", strconv.Itoa(reportCycles), host}, args...)

	tstart := time.Now()
	mtr := findMtrBin()
	rawOutput, err := exec.Command(mtr, args...).Output()

	if err != nil {
		panic("Error running the mtr command")
	}

	buf := bytes.NewBuffer(rawOutput)
	scanner := bufio.NewScanner(buf)
	scanner.Split(bufio.ScanLines)

	for scanner.Scan() {
		r, _ := regexp.Compile(`^\s+\d+\.`)

		line := scanner.Text()
		if !r.MatchString(line) {
			continue
		}

		tokens := strings.Fields(line)
		sent, err := strconv.Atoi(tokens[3])
		if err != nil {
			panic("Error parsing sent field")
		}

		host := Host{
			IP:   tokens[1],
			Sent: sent,
		}

		f2F(strings.Replace(tokens[2], "%", "", -1), &host.LostPercent)
		f2F(tokens[4], &host.Last)
		f2F(tokens[5], &host.Avg)
		f2F(tokens[6], &host.Best)
		f2F(tokens[7], &host.Worst)
		f2F(tokens[8], &host.StDev)

		report.Hosts = append(report.Hosts, &host)
	}

	report.Hops = len(report.Hosts)
	report.ElapsedTime = time.Since(tstart)
	report.Location = loc

	return report
}

func f2F(val string, field *float64) {
	f, err := strconv.ParseFloat(val, 64)
	*field = f
	if err != nil {
		panic("Error parsing field")
	}
}

func findMtrBin() string {
	paths := os.Getenv("PATH")
	if paths == "" {
		return ""
	}

	for _, path := range strings.Split(paths, ":") {
		if _, err := os.Stat(path + "/mtr"); err == nil {
			return path + "/mtr"
		}
	}

	return ""
}

func run(r *Report, count int, host string, stdout bool, args *mqttc.Args) error {

	var err error = nil
	if stdout {
		msg, _ := json.MarshalIndent(r, "", "  ")
		fmt.Println(string(msg))
	} else {
		msg, _ := json.Marshal(r)
		err = mqttc.PushMsg(string(msg), args)
	}

	return err
}

func parseBrokerUrls(brokerUrls string) []string {
	tokens := strings.Split(brokerUrls, ",")
	for i, url := range tokens {
		tokens[i] = strings.TrimSpace(url)
	}

	return tokens
}

func handleError(err error, fail bool) {
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error()+"\n")
	}

	if fail {
		os.Exit(1)
	}
}

func main() {
	kingpin.Version("0.2")

	count := kingpin.Flag("count", "Report cycles (mtr -c)").
		Default("10").Int()

	topic := kingpin.Flag("topic", "MTTQ topic").Default("/metrics/mtr").
		String()

	host := kingpin.Arg("host", "Target host").Required().String()

	repeat := kingpin.Flag("repeat", "Send the report every X seconds").
		Default("0").Int()

	brokerUrls := kingpin.Flag("broker-urls", "Comman separated MQTT broker URLs").
		Required().Default("").OverrideDefaultFromEnvar("MQTT_URLS").String()

	stdout := kingpin.Flag("stdout", "Print the report to stdout").
		Default("false").Bool()

	cafile := kingpin.Flag("cafile", "CA certificate when using TLS (optional)").
		String()

	location := kingpin.Flag("location", "Geocode the location of the server").
		String()

	insecure := kingpin.Flag("insecure", "Don't verify the server's certificate chain and host name.").
		Default("false").Bool()

	kingpin.Parse()

	if *cafile != "" {
		if _, err := os.Stat(*cafile); err != nil {
			fmt.Fprintf(os.Stderr, "Error reading CA certificate %s", err.Error())
			os.Exit(1)
		}
	}

	if findMtrBin() == "" {
		fmt.Fprintf(os.Stderr, "mtr binary not found in path\n")
		os.Exit(1)
	}

	var report *Report
	if *location != "" {
		loc, err := findLocation(*location)
		if err != nil {
			log.Fatalf("Geocoding of location %s failed: %s", *location, err)
			os.Exit(1)
		}
		report = NewReportWithLoc(*count, *host, loc)
	} else {
		report = NewReport(*count, *host)
	}

	urlList := parseBrokerUrls(*brokerUrls)

	var err error
	args := mqttc.Args{
		BrokerURLs:    urlList,
		ClientID:      "push-mtr",
		Topic:         *topic,
		TLSCACertPath: *cafile,
		TLSSkipVerify: *insecure,
	}

	if *repeat != 0 {
		timer := time.NewTicker(time.Duration(*repeat) * time.Second)
		for _ = range timer.C {
			err = run(report, *count, *host, *stdout, &args)
			handleError(err, false)
		}
	} else {
		err := run(report, *count, *host, *stdout, &args)
		handleError(err, true)
	}

}

func findLocation(query string) (*ReportLocation, error) {

	gominatim.SetServer("https://nominatim.openstreetmap.org/")

	iplocChan := make(chan geoipc.Location)
	go func() {
		// retry once
		for i := 0; i < 2; i++ {
			loc, err := geoipc.GetLocation()
			if err == nil {
				iplocChan <- loc
				break
			}
			time.Sleep(2 * time.Second)
		}
		close(iplocChan)
	}()

	gomiChan := make(chan []gominatim.SearchResult)
	go func(query string) {
		qry := gominatim.SearchQuery{
			Q:              query,
			Addressdetails: true,
			Limit:          1,
		}
		for i := 0; i < 2; i++ {
			r, err := qry.Get()
			if err == nil {
				gomiChan <- r
				break
			}
			time.Sleep(2 * time.Second)
		}
		close(gomiChan)
	}(query)

	iploc := <-iplocChan
	responses := <-gomiChan

	if (iploc == geoipc.Location{}) && len(responses) == 0 {
		return nil, fmt.Errorf("No locations found for '%s'", query)
	}

	if len(responses) > 0 {
		resp := responses[0]
		addr := resp.Address
		lat, _ := strconv.ParseFloat(resp.Lat, 64)
		lon, _ := strconv.ParseFloat(resp.Lon, 64)
		loc := ReportLocation{
			CountryCode: addr.CountryCode,
			CountryName: addr.Country,
			City:        addr.City,
			Latitude:    lat,
			Longitude:   lon,
			IP:          iploc.IP,
		}
		return &loc, nil
	} else {
		return nil, fmt.Errorf("Geocoding %s failed", query)
	}
}
<commit_msg>Bump version<commit_after>package main

import (
	mqttc "./mqttc"
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/grindhold/gominatim"
	geoipc "github.com/rubiojr/freegeoip-client"
	"gopkg.in/alecthomas/kingpin.v1"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Host struct {
	IP          string  `json:"ip"`
	Name        string  `json:"hostname"`
	Hop         int     `json:"hop-number"`
	Sent        int     `json:"sent"`
	LostPercent float64 `json:"lost-percent"`
	Last        float64 `json:"mean"`
	Avg         float64 `json:"mean"`
	Best        float64 `json:"best"`
	Worst       float64 `json:"worst"`
	StDev       float64 `json:"standard-dev"`
}

type Report struct {
	Time        time.Time       `json:"time"`
	Hosts       []*Host         `json:"hosts"`
	Hops        int             `json:"hops"`
	ElapsedTime time.Duration   `json:"elapsed_time"`
	Location    *ReportLocation `json:"location"`
}

// slightly simpler struct than the one provided by geoipc
type ReportLocation struct {
	IP          string  `json:"ip"`
	CountryCode string  `json:"country_code"`
	CountryName string  `json:"country_name"`
	City        string  `json:"city"`
	Latitude    float64 `json:"latitude"`
	Longitude   float64 `json:"longitude"`
}

func NewReport(reportCycles int, host string, args ...string) *Report {
	loc, err := geoipc.GetLocation()
	if err != nil {
		log.Panicf("Error getting location from geoip server: %s", err)
	}

	l := ReportLocation{
		CountryCode: loc.CountryCode,
		CountryName: loc.CountryName,
		Latitude:    loc.Latitude,
		Longitude:   loc.Longitude,
		IP:          loc.IP,
	}
	return NewReportWithLoc(reportCycles, host, &l, args...)
}

func NewReportWithLoc(reportCycles int, host string, loc *ReportLocation, args ...string) *Report {
	report := &Report{}
	report.Time = time.Now()
	args = append([]string{"--report", "-n", "-c", strconv.Itoa(reportCycles), host}, args...)

	tstart := time.Now()
	mtr := findMtrBin()
	rawOutput, err := exec.Command(mtr, args...).Output()

	if err != nil {
		panic("Error running the mtr command")
	}

	buf := bytes.NewBuffer(rawOutput)
	scanner := bufio.NewScanner(buf)
	scanner.Split(bufio.ScanLines)

	for scanner.Scan() {
		r, _ := regexp.Compile(`^\s+\d+\.`)

		line := scanner.Text()
		if !r.MatchString(line) {
			continue
		}

		tokens := strings.Fields(line)
		sent, err := strconv.Atoi(tokens[3])
		if err != nil {
			panic("Error parsing sent field")
		}

		host := Host{
			IP:   tokens[1],
			Sent: sent,
		}

		f2F(strings.Replace(tokens[2], "%", "", -1), &host.LostPercent)
		f2F(tokens[4], &host.Last)
		f2F(tokens[5], &host.Avg)
		f2F(tokens[6], &host.Best)
		f2F(tokens[7], &host.Worst)
		f2F(tokens[8], &host.StDev)

		report.Hosts = append(report.Hosts, &host)
	}

	report.Hops = len(report.Hosts)
	report.ElapsedTime = time.Since(tstart)
	report.Location = loc

	return report
}

func f2F(val string, field *float64) {
	f, err := strconv.ParseFloat(val, 64)
	*field = f
	if err != nil {
		panic("Error parsing field")
	}
}

func findMtrBin() string {
	paths := os.Getenv("PATH")
	if paths == "" {
		return ""
	}

	for _, path := range strings.Split(paths, ":") {
		if _, err := os.Stat(path + "/mtr"); err == nil {
			return path + "/mtr"
		}
	}

	return ""
}

func run(r *Report, count int, host string, stdout bool, args *mqttc.Args) error {

	var err error = nil
	if stdout {
		msg, _ := json.MarshalIndent(r, "", "  ")
		fmt.Println(string(msg))
	} else {
		msg, _ := json.Marshal(r)
		err = mqttc.PushMsg(string(msg), args)
	}

	return err
}

func parseBrokerUrls(brokerUrls string) []string {
	tokens := strings.Split(brokerUrls, ",")
	for i, url := range tokens {
		tokens[i] = strings.TrimSpace(url)
	}

	return tokens
}

func handleError(err error, fail bool) {
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error()+"\n")
	}

	if fail {
		os.Exit(1)
	}
}

func main() {
	kingpin.Version("0.2.1")

	count := kingpin.Flag("count", "Report cycles (mtr -c)").
		Default("10").Int()

	topic := kingpin.Flag("topic", "MTTQ topic").Default("/metrics/mtr").
		String()

	host := kingpin.Arg("host", "Target host").Required().String()

	repeat := kingpin.Flag("repeat", "Send the report every X seconds").
		Default("0").Int()

	brokerUrls := kingpin.Flag("broker-urls", "Comman separated MQTT broker URLs").
		Required().Default("").OverrideDefaultFromEnvar("MQTT_URLS").String()

	stdout := kingpin.Flag("stdout", "Print the report to stdout").
		Default("false").Bool()

	cafile := kingpin.Flag("cafile", "CA certificate when using TLS (optional)").
		String()

	location := kingpin.Flag("location", "Geocode the location of the server").
		String()

	insecure := kingpin.Flag("insecure", "Don't verify the server's certificate chain and host name.").
		Default("false").Bool()

	kingpin.Parse()

	if *cafile != "" {
		if _, err := os.Stat(*cafile); err != nil {
			fmt.Fprintf(os.Stderr, "Error reading CA certificate %s", err.Error())
			os.Exit(1)
		}
	}

	if findMtrBin() == "" {
		fmt.Fprintf(os.Stderr, "mtr binary not found in path\n")
		os.Exit(1)
	}

	var report *Report
	if *location != "" {
		loc, err := findLocation(*location)
		if err != nil {
			log.Fatalf("Geocoding of location %s failed: %s", *location, err)
			os.Exit(1)
		}
		report = NewReportWithLoc(*count, *host, loc)
	} else {
		report = NewReport(*count, *host)
	}

	urlList := parseBrokerUrls(*brokerUrls)

	var err error
	args := mqttc.Args{
		BrokerURLs:    urlList,
		ClientID:      "push-mtr",
		Topic:         *topic,
		TLSCACertPath: *cafile,
		TLSSkipVerify: *insecure,
	}

	if *repeat != 0 {
		timer := time.NewTicker(time.Duration(*repeat) * time.Second)
		for _ = range timer.C {
			err = run(report, *count, *host, *stdout, &args)
			handleError(err, false)
		}
	} else {
		err := run(report, *count, *host, *stdout, &args)
		handleError(err, true)
	}

}

func findLocation(query string) (*ReportLocation, error) {

	gominatim.SetServer("https://nominatim.openstreetmap.org/")

	iplocChan := make(chan geoipc.Location)
	go func() {
		// retry once
		for i := 0; i < 2; i++ {
			loc, err := geoipc.GetLocation()
			if err == nil {
				iplocChan <- loc
				break
			}
			time.Sleep(2 * time.Second)
		}
		close(iplocChan)
	}()

	gomiChan := make(chan []gominatim.SearchResult)
	go func(query string) {
		qry := gominatim.SearchQuery{
			Q:              query,
			Addressdetails: true,
			Limit:          1,
		}
		for i := 0; i < 2; i++ {
			r, err := qry.Get()
			if err == nil {
				gomiChan <- r
				break
			}
			time.Sleep(2 * time.Second)
		}
		close(gomiChan)
	}(query)

	iploc := <-iplocChan
	responses := <-gomiChan

	if (iploc == geoipc.Location{}) && len(responses) == 0 {
		return nil, fmt.Errorf("No locations found for '%s'", query)
	}

	if len(responses) > 0 {
		resp := responses[0]
		addr := resp.Address
		lat, _ := strconv.ParseFloat(resp.Lat, 64)
		lon, _ := strconv.ParseFloat(resp.Lon, 64)
		loc := ReportLocation{
			CountryCode: addr.CountryCode,
			CountryName: addr.Country,
			City:        addr.City,
			Latitude:    lat,
			Longitude:   lon,
			IP:          iploc.IP,
		}
		return &loc, nil
	} else {
		return nil, fmt.Errorf("Geocoding %s failed", query)
	}
}
<|endoftext|>