<commit_before>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"sync"
	"testing"
	"time"

	"github.com/letsencrypt/boulder/core"
	"github.com/letsencrypt/boulder/sa"
	"github.com/letsencrypt/boulder/test"
)

var dbConnStr = "mysql+tcp://boulder@localhost:3306/boulder_test"

func BenchmarkCheckCert(b *testing.B) {
	dbMap, err := sa.NewDbMap(dbConnStr)
	if err != nil {
		fmt.Println("Couldn't connect to database")
		return
	}

	checker := newChecker(dbMap)
	testKey, _ := rsa.GenerateKey(rand.Reader, 1024)
	expiry := time.Now().AddDate(0, 0, 1)
	serial := big.NewInt(1337)
	rawCert := x509.Certificate{
		Subject: pkix.Name{
			CommonName: "example.com",
		},
		NotAfter:     expiry,
		DNSNames:     []string{"example-a.com"},
		SerialNumber: serial,
	}
	certDer, _ := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
	cert := core.Certificate{
		Status:  core.StatusValid,
		Serial:  core.SerialToString(serial),
		Digest:  core.Fingerprint256(certDer),
		DER:     certDer,
		Issued:  time.Now(),
		Expires: expiry,
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		checker.checkCert(cert)
	}
}

func TestCheckCert(t *testing.T) {
	testKey, _ := rsa.GenerateKey(rand.Reader, 1024)
	now := time.Now()
	expiry := now.Add(checkPeriod)
	serial := big.NewInt(1337)
	// Problems
	//   Blacklsited common name
	//   Expiry period is too long
	//   Basic Constraints aren't set
	//   Wrong key usage (none)
	rawCert := x509.Certificate{
		Subject: pkix.Name{
			CommonName: "example.com",
		},
		NotAfter:              expiry.AddDate(0, 0, 1),
		DNSNames:              []string{"example-a.com"},
		SerialNumber:          serial,
		BasicConstraintsValid: false,
	}
	brokenCertDer, err := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
	test.AssertNotError(t, err, "Couldn't create certificate")
	// Problems
	//   Digest doesn't match
	//   Serial doesn't match
	//   Expiry doesn't match
	cert := core.Certificate{
		Status:  core.StatusValid,
		DER:     brokenCertDer,
		Issued:  now,
		Expires: expiry.AddDate(0, 0, 2),
	}

	checker := newChecker(nil)

	problems := checker.checkCert(cert)
	test.AssertEquals(t, len(problems), 7)

	// Fix the problems
	rawCert.Subject.CommonName = "example-a.com"
	rawCert.NotAfter = expiry
	rawCert.BasicConstraintsValid = true
	rawCert.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth}
	goodCertDer, err := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
	test.AssertNotError(t, err, "Couldn't create certificate")
	parsed, err := x509.ParseCertificate(goodCertDer)
	test.AssertNotError(t, err, "Couldn't parse created certificate")
	cert.Serial = core.SerialToString(serial)
	cert.Digest = core.Fingerprint256(goodCertDer)
	cert.DER = goodCertDer
	cert.Expires = parsed.NotAfter
	problems = checker.checkCert(cert)
	test.AssertEquals(t, len(problems), 0)
}

func TestGetAndProcessCerts(t *testing.T) {
	dbMap, err := sa.NewDbMap(dbConnStr)
	test.AssertNotError(t, err, "Couldn't connect to database")
	checker := newChecker(dbMap)
	sa, err := sa.NewSQLStorageAuthority(dbMap)
	test.AssertNotError(t, err, "Couldn't create SA to insert certificates")
	defer func() {
		dbMap.TruncateTables()
		test.AssertNotError(t, err, "Failed to truncate tables")
	}()

	testKey, _ := rsa.GenerateKey(rand.Reader, 1024)
	// Problems
	//   Blacklsited common name
	//   Expiry period is too long
	//   Basic Constraints aren't set
	//   Wrong key usage (none)
	rawCert := x509.Certificate{
		Subject: pkix.Name{
			CommonName: "not-blacklisted.com",
		},
		BasicConstraintsValid: true,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
	}
	for i := int64(0); i < 5; i++ {
		rawCert.SerialNumber = big.NewInt(i)
		certDER, err := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
		test.AssertNotError(t, err, "Couldn't create certificate")
		_, err = sa.AddCertificate(certDER, 0)
		test.AssertNotError(t, err, "Couldn't add certificate")
	}

	err = checker.getCerts()
	test.AssertNotError(t, err, "Failed to retrieve certificates")
	test.AssertEquals(t, len(checker.certs), 5)
	wg := new(sync.WaitGroup)
	wg.Add(1)
	checker.processCerts(wg)
	test.AssertEquals(t, checker.issuedReport.BadCerts, int64(5))
	test.AssertEquals(t, len(checker.issuedReport.Entries), 5)
}
<commit_msg>DB typo<commit_after>// Copyright 2015 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"sync"
	"testing"
	"time"

	"github.com/letsencrypt/boulder/core"
	"github.com/letsencrypt/boulder/sa"
	"github.com/letsencrypt/boulder/test"
)

var dbConnStr = "mysql+tcp://boulder@localhost:3306/boulder_sa_test"

func BenchmarkCheckCert(b *testing.B) {
	dbMap, err := sa.NewDbMap(dbConnStr)
	if err != nil {
		fmt.Println("Couldn't connect to database")
		return
	}

	checker := newChecker(dbMap)
	testKey, _ := rsa.GenerateKey(rand.Reader, 1024)
	expiry := time.Now().AddDate(0, 0, 1)
	serial := big.NewInt(1337)
	rawCert := x509.Certificate{
		Subject: pkix.Name{
			CommonName: "example.com",
		},
		NotAfter:     expiry,
		DNSNames:     []string{"example-a.com"},
		SerialNumber: serial,
	}
	certDer, _ := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
	cert := core.Certificate{
		Status:  core.StatusValid,
		Serial:  core.SerialToString(serial),
		Digest:  core.Fingerprint256(certDer),
		DER:     certDer,
		Issued:  time.Now(),
		Expires: expiry,
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		checker.checkCert(cert)
	}
}

func TestCheckCert(t *testing.T) {
	testKey, _ := rsa.GenerateKey(rand.Reader, 1024)
	now := time.Now()
	expiry := now.Add(checkPeriod)
	serial := big.NewInt(1337)
	// Problems
	//   Blacklsited common name
	//   Expiry period is too long
	//   Basic Constraints aren't set
	//   Wrong key usage (none)
	rawCert := x509.Certificate{
		Subject: pkix.Name{
			CommonName: "example.com",
		},
		NotAfter:              expiry.AddDate(0, 0, 1),
		DNSNames:              []string{"example-a.com"},
		SerialNumber:          serial,
		BasicConstraintsValid: false,
	}
	brokenCertDer, err := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
	test.AssertNotError(t, err, "Couldn't create certificate")
	// Problems
	//   Digest doesn't match
	//   Serial doesn't match
	//   Expiry doesn't match
	cert := core.Certificate{
		Status:  core.StatusValid,
		DER:     brokenCertDer,
		Issued:  now,
		Expires: expiry.AddDate(0, 0, 2),
	}

	checker := newChecker(nil)

	problems := checker.checkCert(cert)
	test.AssertEquals(t, len(problems), 7)

	// Fix the problems
	rawCert.Subject.CommonName = "example-a.com"
	rawCert.NotAfter = expiry
	rawCert.BasicConstraintsValid = true
	rawCert.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth}
	goodCertDer, err := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
	test.AssertNotError(t, err, "Couldn't create certificate")
	parsed, err := x509.ParseCertificate(goodCertDer)
	test.AssertNotError(t, err, "Couldn't parse created certificate")
	cert.Serial = core.SerialToString(serial)
	cert.Digest = core.Fingerprint256(goodCertDer)
	cert.DER = goodCertDer
	cert.Expires = parsed.NotAfter
	problems = checker.checkCert(cert)
	test.AssertEquals(t, len(problems), 0)
}

func TestGetAndProcessCerts(t *testing.T) {
	dbMap, err := sa.NewDbMap(dbConnStr)
	test.AssertNotError(t, err, "Couldn't connect to database")
	checker := newChecker(dbMap)
	sa, err := sa.NewSQLStorageAuthority(dbMap)
	test.AssertNotError(t, err, "Couldn't create SA to insert certificates")
	defer func() {
		dbMap.TruncateTables()
		test.AssertNotError(t, err, "Failed to truncate tables")
	}()

	testKey, _ := rsa.GenerateKey(rand.Reader, 1024)
	// Problems
	//   Blacklsited common name
	//   Expiry period is too long
	//   Basic Constraints aren't set
	//   Wrong key usage (none)
	rawCert := x509.Certificate{
		Subject: pkix.Name{
			CommonName: "not-blacklisted.com",
		},
		BasicConstraintsValid: true,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
	}
	for i := int64(0); i < 5; i++ {
		rawCert.SerialNumber = big.NewInt(i)
		certDER, err := x509.CreateCertificate(rand.Reader, &rawCert, &rawCert, &testKey.PublicKey, testKey)
		test.AssertNotError(t, err, "Couldn't create certificate")
		_, err = sa.AddCertificate(certDER, 0)
		test.AssertNotError(t, err, "Couldn't add certificate")
	}

	err = checker.getCerts()
	test.AssertNotError(t, err, "Failed to retrieve certificates")
	test.AssertEquals(t, len(checker.certs), 5)
	wg := new(sync.WaitGroup)
	wg.Add(1)
	checker.processCerts(wg)
	test.AssertEquals(t, checker.issuedReport.BadCerts, int64(5))
	test.AssertEquals(t, len(checker.issuedReport.Entries), 5)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"
	"path/filepath"

	"github.com/spf13/cobra"

	"github.com/NebulousLabs/Sia/api"
	"github.com/NebulousLabs/Sia/modules"
)

// filesize returns a string that displays a filesize in human-readable units.
func filesizeUnits(size int64) string {
	if size == 0 {
		return "0 B"
	}
	sizes := []string{"B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}
	i := int(math.Log10(float64(size)) / 3)
	return fmt.Sprintf("%.*f %s", i, float64(size)/math.Pow10(3*i), sizes[i])
}

var (
	renterCmd = &cobra.Command{
		Use:   "renter",
		Short: "Perform renter actions",
		Long:  "Upload, download, rename, delete, load, or share files.",
		Run:   wrap(renterfileslistcmd),
	}

	renterUploadsCmd = &cobra.Command{
		Use:   "uploads",
		Short: "View the upload queue",
		Long:  "View the list of files currently uploading.",
		Run:   wrap(renteruploadscmd),
	}

	renterDownloadsCmd = &cobra.Command{
		Use:   "downloads",
		Short: "View the download queue",
		Long:  "View the list of files currently downloading.",
		Run:   wrap(renterdownloadscmd),
	}

	renterFilesDeleteCmd = &cobra.Command{
		Use:   "delete [path]",
		Short: "Delete a file",
		Long:  "Delete a file. Does not delete the file on disk.",
		Run:   wrap(renterfilesdeletecmd),
	}

	renterFilesDownloadCmd = &cobra.Command{
		Use:   "download [path] [destination]",
		Short: "Download a file",
		Long:  "Download a previously-uploaded file to a specified destination.",
		Run:   wrap(renterfilesdownloadcmd),
	}

	renterFilesListCmd = &cobra.Command{
		Use:   "list",
		Short: "List the status of all files",
		Long:  "List the status of all files known to the renter on the Sia network.",
		Run:   wrap(renterfileslistcmd),
	}

	renterFilesLoadCmd = &cobra.Command{
		Use:   "load [source]",
		Short: "Load a .sia file",
		Long:  "Load a .sia file, adding the file entries contained within.",
		Run:   wrap(renterfilesloadcmd),
	}

	renterFilesLoadASCIICmd = &cobra.Command{
		Use:   "loadascii [ascii]",
		Short: "Load an ASCII-encoded .sia file",
		Long:  "Load an ASCII-encoded .sia file.",
		Run:   wrap(renterfilesloadasciicmd),
	}

	renterFilesRenameCmd = &cobra.Command{
		Use:   "rename [path] [newpath]",
		Short: "Rename a file",
		Long:  "Rename a file.",
		Run:   wrap(renterfilesrenamecmd),
	}

	renterFilesShareCmd = &cobra.Command{
		Use:   "share [path] [destination]",
		Short: "Export a file to a .sia for sharing",
		Long:  "Export a file to a .sia for sharing.",
		Run:   wrap(renterfilessharecmd),
	}

	renterFilesShareASCIICmd = &cobra.Command{
		Use:   "shareascii [path]",
		Short: "Export a file as an ASCII-encoded .sia file",
		Long:  "Export a file as an ASCII-encoded .sia file.",
		Run:   wrap(renterfilesshareasciicmd),
	}

	renterFilesUploadCmd = &cobra.Command{
		Use:   "upload [source] [path]",
		Short: "Upload a file",
		Long:  "Upload a file using a given nickname.",
		Run:   wrap(renterfilesuploadcmd),
	}
)

// abs returns the absolute representation of a path.
// TODO: bad things can happen if you run siac from a non-existent directory.
// Implement some checks to catch this problem.
func abs(path string) string {
	abspath, err := filepath.Abs(path)
	if err != nil {
		return path
	}
	return abspath
}

func renteruploadscmd() {
	var rf api.RenterFiles
	err := getAPI("/renter/files", &rf)
	if err != nil {
		fmt.Println("Could not get upload queue:", err)
		return
	}

	// TODO: add a --history flag to the uploads command to mirror the --history
	//       flag in the downloads command. This hasn't been done yet because the
	//       call to /renter/files includes files that have been shared with you,
	//       not just files you've uploaded.

	// Filter out files that have been uploaded.
	var filteredFiles []modules.FileInfo
	for _, fi := range rf.Files {
		if !fi.Available {
			filteredFiles = append(filteredFiles, fi)
		}
	}
	if len(filteredFiles) == 0 {
		fmt.Println("No files are uploading.")
		return
	}
	fmt.Println("Uploading", len(filteredFiles), "files:")
	for _, file := range filteredFiles {
		fmt.Printf("%13s  %s (uploading, %0.2f%%)\n", filesizeUnits(int64(file.Filesize)), file.SiaPath, file.UploadProgress)
	}
}

func renterdownloadscmd() {
	var queue api.RenterDownloadQueue
	err := getAPI("/renter/downloads", &queue)
	if err != nil {
		fmt.Println("Could not get download queue:", err)
		return
	}
	// Filter out files that have been downloaded.
	var downloading []modules.DownloadInfo
	for _, file := range queue.Downloads {
		if file.Received != file.Filesize {
			downloading = append(downloading, file)
		}
	}
	if len(downloading) == 0 {
		fmt.Println("No files are downloading.")
	} else {
		fmt.Println("Downloading", len(downloading), "files:")
		for _, file := range downloading {
			fmt.Printf("%s: %5.1f%% %s -> %s\n", file.StartTime.Format("Jan 02 03:04 PM"), 100*float64(file.Received)/float64(file.Filesize), file.SiaPath, file.Destination)
		}
	}
	if !renterShowHistory {
		return
	}
	fmt.Println()
	// Filter out files that are downloading.
	var downloaded []modules.DownloadInfo
	for _, file := range queue.Downloads {
		if file.Received == file.Filesize {
			downloaded = append(downloaded, file)
		}
	}
	if len(downloaded) == 0 {
		fmt.Println("No files downloaded.")
	} else {
		fmt.Println("Downloaded", len(downloaded), "files:")
		for _, file := range downloaded {
			fmt.Printf("%s: %s -> %s\n", file.StartTime.Format("Jan 02 03:04 PM"), file.SiaPath, file.Destination)
		}
	}
}

func renterfilesdeletecmd(path string) {
	err := post("/renter/delete/"+path, "")
	if err != nil {
		fmt.Println("Could not delete file:", err)
		return
	}
	fmt.Println("Deleted", path)
}

func renterfilesdownloadcmd(path, destination string) {
	err := get("/renter/download/" + path + "?destination=" + abs(destination))
	if err != nil {
		fmt.Println("Could not download file:", err)
		return
	}
	fmt.Printf("Downloaded '%s' to %s.\n", path, abs(destination))
}

func renterfileslistcmd() {
	var rf api.RenterFiles
	err := getAPI("/renter/files", &rf)
	if err != nil {
		fmt.Println("Could not get file list:", err)
		return
	}
	if len(rf.Files) == 0 {
		fmt.Println("No files have been uploaded.")
		return
	}
	fmt.Println("Tracking", len(rf.Files), "files:")
	for _, file := range rf.Files {
		if file.Available {
			fmt.Printf("%13s  %s\n", filesizeUnits(int64(file.Filesize)), file.SiaPath)
		} else {
			fmt.Printf("%13s  %s (uploading, %0.2f%%)\n", filesizeUnits(int64(file.Filesize)), file.SiaPath, file.UploadProgress)
		}
	}
}

func renterfilesloadcmd(source string) {
	var info api.RenterLoad
	err := postResp("/renter/load", "source="+abs(source), &info)
	if err != nil {
		fmt.Println("Could not load file:", err)
		return
	}
	fmt.Printf("Loaded %d file(s):\n", len(info.FilesAdded))
	for _, file := range info.FilesAdded {
		fmt.Printf("\t%s\n", file)
	}
}

func renterfilesloadasciicmd(ascii string) {
	var info api.RenterLoad
	err := postResp("/renter/loadascii", "asciisia="+ascii, &info)
	if err != nil {
		fmt.Println("Could not load file:", err)
		return
	}
	fmt.Printf("Loaded %d file(s):\n", len(info.FilesAdded))
	for _, file := range info.FilesAdded {
		fmt.Printf("\t%s\n", file)
	}
}

func renterfilesrenamecmd(path, newpath string) {
	err := post("/renter/rename/"+path, "newsiapath="+newpath)
	if err != nil {
		fmt.Println("Could not rename file:", err)
		return
	}
	fmt.Printf("Renamed %s to %s\n", path, newpath)
}

func renterfilessharecmd(path, destination string) {
	err := get(fmt.Sprintf("/renter/share?siapaths=%s&destination=%s", path, abs(destination)))
	if err != nil {
		fmt.Println("Could not share file:", err)
		return
	}
	fmt.Printf("Exported %s to %s\n", path, abs(destination))
}

func renterfilesshareasciicmd(path string) {
	var data api.RenterShareASCII
	err := getAPI("/renter/shareascii?siapaths="+path, &data)
	if err != nil {
		fmt.Println("Could not share file:", err)
		return
	}
	fmt.Println(data.ASCIIsia)
}

func renterfilesuploadcmd(source, path string) {
	err := post("/renter/upload/"+path, "source="+abs(source))
	if err != nil {
		fmt.Println("Could not upload file:", err)
		return
	}
	fmt.Printf("Uploaded '%s' as %s.\n", abs(source), path)
}
<commit_msg>Add docstrings to renter command handlers<commit_after>package main

import (
	"fmt"
	"math"
	"path/filepath"

	"github.com/spf13/cobra"

	"github.com/NebulousLabs/Sia/api"
	"github.com/NebulousLabs/Sia/modules"
)

// filesize returns a string that displays a filesize in human-readable units.
func filesizeUnits(size int64) string {
	if size == 0 {
		return "0 B"
	}
	sizes := []string{"B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}
	i := int(math.Log10(float64(size)) / 3)
	return fmt.Sprintf("%.*f %s", i, float64(size)/math.Pow10(3*i), sizes[i])
}

var (
	renterCmd = &cobra.Command{
		Use:   "renter",
		Short: "Perform renter actions",
		Long:  "Upload, download, rename, delete, load, or share files.",
		Run:   wrap(renterfileslistcmd),
	}

	renterUploadsCmd = &cobra.Command{
		Use:   "uploads",
		Short: "View the upload queue",
		Long:  "View the list of files currently uploading.",
		Run:   wrap(renteruploadscmd),
	}

	renterDownloadsCmd = &cobra.Command{
		Use:   "downloads",
		Short: "View the download queue",
		Long:  "View the list of files currently downloading.",
		Run:   wrap(renterdownloadscmd),
	}

	renterFilesDeleteCmd = &cobra.Command{
		Use:   "delete [path]",
		Short: "Delete a file",
		Long:  "Delete a file. Does not delete the file on disk.",
		Run:   wrap(renterfilesdeletecmd),
	}

	renterFilesDownloadCmd = &cobra.Command{
		Use:   "download [path] [destination]",
		Short: "Download a file",
		Long:  "Download a previously-uploaded file to a specified destination.",
		Run:   wrap(renterfilesdownloadcmd),
	}

	renterFilesListCmd = &cobra.Command{
		Use:   "list",
		Short: "List the status of all files",
		Long:  "List the status of all files known to the renter on the Sia network.",
		Run:   wrap(renterfileslistcmd),
	}

	renterFilesLoadCmd = &cobra.Command{
		Use:   "load [source]",
		Short: "Load a .sia file",
		Long:  "Load a .sia file, adding the file entries contained within.",
		Run:   wrap(renterfilesloadcmd),
	}

	renterFilesLoadASCIICmd = &cobra.Command{
		Use:   "loadascii [ascii]",
		Short: "Load an ASCII-encoded .sia file",
		Long:  "Load an ASCII-encoded .sia file.",
		Run:   wrap(renterfilesloadasciicmd),
	}

	renterFilesRenameCmd = &cobra.Command{
		Use:   "rename [path] [newpath]",
		Short: "Rename a file",
		Long:  "Rename a file.",
		Run:   wrap(renterfilesrenamecmd),
	}

	renterFilesShareCmd = &cobra.Command{
		Use:   "share [path] [destination]",
		Short: "Export a file to a .sia for sharing",
		Long:  "Export a file to a .sia for sharing.",
		Run:   wrap(renterfilessharecmd),
	}

	renterFilesShareASCIICmd = &cobra.Command{
		Use:   "shareascii [path]",
		Short: "Prints an ASCII-encoded .sia file for sharing",
		Long:  "Prints an ASCII-encoded .sia file for sharing, but does not save the .sia file to disk.",
		Run:   wrap(renterfilesshareasciicmd),
	}

	renterFilesUploadCmd = &cobra.Command{
		Use:   "upload [source] [path]",
		Short: "Upload a file",
		Long:  "Upload a file to [path] on the Sia network.",
		Run:   wrap(renterfilesuploadcmd),
	}
)

// abs returns the absolute representation of a path.
// TODO: bad things can happen if you run siac from a non-existent directory.
// Implement some checks to catch this problem.
func abs(path string) string {
	abspath, err := filepath.Abs(path)
	if err != nil {
		return path
	}
	return abspath
}

// renteruploadscmd is the handler for the command `siac renter uploads`.
// Lists files currently uploading.
func renteruploadscmd() {
	var rf api.RenterFiles
	err := getAPI("/renter/files", &rf)
	if err != nil {
		fmt.Println("Could not get upload queue:", err)
		return
	}

	// TODO: add a --history flag to the uploads command to mirror the --history
	//       flag in the downloads command. This hasn't been done yet because the
	//       call to /renter/files includes files that have been shared with you,
	//       not just files you've uploaded.

	// Filter out files that have been uploaded.
	var filteredFiles []modules.FileInfo
	for _, fi := range rf.Files {
		if !fi.Available {
			filteredFiles = append(filteredFiles, fi)
		}
	}
	if len(filteredFiles) == 0 {
		fmt.Println("No files are uploading.")
		return
	}
	fmt.Println("Uploading", len(filteredFiles), "files:")
	for _, file := range filteredFiles {
		fmt.Printf("%13s  %s (uploading, %0.2f%%)\n", filesizeUnits(int64(file.Filesize)), file.SiaPath, file.UploadProgress)
	}
}

// renterdownloadscmd is the handler for the command `siac renter downloads`.
// Lists files currently downloading, and optionally previously downloaded
// files if the -H or --history flag is specified.
func renterdownloadscmd() {
	var queue api.RenterDownloadQueue
	err := getAPI("/renter/downloads", &queue)
	if err != nil {
		fmt.Println("Could not get download queue:", err)
		return
	}
	// Filter out files that have been downloaded.
	var downloading []modules.DownloadInfo
	for _, file := range queue.Downloads {
		if file.Received != file.Filesize {
			downloading = append(downloading, file)
		}
	}
	if len(downloading) == 0 {
		fmt.Println("No files are downloading.")
	} else {
		fmt.Println("Downloading", len(downloading), "files:")
		for _, file := range downloading {
			fmt.Printf("%s: %5.1f%% %s -> %s\n", file.StartTime.Format("Jan 02 03:04 PM"), 100*float64(file.Received)/float64(file.Filesize), file.SiaPath, file.Destination)
		}
	}
	if !renterShowHistory {
		return
	}
	fmt.Println()
	// Filter out files that are downloading.
	var downloaded []modules.DownloadInfo
	for _, file := range queue.Downloads {
		if file.Received == file.Filesize {
			downloaded = append(downloaded, file)
		}
	}
	if len(downloaded) == 0 {
		fmt.Println("No files downloaded.")
	} else {
		fmt.Println("Downloaded", len(downloaded), "files:")
		for _, file := range downloaded {
			fmt.Printf("%s: %s -> %s\n", file.StartTime.Format("Jan 02 03:04 PM"), file.SiaPath, file.Destination)
		}
	}
}

// renterfilesdeletecmd is the handler for the command `siac renter delete [path]`.
// Removes the specified path from the Sia network.
func renterfilesdeletecmd(path string) {
	err := post("/renter/delete/"+path, "")
	if err != nil {
		fmt.Println("Could not delete file:", err)
		return
	}
	fmt.Println("Deleted", path)
}

// renterfilesdownloadcmd is the handler for the comand `siac renter download [path] [destination]`.
// Downloads a path from the Sia network to the local specified destination.
func renterfilesdownloadcmd(path, destination string) {
	err := get("/renter/download/" + path + "?destination=" + abs(destination))
	if err != nil {
		fmt.Println("Could not download file:", err)
		return
	}
	fmt.Printf("Downloaded '%s' to %s.\n", path, abs(destination))
}

// renterfileslistcmd is the handler for the command `siac renter list`.
// Lists files known to the renter on the network.
func renterfileslistcmd() {
	var rf api.RenterFiles
	err := getAPI("/renter/files", &rf)
	if err != nil {
		fmt.Println("Could not get file list:", err)
		return
	}
	if len(rf.Files) == 0 {
		fmt.Println("No files have been uploaded.")
		return
	}
	fmt.Println("Tracking", len(rf.Files), "files:")
	for _, file := range rf.Files {
		if file.Available {
			fmt.Printf("%13s  %s\n", filesizeUnits(int64(file.Filesize)), file.SiaPath)
		} else {
			fmt.Printf("%13s  %s (uploading, %0.2f%%)\n", filesizeUnits(int64(file.Filesize)), file.SiaPath, file.UploadProgress)
		}
	}
}

// renterfilesloadcmd is the handler for the command `siac renter load [source]`.
// Loads a .sia file, adding the file entries contained within.
func renterfilesloadcmd(source string) {
	var info api.RenterLoad
	err := postResp("/renter/load", "source="+abs(source), &info)
	if err != nil {
		fmt.Println("Could not load file:", err)
		return
	}
	fmt.Printf("Loaded %d file(s):\n", len(info.FilesAdded))
	for _, file := range info.FilesAdded {
		fmt.Printf("\t%s\n", file)
	}
}

// renterfilesloadasciicmd is the handler for the command `siac renter loadascii [ascii]`.
// Load an ASCII-encoded .sia file.
func renterfilesloadasciicmd(ascii string) {
	var info api.RenterLoad
	err := postResp("/renter/loadascii", "asciisia="+ascii, &info)
	if err != nil {
		fmt.Println("Could not load file:", err)
		return
	}
	fmt.Printf("Loaded %d file(s):\n", len(info.FilesAdded))
	for _, file := range info.FilesAdded {
		fmt.Printf("\t%s\n", file)
	}
}

// renterfilesrenamecmd is the handler for the command `siac renter rename [path] [newpath]`.
// Renames a file on the Sia network.
func renterfilesrenamecmd(path, newpath string) {
	err := post("/renter/rename/"+path, "newsiapath="+newpath)
	if err != nil {
		fmt.Println("Could not rename file:", err)
		return
	}
	fmt.Printf("Renamed %s to %s\n", path, newpath)
}

// renterfilessharecmd is the handler for the command `siac renter share [path] [destination]`.
// Export a file to a .sia for sharing.
func renterfilessharecmd(path, destination string) {
	err := get(fmt.Sprintf("/renter/share?siapaths=%s&destination=%s", path, abs(destination)))
	if err != nil {
		fmt.Println("Could not share file:", err)
		return
	}
	fmt.Printf("Exported %s to %s\n", path, abs(destination))
}

// renterfilesshareasciicmd is the handler for the command `siac renter shareascii [path]`.
// Prints an ascii-encoded sia file.
func renterfilesshareasciicmd(path string) {
	var data api.RenterShareASCII
	err := getAPI("/renter/shareascii?siapaths="+path, &data)
	if err != nil {
		fmt.Println("Could not share file:", err)
		return
	}
	fmt.Println(data.ASCIIsia)
}

// renterfilesuploadcmd is the handler for the command `siac renter upload [source] [path]`.
// Uploads the [source] file to [path] on the Sia network.
func renterfilesuploadcmd(source, path string) {
	err := post("/renter/upload/"+path, "source="+abs(source))
	if err != nil {
		fmt.Println("Could not upload file:", err)
		return
	}
	fmt.Printf("Uploaded '%s' as %s.\n", abs(source), path)
}
<|endoftext|>


<commit_before>package main_test

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path"
	"strings"

	"google.golang.org/grpc/grpclog"

	"code.cloudfoundry.org/bbs"
	bbsconfig "code.cloudfoundry.org/bbs/cmd/bbs/config"
	bbstestrunner "code.cloudfoundry.org/bbs/cmd/bbs/testrunner"
	"code.cloudfoundry.org/bbs/encryption"
	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/bbs/test_helpers"
	"code.cloudfoundry.org/bbs/test_helpers/sqlrunner"
	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/consuladapter"
	"code.cloudfoundry.org/consuladapter/consulrunner"
	"code.cloudfoundry.org/inigo/helpers/portauthority"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"code.cloudfoundry.org/rep/maintain"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	"testing"
	"time"
)

var (
	auctioneerPath       string
	auctioneerServerPort uint16
	auctioneerLocation   string

	dotNetStack           = "dot-net"
	dotNetRootFSURL       = models.PreloadedRootFS(dotNetStack)
	linuxStack            = "linux"
	linuxRootFSURL        = models.PreloadedRootFS(linuxStack)
	dotNetCell, linuxCell *FakeCell

	consulRunner *consulrunner.ClusterRunner
	consulClient consuladapter.Client

	bbsConfig  bbsconfig.BBSConfig
	bbsBinPath string
	bbsURL     *url.URL
	bbsRunner  *ginkgomon.Runner
	bbsProcess ifrit.Process
	bbsClient  bbs.InternalClient

	locketBinPath string

	sqlProcess ifrit.Process
	sqlRunner  sqlrunner.SQLRunner

	logger        lager.Logger
	portAllocator portauthority.PortAllocator
)

func TestAuctioneer(t *testing.T) {
	// these integration tests can take a bit, especially under load;
	// 1 second is too harsh
	SetDefaultEventuallyTimeout(10 * time.Second)

	RegisterFailHandler(Fail)
	RunSpecs(t, "Auctioneer Cmd Suite")
}

var _ = SynchronizedBeforeSuite(func() []byte {
	compiledBBSPath, err := gexec.Build("code.cloudfoundry.org/bbs/cmd/bbs", "-race")
	Expect(err).NotTo(HaveOccurred())

	compiledAuctioneerPath, err := gexec.Build("code.cloudfoundry.org/auctioneer/cmd/auctioneer", "-race")
	Expect(err).NotTo(HaveOccurred())

	locketPath, err := gexec.Build("code.cloudfoundry.org/locket/cmd/locket", "-race")
	Expect(err).NotTo(HaveOccurred())

	return []byte(strings.Join([]string{compiledAuctioneerPath, compiledBBSPath, locketPath}, ","))
}, func(pathsByte []byte) {
	grpclog.SetLogger(log.New(ioutil.Discard, "", 0))
	node := GinkgoParallelNode()
	startPort := 1050 * node // make sure we don't conflict with etcd ports 4000+GinkgoParallelNode & 7000+GinkgoParallelNode (4000,7000,40001,70001...)
	portRange := 1000
	endPort := startPort + portRange

	var err error
	portAllocator, err = portauthority.New(startPort, endPort)
	Expect(err).NotTo(HaveOccurred())

	paths := string(pathsByte)
	auctioneerPath = strings.Split(paths, ",")[0]
	bbsBinPath = strings.Split(paths, ",")[1]
	locketBinPath = strings.Split(paths, ",")[2]

	dbName := fmt.Sprintf("diego_%d", GinkgoParallelNode())
	sqlRunner = test_helpers.NewSQLRunner(dbName)
	sqlProcess = ginkgomon.Invoke(sqlRunner)

	consulStartingPort, err := portAllocator.ClaimPorts(consulrunner.PortOffsetLength)
	Expect(err).NotTo(HaveOccurred())
	consulRunner = consulrunner.NewClusterRunner(
		consulrunner.ClusterRunnerConfig{
			StartingPort: int(consulStartingPort),
			NumNodes:     1,
			Scheme:       "http",
		},
	)

	auctioneerServerPort, err = portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	auctioneerLocation = fmt.Sprintf("127.0.0.1:%d", auctioneerServerPort)

	logger = lagertest.NewTestLogger("test")

	consulRunner.Start()
	consulRunner.WaitUntilReady()

	bbsPort, err := portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	healthPort, err := portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	bbsAddress := fmt.Sprintf("127.0.0.1:%d", bbsPort)
	healthAddress := fmt.Sprintf("127.0.0.1:%d", healthPort)

	bbsURL = &url.URL{
		Scheme: "https",
		Host:   bbsAddress,
	}

	fixturesPath := path.Join(os.Getenv("GOPATH"), "src/code.cloudfoundry.org/auctioneer/cmd/auctioneer/fixtures")

	bbsConfig = bbsconfig.BBSConfig{
		ListenAddress:     bbsAddress,
		AdvertiseURL:      bbsURL.String(),
		AuctioneerAddress: "http://" + auctioneerLocation,
		ConsulCluster:     consulRunner.ConsulCluster(),
		HealthAddress:     healthAddress,

		EncryptionConfig: encryption.EncryptionConfig{
			EncryptionKeys: map[string]string{"label": "key"},
			ActiveKeyLabel: "label",
		},

		DatabaseDriver:                sqlRunner.DriverName(),
		DatabaseConnectionString:      sqlRunner.ConnectionString(),
		DetectConsulCellRegistrations: true,

		CaFile:   path.Join(fixturesPath, "green-certs", "ca.crt"),
		CertFile: path.Join(fixturesPath, "green-certs", "server.crt"),
		KeyFile:  path.Join(fixturesPath, "green-certs", "server.key"),
	}
})

var _ = BeforeEach(func() {
	consulRunner.Reset()

	bbsRunner = bbstestrunner.New(bbsBinPath, bbsConfig)
	bbsProcess = ginkgomon.Invoke(bbsRunner)

	consulClient = consulRunner.NewClient()

	cellPresenceClient := maintain.NewCellPresenceClient(consulClient, clock.NewClock())

	dotNetCell = SpinUpFakeCell(cellPresenceClient, "dot-net-cell", "", dotNetStack)
	linuxCell = SpinUpFakeCell(cellPresenceClient, "linux-cell", "", linuxStack)
})

var _ = AfterEach(func() {
	ginkgomon.Kill(bbsProcess)
	dotNetCell.Stop()
	linuxCell.Stop()

	sqlRunner.Reset()
})

var _ = SynchronizedAfterSuite(func() {
	if consulRunner != nil {
		consulRunner.Stop()
	}

	ginkgomon.Kill(sqlProcess)
}, func() {
	gexec.CleanupBuildArtifacts()
})
<commit_msg>get rid of etcd references<commit_after>package main_test

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path"
	"strings"

	"google.golang.org/grpc/grpclog"

	"code.cloudfoundry.org/bbs"
	bbsconfig "code.cloudfoundry.org/bbs/cmd/bbs/config"
	bbstestrunner "code.cloudfoundry.org/bbs/cmd/bbs/testrunner"
	"code.cloudfoundry.org/bbs/encryption"
	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/bbs/test_helpers"
	"code.cloudfoundry.org/bbs/test_helpers/sqlrunner"
	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/consuladapter"
	"code.cloudfoundry.org/consuladapter/consulrunner"
	"code.cloudfoundry.org/inigo/helpers/portauthority"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"code.cloudfoundry.org/rep/maintain"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	"testing"
	"time"
)

var (
	auctioneerPath       string
	auctioneerServerPort uint16
	auctioneerLocation   string

	dotNetStack           = "dot-net"
	dotNetRootFSURL       = models.PreloadedRootFS(dotNetStack)
	linuxStack            = "linux"
	linuxRootFSURL        = models.PreloadedRootFS(linuxStack)
	dotNetCell, linuxCell *FakeCell

	consulRunner *consulrunner.ClusterRunner
	consulClient consuladapter.Client

	bbsConfig  bbsconfig.BBSConfig
	bbsBinPath string
	bbsURL     *url.URL
	bbsRunner  *ginkgomon.Runner
	bbsProcess ifrit.Process
	bbsClient  bbs.InternalClient

	locketBinPath string

	sqlProcess ifrit.Process
	sqlRunner  sqlrunner.SQLRunner

	logger        lager.Logger
	portAllocator portauthority.PortAllocator
)

func TestAuctioneer(t *testing.T) {
	// these integration tests can take a bit, especially under load;
	// 1 second is too harsh
	SetDefaultEventuallyTimeout(10 * time.Second)

	RegisterFailHandler(Fail)
	RunSpecs(t, "Auctioneer Cmd Suite")
}

var _ = SynchronizedBeforeSuite(func() []byte {
	compiledBBSPath, err := gexec.Build("code.cloudfoundry.org/bbs/cmd/bbs", "-race")
	Expect(err).NotTo(HaveOccurred())

	compiledAuctioneerPath, err := gexec.Build("code.cloudfoundry.org/auctioneer/cmd/auctioneer", "-race")
	Expect(err).NotTo(HaveOccurred())

	locketPath, err := gexec.Build("code.cloudfoundry.org/locket/cmd/locket", "-race")
	Expect(err).NotTo(HaveOccurred())

	return []byte(strings.Join([]string{compiledAuctioneerPath, compiledBBSPath, locketPath}, ","))
}, func(pathsByte []byte) {
	grpclog.SetLogger(log.New(ioutil.Discard, "", 0))
	node := GinkgoParallelNode()
	startPort := 1050 * node
	portRange := 1000
	endPort := startPort + portRange

	var err error
	portAllocator, err = portauthority.New(startPort, endPort)
	Expect(err).NotTo(HaveOccurred())

	paths := string(pathsByte)
	auctioneerPath = strings.Split(paths, ",")[0]
	bbsBinPath = strings.Split(paths, ",")[1]
	locketBinPath = strings.Split(paths, ",")[2]

	dbName := fmt.Sprintf("diego_%d", GinkgoParallelNode())
	sqlRunner = test_helpers.NewSQLRunner(dbName)
	sqlProcess = ginkgomon.Invoke(sqlRunner)

	consulStartingPort, err := portAllocator.ClaimPorts(consulrunner.PortOffsetLength)
	Expect(err).NotTo(HaveOccurred())
	consulRunner = consulrunner.NewClusterRunner(
		consulrunner.ClusterRunnerConfig{
			StartingPort: int(consulStartingPort),
			NumNodes:     1,
			Scheme:       "http",
		},
	)

	auctioneerServerPort, err = portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	auctioneerLocation = fmt.Sprintf("127.0.0.1:%d", auctioneerServerPort)

	logger = lagertest.NewTestLogger("test")

	consulRunner.Start()
	consulRunner.WaitUntilReady()

	bbsPort, err := portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	healthPort, err := portAllocator.ClaimPorts(1)
	Expect(err).NotTo(HaveOccurred())
	bbsAddress := fmt.Sprintf("127.0.0.1:%d", bbsPort)
	healthAddress := fmt.Sprintf("127.0.0.1:%d", healthPort)

	bbsURL = &url.URL{
		Scheme: "https",
		Host:   bbsAddress,
	}

	fixturesPath := path.Join(os.Getenv("GOPATH"), "src/code.cloudfoundry.org/auctioneer/cmd/auctioneer/fixtures")

	bbsConfig = bbsconfig.BBSConfig{
		ListenAddress:     bbsAddress,
		AdvertiseURL:      bbsURL.String(),
		AuctioneerAddress: "http://" + auctioneerLocation,
		ConsulCluster:     consulRunner.ConsulCluster(),
		HealthAddress:     healthAddress,

		EncryptionConfig: encryption.EncryptionConfig{
			EncryptionKeys: map[string]string{"label": "key"},
			ActiveKeyLabel: "label",
		},

		DatabaseDriver:                sqlRunner.DriverName(),
		DatabaseConnectionString:      sqlRunner.ConnectionString(),
		DetectConsulCellRegistrations: true,

		CaFile:   path.Join(fixturesPath, "green-certs", "ca.crt"),
		CertFile: path.Join(fixturesPath, "green-certs", "server.crt"),
		KeyFile:  path.Join(fixturesPath, "green-certs", "server.key"),
	}
})

var _ = BeforeEach(func() {
	consulRunner.Reset()

	bbsRunner = bbstestrunner.New(bbsBinPath, bbsConfig)
	bbsProcess = ginkgomon.Invoke(bbsRunner)

	consulClient = consulRunner.NewClient()

	cellPresenceClient := maintain.NewCellPresenceClient(consulClient, clock.NewClock())

	dotNetCell = SpinUpFakeCell(cellPresenceClient, "dot-net-cell", "", dotNetStack)
	linuxCell = SpinUpFakeCell(cellPresenceClient, "linux-cell", "", linuxStack)
})

var _ = AfterEach(func() {
	ginkgomon.Kill(bbsProcess)
	dotNetCell.Stop()
	linuxCell.Stop()

	sqlRunner.Reset()
})

var _ = SynchronizedAfterSuite(func() {
	if consulRunner != nil {
		consulRunner.Stop()
	}

	ginkgomon.Kill(sqlProcess)
}, func() {
	gexec.CleanupBuildArtifacts()
})
<|endoftext|>


<commit_before><commit_msg>Added more permissions to get team and user info<commit_after><|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/fujiwara/fluent-agent-hydra/fluent"
)

var (
	MessageCh        = make(map[string]map[int64]chan fluent.FluentRecordSet)
	MessageBufferLen = 100
	MessageChMutex   sync.Mutex
	ConnectionId     int64
)

type OutputOption struct {
	IncludeTime bool
	IncludeTag  bool
}

type Encoder interface {
	Encode(interface{}) error
}

type LTSVEncoder struct {
	writer io.Writer
}

func NewLTSVEncoder(w io.Writer) *LTSVEncoder {
	return &LTSVEncoder{writer: w}
}

func (w *LTSVEncoder) Encode(data interface{}) error {
	first := true
	switch record := data.(type) {
	default:
		return fmt.Errorf("unsupported type")
	case map[string]interface{}:
		for key, value := range record {
			if !first {
				_, err := fmt.Fprint(w.writer, "\t")
				if err != nil {
					return err
				}
			} else {
				first = false
			}
			_, err := fmt.Fprint(w.writer, key, ":", value)
			if err != nil {
				return err
			}
		}
		_, err := fmt.Fprint(w.writer, "\n")
		if err != nil {
			return err
		}
	}
	return nil
}

func main() {
	var (
		forwardPort int
		httpPort    int
	)
	flag.IntVar(&forwardPort, "forward-port", 24224, "fluentd forward listen port")
	flag.IntVar(&httpPort, "http-port", 24225, "http listen port")
	flag.Parse()

	var err error
	err = runForwardServer(forwardPort)
	if err != nil {
		log.Fatal(err)
	}
	err = runHTTPServer(httpPort)
	if err != nil {
		log.Fatal(err)
	}
}

func runForwardServer(port int) error {
	addr := fmt.Sprintf(":%d", port)
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	log.Println("[info] forward server listing", l.Addr())
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				log.Println("[error] accept error", err)
			}
			go handleForwardConn(conn)
		}
	}()
	return nil
}

func handleForwardConn(conn net.Conn) {
	defer conn.Close()
	for {
		recordSets, err := fluent.DecodeEntries(conn)
		if err == io.EOF {
			return
		} else if err != nil {
			log.Println("decode entries failed", err, conn.LocalAddr())
			return
		}
		for _, recordSet := range recordSets {
			for tag, channels := range MessageCh {
				if !matchTag(tag, recordSet.Tag) {
					continue
				}
				for _, ch := range channels {
					select {
					case ch <- recordSet:
					}
				}
			}
		}
	}
}

func runHTTPServer(port int) error {
	http.HandleFunc("/", httpHandler)
	addr := fmt.Sprintf(":%d", port)
	log.Println("[info] http server listing", addr)
	return http.ListenAndServe(addr, nil)
}

func httpHandler(w http.ResponseWriter, r *http.Request) {
	tag := strings.Trim(r.URL.Path, "/")
	if tag == "" {
		http.NotFound(w, r)
		return
	}

	option := OutputOption{false, false}
	option.IncludeTime, _ = strconv.ParseBool(r.FormValue("time"))
	option.IncludeTag, _ = strconv.ParseBool(r.FormValue("tag"))
	var encoder Encoder
	switch t, _ := strconv.ParseBool(r.FormValue("ltsv")); t {
	case true:
		encoder = NewLTSVEncoder(w)
	default:
		encoder = json.NewEncoder(w)
	}

	id := atomic.AddInt64(&ConnectionId, 1)
	ch := subscribe(tag, id)
	defer unsubscribe(tag, id)

	log.Printf(
		"[info] client %s tag:%s",
		r.RemoteAddr,
		tag,
	)

	w.WriteHeader(http.StatusOK)
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
	for {
		recordSet := <-ch
		err := writeResponse(encoder, w, recordSet, option)
		if err != nil {
			return
		}
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
	}
	return
}

func matchTag(matcher, tag string) bool {
	matcher = strings.TrimRight(matcher, ".*")
	return strings.HasPrefix(tag, matcher)
}

func subscribe(tag string, id int64) chan fluent.FluentRecordSet {
	MessageChMutex.Lock()
	defer MessageChMutex.Unlock()

	if m := MessageCh[tag]; m == nil {
		MessageCh[tag] = make(map[int64]chan fluent.FluentRecordSet)
	}
	ch := make(chan fluent.FluentRecordSet, MessageBufferLen)
	MessageCh[tag][id] = ch
	return ch
}

func unsubscribe(tag string, id int64) {
	MessageChMutex.Lock()
	defer MessageChMutex.Unlock()

	delete(MessageCh[tag], id)
}

func writeResponse(encoder Encoder, w http.ResponseWriter, rs fluent.FluentRecordSet, option OutputOption) error {
	for _, record := range rs.Records {
		record, _ := record.(*fluent.TinyFluentRecord)
		data := record.GetAllData()
		if option.IncludeTime {
			fmt.Fprint(w, time.Unix(record.Timestamp, 0).Format(time.RFC3339), "\t")
		}
		if option.IncludeTag {
			fmt.Fprint(w, rs.Tag, "\t")
		}
		err := encoder.Encode(data)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>add message count report log<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/fujiwara/fluent-agent-hydra/fluent"
)

var (
	MessageCh        = make(map[string]map[int64]chan fluent.FluentRecordSet)
	MessageBufferLen = 100
	MessageChMutex   sync.Mutex
	ConnectionId     int64
	Counter          int64
)

type OutputOption struct {
	IncludeTime bool
	IncludeTag  bool
}

type Encoder interface {
	Encode(interface{}) error
}

type LTSVEncoder struct {
	writer io.Writer
}

func NewLTSVEncoder(w io.Writer) *LTSVEncoder {
	return &LTSVEncoder{writer: w}
}

func (w *LTSVEncoder) Encode(data interface{}) error {
	first := true
	switch record := data.(type) {
	default:
		return fmt.Errorf("unsupported type")
	case map[string]interface{}:
		for key, value := range record {
			if !first {
				_, err := fmt.Fprint(w.writer, "\t")
				if err != nil {
					return err
				}
			} else {
				first = false
			}
			_, err := fmt.Fprint(w.writer, key, ":", value)
			if err != nil {
				return err
			}
		}
		_, err := fmt.Fprint(w.writer, "\n")
		if err != nil {
			return err
		}
	}
	return nil
}

func main() {
	var (
		forwardPort   int
		httpPort      int
		countInterval int
	)
	flag.IntVar(&forwardPort, "forward-port", 24224, "fluentd forward listen port")
	flag.IntVar(&httpPort, "http-port", 24225, "http listen port")
	flag.IntVar(&countInterval, "count-interval", 60, "log counter output interval(sec)")
	flag.Parse()

	go runReporter(time.Duration(countInterval) * time.Second)

	var err error
	err = runForwardServer(forwardPort)
	if err != nil {
		log.Fatal(err)
	}
	err = runHTTPServer(httpPort)
	if err != nil {
		log.Fatal(err)
	}
}

func runReporter(t time.Duration) {
	ticker := time.Tick(t)
	for _ = range ticker {
		c := atomic.SwapInt64(&Counter, 0)
		if c > 0 {
			log.Println("count:", c, "cps:", float64(c)/float64(t/time.Second))
		}
	}
}

func runForwardServer(port int) error {
	addr := fmt.Sprintf(":%d", port)
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	log.Println("[info] forward server listing", l.Addr())
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				log.Println("[error] accept error", err)
			}
			go handleForwardConn(conn)
		}
	}()
	return nil
}

func handleForwardConn(conn net.Conn) {
	defer conn.Close()
	for {
		recordSets, err := fluent.DecodeEntries(conn)
		if err == io.EOF {
			return
		} else if err != nil {
			log.Println("decode entries failed", err, conn.LocalAddr())
			return
		}
		for _, recordSet := range recordSets {
			atomic.AddInt64(&Counter, int64(len(recordSet.Records)))
			for tag, channels := range MessageCh {
				if !matchTag(tag, recordSet.Tag) {
					continue
				}
				for _, ch := range channels {
					select {
					case ch <- recordSet:
					}
				}
			}
		}
	}
}

func runHTTPServer(port int) error {
	http.HandleFunc("/", httpHandler)
	addr := fmt.Sprintf(":%d", port)
	log.Println("[info] http server listing", addr)
	return http.ListenAndServe(addr, nil)
}

func httpHandler(w http.ResponseWriter, r *http.Request) {
	tag := strings.Trim(r.URL.Path, "/")
	if tag == "" {
		http.NotFound(w, r)
		return
	}

	option := OutputOption{false, false}
	option.IncludeTime, _ = strconv.ParseBool(r.FormValue("time"))
	option.IncludeTag, _ = strconv.ParseBool(r.FormValue("tag"))
	var encoder Encoder
	switch t, _ := strconv.ParseBool(r.FormValue("ltsv")); t {
	case true:
		encoder = NewLTSVEncoder(w)
	default:
		encoder = json.NewEncoder(w)
	}

	id := atomic.AddInt64(&ConnectionId, 1)
	ch := subscribe(tag, id)
	defer unsubscribe(tag, id)

	log.Printf(
		"[info] client %s tag:%s",
		r.RemoteAddr,
		tag,
	)

	w.WriteHeader(http.StatusOK)
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
	for {
		recordSet := <-ch
		err := writeResponse(encoder, w, recordSet, option)
		if err != nil {
			return
		}
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
	}
	return
}

func matchTag(matcher, tag string) bool {
	matcher = strings.TrimRight(matcher, ".*")
	return strings.HasPrefix(tag, matcher)
}

func subscribe(tag string, id int64) chan fluent.FluentRecordSet {
	MessageChMutex.Lock()
	defer MessageChMutex.Unlock()

	if m := MessageCh[tag]; m == nil {
		MessageCh[tag] = make(map[int64]chan fluent.FluentRecordSet)
	}
	ch := make(chan fluent.FluentRecordSet, MessageBufferLen)
	MessageCh[tag][id] = ch
	return ch
}

func unsubscribe(tag string, id int64) {
	MessageChMutex.Lock()
	defer MessageChMutex.Unlock()

	delete(MessageCh[tag], id)
}

func writeResponse(encoder Encoder, w http.ResponseWriter, rs fluent.FluentRecordSet, option OutputOption) error {
	for _, record := range rs.Records {
		record, _ := record.(*fluent.TinyFluentRecord)
		data := record.GetAllData()
		if option.IncludeTime {
			fmt.Fprint(w, time.Unix(record.Timestamp, 0).Format(time.RFC3339), "\t")
		}
		if option.IncludeTag {
			fmt.Fprint(w, rs.Tag, "\t")
		}
		err := encoder.Encode(data)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package daemon

import (
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/socketplane/socketplane/Godeps/_workspace/src/github.com/samalba/dockerclient"
)

type adapterRequest struct {
	PowerstripProtocolVersion int
	Type                      string
	ClientRequest             struct {
		Method  string
		Request string
		Body    string
	}
	ServerResponse struct {
		Body        string
		Code        int
		ContentType string
	}
}

type adapterPreResponse struct {
	PowerstripProtocolVersion int
	ModifiedClientRequest     struct {
		Method  string
		Request string
		Body    string
	}
}

type adapterPostResponse struct {
	PowerstripProtocolVersion int
	ModifiedServerResponse    struct {
		Body        string
		Code        int
		ContentType string
	}
}

func psAdapterPreHook(d *Daemon, reqParams adapterRequest) *adapterPreResponse {
	if reqParams.ClientRequest.Body != "" {
		jsonBody := &dockerclient.ContainerConfig{}
		err := json.Unmarshal([]byte(reqParams.ClientRequest.Body), &jsonBody)
		if err != nil {
			fmt.Println("Body JSON unmarshall failed", err)
		}

		jsonBody.HostConfig.NetworkMode = "none"

		preResp := &adapterPreResponse{}
		preResp.PowerstripProtocolVersion = reqParams.PowerstripProtocolVersion
		preResp.ModifiedClientRequest.Method = reqParams.ClientRequest.Method
		preResp.ModifiedClientRequest.Request = reqParams.ClientRequest.Request

		body, _ := json.Marshal(jsonBody)
		preResp.ModifiedClientRequest.Body = string(body)

		return preResp
	}
	return nil
}

func psAdapterPostHook(d *Daemon, reqParams adapterRequest) *adapterPostResponse {
	if reqParams.ClientRequest.Method != "POST" &&
		reqParams.ClientRequest.Method != "DELETE" {
		fmt.Println("Invalid method: ", reqParams.ClientRequest.Method)
		return nil
	}

	if reqParams.ClientRequest.Request != "" {
		// start api looks like this /<version>/containers/<cid>/start
		s := regexp.MustCompile("/").Split(reqParams.ClientRequest.Request, 5)
		cid := s[3]

		var cfg = &Connection{}
		var op = ConnectionAdd

		switch reqParams.ClientRequest.Method {
		case "POST":
			docker, _ := dockerclient.NewDockerClient(
				"unix:///var/run/docker.sock", nil)
			info, err := docker.InspectContainer(cid)
			if err != nil {
				fmt.Println("InspectContainer failed", err)
				return nil
			}

			cfg.ContainerID = string(cid)
			cfg.ContainerName = info.Name
			cfg.ContainerPID = strconv.Itoa(info.State.Pid)
			cfg.Network = DefaultNetworkName
			for _, env := range info.Config.Env {
				val := regexp.MustCompile("=").Split(env, 3)
				if val[0] == "SP_NETWORK" {
					cfg.Network = strings.Trim(val[1], " ")
				}
			}

			op = ConnectionAdd
		case "DELETE":

			var ok bool
			if cfg, ok = d.Connections[cid]; !ok {
				fmt.Println("Couldn't find the connection", cid)
				return nil
			}

			op = ConnectionDelete
		}

		context := &ConnectionContext{
			op,
			cfg,
			make(chan *Connection),
		}

		d.cC <- context

		<-context.Result

		postResp := &adapterPostResponse{}
		postResp.PowerstripProtocolVersion = reqParams.PowerstripProtocolVersion
		postResp.ModifiedServerResponse.ContentType = "application/json"
		postResp.ModifiedServerResponse.Body = reqParams.ServerResponse.Body
		postResp.ModifiedServerResponse.Code = reqParams.ServerResponse.Code

		return postResp
	}

	return nil
}

func psAdapter(d *Daemon, w http.ResponseWriter, r *http.Request) *apiError {
	var reqParams adapterRequest
	decoder := json.NewDecoder(r.Body)
	err := decoder.Decode(&reqParams)
	if err != nil {
		fmt.Println("Error decodeing JSON", err)
		//return &apiError{http.StatusInternalServerError, err.Error()}
	}

	var data []byte
	switch reqParams.Type {
	case "pre-hook":
		data, _ = json.Marshal(psAdapterPreHook(d, reqParams))
	case "post-hook":
		data, _ = json.Marshal(psAdapterPostHook(d, reqParams))
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(data)
	return nil
}
<commit_msg>Added proper error response<commit_after>package daemon

import (
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/socketplane/socketplane/Godeps/_workspace/src/github.com/samalba/dockerclient"
)

type adapterRequest struct {
	PowerstripProtocolVersion int
	Type                      string
	ClientRequest             struct {
		Method  string
		Request string
		Body    string
	}
	ServerResponse struct {
		Body        string
		Code        int
		ContentType string
	}
}

type adapterPreResponse struct {
	PowerstripProtocolVersion int
	ModifiedClientRequest     struct {
		Method  string
		Request string
		Body    string
	}
}

type adapterPostResponse struct {
	PowerstripProtocolVersion int
	ModifiedServerResponse    struct {
		Body        string
		Code        int
		ContentType string
	}
}

func psAdapterPreHook(d *Daemon, reqParams adapterRequest) (preResp *adapterPreResponse) {
	preResp = &adapterPreResponse{}
	preResp.PowerstripProtocolVersion = reqParams.PowerstripProtocolVersion
	preResp.ModifiedClientRequest.Method = reqParams.ClientRequest.Method
	preResp.ModifiedClientRequest.Request = reqParams.ClientRequest.Request
	preResp.ModifiedClientRequest.Body = reqParams.ClientRequest.Body

	if reqParams.ClientRequest.Body != "" {
		jsonBody := &dockerclient.ContainerConfig{}
		err := json.Unmarshal([]byte(reqParams.ClientRequest.Body), &jsonBody)
		if err != nil {
			fmt.Println("Body JSON unmarshall failed", err)
		}

		jsonBody.HostConfig.NetworkMode = "none"

		body, _ := json.Marshal(jsonBody)
		preResp.ModifiedClientRequest.Body = string(body)

	}

	return
}

func psAdapterPostHook(d *Daemon, reqParams adapterRequest) (postResp *adapterPostResponse) {
	postResp = &adapterPostResponse{}
	postResp.PowerstripProtocolVersion = reqParams.PowerstripProtocolVersion
	postResp.ModifiedServerResponse.ContentType = "application/json"
	postResp.ModifiedServerResponse.Body = reqParams.ServerResponse.Body
	postResp.ModifiedServerResponse.Code = reqParams.ServerResponse.Code

	if reqParams.ClientRequest.Method != "POST" &&
		reqParams.ClientRequest.Method != "DELETE" {
		fmt.Println("Invalid method: ", reqParams.ClientRequest.Method)
		postResp.ModifiedServerResponse.Code = 500
		return
	}

	if reqParams.ClientRequest.Request != "" {
		// start api looks like this /<version>/containers/<cid>/start
		s := regexp.MustCompile("/").Split(reqParams.ClientRequest.Request, 5)
		cid := s[3]

		var cfg = &Connection{}
		var op = ConnectionAdd

		switch reqParams.ClientRequest.Method {
		case "POST":
			docker, _ := dockerclient.NewDockerClient(
				"unix:///var/run/docker.sock", nil)
			info, err := docker.InspectContainer(cid)
			if err != nil {
				fmt.Println("InspectContainer failed", err)
				postResp.ModifiedServerResponse.Code = 404
				return
			}

			cfg.ContainerID = string(cid)
			cfg.ContainerName = info.Name
			cfg.ContainerPID = strconv.Itoa(info.State.Pid)
			cfg.Network = DefaultNetworkName
			for _, env := range info.Config.Env {
				val := regexp.MustCompile("=").Split(env, 3)
				if val[0] == "SP_NETWORK" {
					cfg.Network = strings.Trim(val[1], " ")
				}
			}

			op = ConnectionAdd
		case "DELETE":
			var ok bool
			if cfg, ok = d.Connections[cid]; !ok {
				fmt.Println("Couldn't find the connection", cid)
				postResp.ModifiedServerResponse.Code = 500
				return
			}

			op = ConnectionDelete
		}

		context := &ConnectionContext{
			op,
			cfg,
			make(chan *Connection),
		}

		d.cC <- context

		<-context.Result
	}

	return
}

func psAdapter(d *Daemon, w http.ResponseWriter, r *http.Request) *apiError {
	var reqParams adapterRequest
	decoder := json.NewDecoder(r.Body)
	err := decoder.Decode(&reqParams)
	if err != nil {
		fmt.Println("Error decodeing JSON", err)
		//return &apiError{http.StatusInternalServerError, err.Error()}
	}

	var data []byte
	switch reqParams.Type {
	case "pre-hook":
		data, _ = json.Marshal(psAdapterPreHook(d, reqParams))
	case "post-hook":
		data, _ = json.Marshal(psAdapterPostHook(d, reqParams))
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(data)
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package commands

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/gnuflag"
	"gopkg.in/juju/names.v2"

	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
)

func newResolvedCommand() cmd.Command {
	return modelcmd.Wrap(&resolvedCommand{})
}

// resolvedCommand marks a unit in an error state as ready to continue.
type resolvedCommand struct {
	modelcmd.ModelCommandBase
	UnitName string
	NoRetry  bool
}

func (c *resolvedCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "resolved",
		Args:    "<unit>",
		Purpose: "Marks unit errors resolved and re-executes failed hooks",
	}
}

func (c *resolvedCommand) SetFlags(f *gnuflag.FlagSet) {
	c.ModelCommandBase.SetFlags(f)
	f.BoolVar(&c.NoRetry, "no-retry", false, "do not re-execute failed hooks on the unit")
}

func (c *resolvedCommand) Init(args []string) error {
	if len(args) > 0 {
		c.UnitName = args[0]
		if !names.IsValidUnit(c.UnitName) {
			return errors.Errorf("invalid unit name %q", c.UnitName)
		}
		args = args[1:]
	} else {
		return errors.Errorf("no unit specified")
	}
	return cmd.CheckEmpty(args)
}

func (c *resolvedCommand) Run(_ *cmd.Context) error {
	client, err := c.NewAPIClient()
	if err != nil {
		return err
	}
	defer client.Close()
	return block.ProcessBlockedError(client.Resolved(c.UnitName, c.NoRetry), block.BlockChange)
}
<commit_msg>fix capitalization for flag description<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package commands

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/gnuflag"
	"gopkg.in/juju/names.v2"

	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
)

func newResolvedCommand() cmd.Command {
	return modelcmd.Wrap(&resolvedCommand{})
}

// resolvedCommand marks a unit in an error state as ready to continue.
type resolvedCommand struct {
	modelcmd.ModelCommandBase
	UnitName string
	NoRetry  bool
}

func (c *resolvedCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "resolved",
		Args:    "<unit>",
		Purpose: "Marks unit errors resolved and re-executes failed hooks",
	}
}

func (c *resolvedCommand) SetFlags(f *gnuflag.FlagSet) {
	c.ModelCommandBase.SetFlags(f)
	f.BoolVar(&c.NoRetry, "no-retry", false, "Do not re-execute failed hooks on the unit")
}

func (c *resolvedCommand) Init(args []string) error {
	if len(args) > 0 {
		c.UnitName = args[0]
		if !names.IsValidUnit(c.UnitName) {
			return errors.Errorf("invalid unit name %q", c.UnitName)
		}
		args = args[1:]
	} else {
		return errors.Errorf("no unit specified")
	}
	return cmd.CheckEmpty(args)
}

func (c *resolvedCommand) Run(_ *cmd.Context) error {
	client, err := c.NewAPIClient()
	if err != nil {
		return err
	}
	defer client.Close()
	return block.ProcessBlockedError(client.Resolved(c.UnitName, c.NoRetry), block.BlockChange)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/spf13/cobra"
)

// publish_cloudwatch_metricsCmd respresents the publish_cloudwatch_metrics command
var publish_cloudwatch_metricsCmd = &cobra.Command{
	Use:   "publish_cloudwatch_metrics",
	Short: "Publish queue metrics to CloudWatch in order to trigger auto-scale alarms",
	Long:  `Publish queue metrics to CloudWatch in order to trigger auto-scale alarms`,
	Run: func(cmd *cobra.Command, args []string) {

		if err := cmd.ParseFlags(args); err != nil {
			log.Printf("err: %v", err)
			return
		}

		awsKey := cmd.Flag("aws_key")

		awsKeyVal := awsKey.Value.String()

		log.Printf("awsKey: %v", awsKeyVal)
		if awsKeyVal == "" {
			log.Printf("ERROR: Missing: --aws_key.\n  %v", cmd.UsageString())
			return
		}

		svc := ec2.New(session.New(), &aws.Config{Region: aws.String("us-east-1")})

		// Call the DescribeInstances Operation
		resp, err := svc.DescribeInstances(nil)
		if err != nil {
			panic(err)
		}

		// resp has all of the response data, pull out instance IDs:
		fmt.Println("> Number of reservation sets: ", len(resp.Reservations))
		for idx, res := range resp.Reservations {
			fmt.Println("  > Number of instances: ", len(res.Instances))
			for _, inst := range resp.Reservations[idx].Instances {
				fmt.Println("    - Instance ID: ", *inst.InstanceId)
				fmt.Println("    - Key Name: ", *inst.KeyName)
			}
		}

		// TODO: push something to CloudWatch
		// CLI example:
		//   aws cloudwatch put-metric-data
		//     --metric-name PageViewCount
		//      --namespace "MyService"
		//     --value 2
		//     --timestamp 2014-02-14T12:00:00.000Z
		cloudwatchSvc := cloudwatch.New(session.New(), &aws.Config{Region: aws.String("us-east-1")})

		log.Printf("cloudwatchSvc: %v", cloudwatchSvc)

		metricName := "NumJobsReadyOrBeingProcessed"
		metricValue := 0.0
		timestamp := time.Now()

		metricDatum := &cloudwatch.MetricDatum{
			MetricName: &metricName,
			Value:      &metricValue,
			Timestamp:  &timestamp,
		}

		metricDatumSlice := []*cloudwatch.MetricDatum{metricDatum}
		namespace := "DeepStyleQueue"

		putMetricDataInput := &cloudwatch.PutMetricDataInput{
			MetricData: metricDatumSlice,
			Namespace:  &namespace,
		}

		out, err := cloudwatchSvc.PutMetricData(putMetricDataInput)
		if err != nil {
			log.Printf("ERROR adding metric data  %v", err)
			return
		}
		log.Printf("Metric data output: %v", out)

	},
}

func init() {
	RootCmd.AddCommand(publish_cloudwatch_metricsCmd)

	// Here you will define your flags and configuration settings

	// Cobra supports Persistent Flags which will work for this command and all subcommands

	publish_cloudwatch_metricsCmd.PersistentFlags().String("aws_key", "", "AWS Key")

	// Cobra supports local flags which will only run when this command is called directly
	// publish_cloudwatch_metricsCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle" )

}
<commit_msg>Try to query view<commit_after>package cmd

import (
	"fmt"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/spf13/cobra"
	"github.com/tleyden/go-couch"
)

// publish_cloudwatch_metricsCmd respresents the publish_cloudwatch_metrics command
var publish_cloudwatch_metricsCmd = &cobra.Command{
	Use:   "publish_cloudwatch_metrics",
	Short: "Publish queue metrics to CloudWatch in order to trigger auto-scale alarms",
	Long:  `Publish queue metrics to CloudWatch in order to trigger auto-scale alarms`,
	Run: func(cmd *cobra.Command, args []string) {

		if err := cmd.ParseFlags(args); err != nil {
			log.Printf("err: %v", err)
			return
		}

		awsKey := cmd.Flag("aws_key")

		awsKeyVal := awsKey.Value.String()

		log.Printf("awsKey: %v", awsKeyVal)
		if awsKeyVal == "" {
			log.Printf("ERROR: Missing: --aws_key.\n  %v", cmd.UsageString())
			return
		}

		urlFlag := cmd.Flag("admin_url")

		urlVal := urlFlag.Value.String()
		if urlVal == "" {
			log.Printf("ERROR: Missing: --url.\n  %v", cmd.UsageString())
			return
		}

		err := addCloudWatchMetrics(urlVal)
		if err != nil {
			log.Printf("ERROR: %v", err)
			return
		}

	},
}

func numJobsReadOrBeingProcessed(syncGwAdminUrl string) (metricValue float64, err error) {

	// try to query view
	//    curl localhost:4985/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs
	// if we get a 404, then install the view and then requery

	// if it has a trailing slash, remove it
	rawUrl := strings.TrimSuffix(syncGwAdminUrl, "/")

	// url validation
	url, err := url.Parse(rawUrl)
	if err != nil {
		return 0.0, err
	}

	db, err := couch.Connect(url.String())
	if err != nil {
		return 0.0, fmt.Errorf("Error connecting to db: %v.  Err: %v", syncGwAdminUrl, err)
	}
	log.Printf("connected to db: %v", db)

	viewUrl := "_design/unprocessed_jobs/_view/unprocessed_jobs"
	options := map[string]interface{}{}
	output := map[string]interface{}{}
	err = db.Query(viewUrl, options, &output)
	if err != nil {
		if strings.Contains(err.Error(), "404") {

			// the view doesn't exist yet, attempt to install view

			// now retry
			errInner := db.Query(viewUrl, options, &output)
			if errInner != nil {
				// failed again, give up
				return 0.0, errInner
			}

		} else {
			return 0.0, err
		}

	}
	log.Printf("output: %+v", output)

	// TODO: count the number of rows / keys in the output and
	// convert to a float and return that value

	return 0.0, nil
}

func addCloudWatchMetrics(syncGwAdminUrl string) error {

	metricValue, err := numJobsReadOrBeingProcessed(syncGwAdminUrl)
	if err != nil {
		return err
	}

	// TODO: push something to CloudWatch
	// CLI example:
	//   aws cloudwatch put-metric-data
	//     --metric-name PageViewCount
	//      --namespace "MyService"
	//     --value 2
	//     --timestamp 2014-02-14T12:00:00.000Z
	cloudwatchSvc := cloudwatch.New(session.New(), &aws.Config{Region: aws.String("us-east-1")})

	log.Printf("cloudwatchSvc: %v", cloudwatchSvc)

	metricName := "NumJobsReadyOrBeingProcessed"
	timestamp := time.Now()

	metricDatum := &cloudwatch.MetricDatum{
		MetricName: &metricName,
		Value:      &metricValue,
		Timestamp:  &timestamp,
	}

	metricDatumSlice := []*cloudwatch.MetricDatum{metricDatum}
	namespace := "DeepStyleQueue"

	putMetricDataInput := &cloudwatch.PutMetricDataInput{
		MetricData: metricDatumSlice,
		Namespace:  &namespace,
	}

	out, err := cloudwatchSvc.PutMetricData(putMetricDataInput)
	if err != nil {
		log.Printf("ERROR adding metric data  %v", err)
		return err
	}
	log.Printf("Metric data output: %v", out)

	return nil

}

func init() {
	RootCmd.AddCommand(publish_cloudwatch_metricsCmd)

	// Here you will define your flags and configuration settings

	// Cobra supports Persistent Flags which will work for this command and all subcommands

	publish_cloudwatch_metricsCmd.PersistentFlags().String("aws_key", "", "AWS Key")

	publish_cloudwatch_metricsCmd.PersistentFlags().String("admin_url", "", "Sync Gateway Admin URL")

	// Cobra supports local flags which will only run when this command is called directly
	// publish_cloudwatch_metricsCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle" )

}
<|endoftext|>


<commit_before>package loadspec

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"
	"time"

	"github.com/danielfireman/esperf/loadspec"
	"github.com/spf13/cobra"
)

var (
	indexOverrides []string
	maxDuration    time.Duration
)

func init() {
	parseSlowlogCmd.Flags().StringSliceVar(&indexOverrides, "index_override", []string{}, "Override slowlog indexes. It is a list, flag could be repeated if you would the loadtest to hit many indexes.")
	parseSlowlogCmd.Flags().DurationVar(&maxDuration, "max_duration", time.Duration(0), "Maximum duration of the generated loadspec. It could be smaller, if the slowlog comprise a smaller time frame.")
}

var parseSlowlogCmd = &cobra.Command{
	Use:   "parseslowlog",
	Short: "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	Long:  "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	RunE: func(cmd *cobra.Command, args []string) error {
		// These constants need to be in sync with the regular expression bellow.
		const (
			logTypeField    = "log_type"
			hostField       = "host"
			timestampField  = "ts"
			indexField      = "index"
			typesField      = "types"
			searchTypeField = "search_type"
			sourceField     = "source"
			numFields       = 6
		)

		var urlArg string
		if len(args) > 0 {
			// To keep in par with gen, we only consider the host or host:port part of the URL.
			urlArg = args[0]
			prefix := ""
			switch {
			case strings.HasPrefix(urlArg, "http://"):
				urlArg = strings.TrimPrefix(urlArg, "http://")
				prefix = "http://"
			case strings.HasPrefix(urlArg, "https://"):
				urlArg = strings.TrimPrefix(urlArg, "https://")
				prefix = "https://"
			}
			i := strings.Index(urlArg, "/")
			if i > 0 {
				urlArg = prefix + urlArg[:i]
			}
		}

		// Regular expression setup.
		// The solution is based on regexp's named matches. For each entry, we build a map of
		// of fields and values. This map is encoded as json and (buffered) written to stdout.
		re, err := regexp.Compile(`\[(?P<ts>[^]]+)\].?\[.*\].?\[(?P<log_type>[^]]+)\].?\[(?P<host>[^]]+)\].?\[(?P<index>[^]]+)\].?\[.*\].*types\[(?P<types>[^]]+)\].*search_type\[(?P<search_type>[^]]+)\].*source\[(?P<source>.*)\], extra_source`)
		if err != nil {
			return err
		}
		subExpNames := re.SubexpNames()

		var entries loadspec.ByDelaySinceLastNanos
		fields := make(map[string]string, numFields)
		scanner := bufio.NewScanner(os.Stdin)
		count := 0
		for scanner.Scan() {
			// Building a map using named matches.
			matches := re.FindAllStringSubmatch(scanner.Text(), -1)[0]
			for i, n := range matches {
				// Removing the first match, which is the whole line.
				if i > 0 {
					fields[subExpNames[i]] = n
				}
			}
			// For now, only processing queries.
			if fields[logTypeField] != "index.search.slowlog.query" {
				continue
			}

			entry := loadspec.Entry{Source: fields[sourceField]}
			// Making timestamp relative to the previous one. Simulate inter-arrival time can be as easy
			// as a time.Sleep and trigger a goroutine.
			t, err := time.Parse(timeLayout, strings.Replace(fields[timestampField], ",", ".", 1))
			if err != nil {
				return err
			}
			// Keeping timestamp here for post-processing bellow.
			entry.DelaySinceLastNanos = t.UnixNano()
			// Host argument is treated as full URL. This keeps consistency between here and gen command.

			host := fields[hostField]
			if urlArg != "" {
				host = urlArg
			}
			index := fields[indexField]
			if len(indexOverrides) > 0 {
				index = indexOverrides[count%len(indexOverrides)]
			}

			// I would love to use url.URL, life is hard.
			// More on that: https://github.com/golang/go/issues/18824
			// TL;DR; We would like to use http://localhost:9200, but since go1.8 it is not allowed anymore.
			path := []string{host, index, fields[typesField], "_search"}
			st := ""
			if fields[searchTypeField] != "" {
				st = fmt.Sprintf("?search_type=%s", strings.ToLower(fields[searchTypeField]))
			}
			entry.URL = fmt.Sprintf("%s%s", strings.Join(path, "/"), st)
			entries = append(entries, &entry)
			count++
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		// Slow log entries are not guaranteed to be timestamp ordered.
		sort.Sort(entries)

		// Writer and encoding configuration.
		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		enc := json.NewEncoder(writer)
		var elapsed, previousTimestamp, currTimestamp int64
		for i, e := range entries {
			// Adjusting from timestamp to delay since last request. That makes a lot easier to replay.
			currTimestamp = e.DelaySinceLastNanos
			if i == 0 {
				e.DelaySinceLastNanos = 0
			} else {
				e.DelaySinceLastNanos -= previousTimestamp
			}
			previousTimestamp = currTimestamp
			if err := enc.Encode(&e); err != nil {
				return err
			}
			elapsed += e.DelaySinceLastNanos
			if elapsed >= maxDuration.Nanoseconds() {
				break
			}
		}
		return nil
	},
}
<commit_msg>Fixing a bug on argument parsing and renaming the index_override flag<commit_after>package loadspec

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"
	"time"

	"github.com/danielfireman/esperf/loadspec"
	"github.com/spf13/cobra"
)

var (
	indexOverride []string
	maxDuration   time.Duration
)

func init() {
	parseSlowlogCmd.Flags().StringSliceVar(&indexOverride, "index_override", []string{}, "Override slowlog indexes. It is a list, flag could be repeated if you would the loadtest to hit many indexes.")
	parseSlowlogCmd.Flags().DurationVar(&maxDuration, "max_duration", time.Duration(0), "Maximum duration of the generated loadspec. It could be smaller, if the slowlog comprise a smaller time frame.")
}

var parseSlowlogCmd = &cobra.Command{
	Use:   "parseslowlog",
	Short: "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	Long:  "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	RunE: func(cmd *cobra.Command, args []string) error {
		// These constants need to be in sync with the regular expression bellow.
		const (
			logTypeField    = "log_type"
			hostField       = "host"
			timestampField  = "ts"
			indexField      = "index"
			typesField      = "types"
			searchTypeField = "search_type"
			sourceField     = "source"
			numFields       = 6
		)

		var urlArg string
		if len(args) > 0 {
			// To keep in par with gen, we only consider the host or host:port part of the URL.
			urlArg = args[0]
			prefix := ""
			switch {
			case strings.HasPrefix(urlArg, "http://"):
				urlArg = strings.TrimPrefix(urlArg, "http://")
				prefix = "http://"
			case strings.HasPrefix(urlArg, "https://"):
				urlArg = strings.TrimPrefix(urlArg, "https://")
				prefix = "https://"
			}
			i := strings.Index(urlArg, "/")
			if i > 0 {
				urlArg = urlArg[:i]
			}
			urlArg = prefix + urlArg
		}

		// Regular expression setup.
		// The solution is based on regexp's named matches. For each entry, we build a map of
		// of fields and values. This map is encoded as json and (buffered) written to stdout.
		re, err := regexp.Compile(`\[(?P<ts>[^]]+)\].?\[.*\].?\[(?P<log_type>[^]]+)\].?\[(?P<host>[^]]+)\].?\[(?P<index>[^]]+)\].?\[.*\].*types\[(?P<types>[^]]+)\].*search_type\[(?P<search_type>[^]]+)\].*source\[(?P<source>.*)\], extra_source`)
		if err != nil {
			return err
		}
		subExpNames := re.SubexpNames()

		var entries loadspec.ByDelaySinceLastNanos
		fields := make(map[string]string, numFields)
		scanner := bufio.NewScanner(os.Stdin)
		count := 0
		for scanner.Scan() {
			// Building a map using named matches.
			matches := re.FindAllStringSubmatch(scanner.Text(), -1)[0]
			for i, n := range matches {
				// Removing the first match, which is the whole line.
				if i > 0 {
					fields[subExpNames[i]] = n
				}
			}
			// For now, only processing queries.
			if fields[logTypeField] != "index.search.slowlog.query" {
				continue
			}

			entry := loadspec.Entry{Source: fields[sourceField]}
			// Making timestamp relative to the previous one. Simulate inter-arrival time can be as easy
			// as a time.Sleep and trigger a goroutine.
			t, err := time.Parse(timeLayout, strings.Replace(fields[timestampField], ",", ".", 1))
			if err != nil {
				return err
			}
			// Keeping timestamp here for post-processing bellow.
			entry.DelaySinceLastNanos = t.UnixNano()
			// Host argument is treated as full URL. This keeps consistency between here and gen command.

			host := fields[hostField]
			if urlArg != "" {
				host = urlArg
			}
			index := fields[indexField]
			if len(indexOverride) > 0 {
				index = indexOverride[count%len(indexOverride)]
			}

			// I would love to use url.URL, life is hard.
			// More on that: https://github.com/golang/go/issues/18824
			// TL;DR; We would like to use http://localhost:9200, but since go1.8 it is not allowed anymore.
			path := []string{host, index, fields[typesField], "_search"}
			st := ""
			if fields[searchTypeField] != "" {
				st = fmt.Sprintf("?search_type=%s", strings.ToLower(fields[searchTypeField]))
			}
			entry.URL = fmt.Sprintf("%s%s", strings.Join(path, "/"), st)
			entries = append(entries, &entry)
			count++
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		// Slow log entries are not guaranteed to be timestamp ordered.
		sort.Sort(entries)

		// Writer and encoding configuration.
		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		enc := json.NewEncoder(writer)
		var elapsed, previousTimestamp, currTimestamp int64
		for i, e := range entries {
			// Adjusting from timestamp to delay since last request. That makes a lot easier to replay.
			currTimestamp = e.DelaySinceLastNanos
			if i == 0 {
				e.DelaySinceLastNanos = 0
			} else {
				e.DelaySinceLastNanos -= previousTimestamp
			}
			previousTimestamp = currTimestamp
			if err := enc.Encode(&e); err != nil {
				return err
			}
			elapsed += e.DelaySinceLastNanos
			if elapsed >= maxDuration.Nanoseconds() {
				break
			}
		}
		return nil
	},
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package scoring

import (
	"context"
	"fmt"

	"github.com/apigee/registry/cmd/registry/core"
	"github.com/apigee/registry/cmd/registry/patch"
	"github.com/apigee/registry/cmd/registry/patterns"
	"github.com/apigee/registry/log"
	"github.com/apigee/registry/rpc"
	"github.com/apigee/registry/server/registry/names"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
)

func scoreCardID(definitionID string) string {
	return fmt.Sprintf("scorecard-%s", definitionID)
}

func FetchScoreCardDefinitions(
	ctx context.Context,
	client artifactClient,
	resource patterns.ResourceName) ([]*rpc.Artifact, error) {
	defArtifacts := make([]*rpc.Artifact, 0)

	project := fmt.Sprintf("%s/locations/global", resource.Project())
	artifact, err := names.ParseArtifact(fmt.Sprintf("%s/artifacts/-", project))
	if err != nil {
		return nil, err
	}
	listFilter := fmt.Sprintf("mime_type == %q", patch.MimeTypeForKind("ScoreCardDefinition"))
	err = client.ListArtifacts(ctx, artifact, listFilter, true,
		func(artifact *rpc.Artifact) error {
			definition := &rpc.ScoreCardDefinition{}
			if err := proto.Unmarshal(artifact.GetContents(), definition); err != nil {
				return err
			}

			// Check if ScoreCardDefinition.TargetResource matches with the supplied resource
			err := matchResourceWithTarget(definition.GetTargetResource(), resource, project)
			if err != nil {
				return err
			}

			defArtifacts = append(defArtifacts, artifact)
			return nil
		})

	if err != nil {
		return nil, err
	}

	return defArtifacts, nil
}

func CalculateScoreCard(
	ctx context.Context,
	client artifactClient,
	defArtifact *rpc.Artifact,
	resource patterns.ResourceInstance,
	dryRun bool) error {
	project := fmt.Sprintf("%s/locations/global", resource.ResourceName().Project())

	// Extract definition
	definition := &rpc.ScoreCardDefinition{}
	if err := proto.Unmarshal(defArtifact.GetContents(), definition); err != nil {
		return err
	}

	var takeAction bool

	// Fetch the to be generated ScoreCard artifact (if present)
	artifactName := fmt.Sprintf("%s/artifacts/%s", resource.ResourceName().String(), scoreCardID(definition.GetId()))
	scoreCardArtifact, err := getArtifact(ctx, client, artifactName, false)
	if err != nil {
		// Generate ScoreCard if the ScoreCard artifact doesn't exist
		if status.Code(err) == codes.NotFound {
			takeAction = true
		} else {
			return fmt.Errorf("failed to fetch artifact %q: %s", artifactName, err)
		}
	}

	// Generate ScoreCard if the definition has been updated
	// This condition is required to avoid the scenario mentioned here: https://github.com/apigee/registry/issues/641
	if scoreCardArtifact != nil && defArtifact.GetUpdateTime().AsTime().Add(patterns.ResourceUpdateThresholdSeconds).After(scoreCardArtifact.GetUpdateTime().AsTime()) {
		takeAction = true
	}

	result := processScorePatterns(ctx, client, definition, resource, scoreCardArtifact, takeAction, project)
	if result.err != nil {
		return err
	}

	if result.needsUpdate {
		if dryRun {
			core.PrintMessage(result.scoreCard)
			return nil
		}
		// upload the scoreCard to registry
		return uploadScoreCard(ctx, client, resource, result.scoreCard)
	}

	log.Debugf(ctx, "ScoreCard %s is already up-to-date.", artifactName)
	return nil
}

// Response returned after fetching all the scoreArtifacts to form a ScoreCard.
type scoreCardResult struct {
	// Represents the ScoreCard generated after fetching all the score artifacts from the score_patterns.
	scoreCard *rpc.ScoreCard
	// Represents if the final scoreCardArtifact needs an update
	// This is determined based on the timestamps of the existing scoreCardArtifact and the dependent scoreArtifacts fetched from score_patterns.
	needsUpdate bool
	// Represents the error encountered while generating the ScoreCard.
	err error
}

func processScorePatterns(
	ctx context.Context,
	client artifactClient,
	definition *rpc.ScoreCardDefinition,
	resource patterns.ResourceInstance,
	scoreCardArtifact *rpc.Artifact,
	takeAction bool,
	project string) scoreCardResult {
	var needsUpdate bool
	scoreArtifacts := make([]*rpc.Score, 0)

	for _, scorePattern := range definition.GetScorePatterns() {
		extendedPattern, err := patterns.SubstituteReferenceEntity(scorePattern, resource.ResourceName())
		if err != nil {
			return scoreCardResult{
				scoreCard:   nil,
				needsUpdate: false,
				err:         fmt.Errorf("invalid pattern %q in score_patterns: %s", scorePattern, err),
			}
		}

		// Fetch scoreArtifact
		artifact, err := getArtifact(ctx, client, extendedPattern.String(), true)
		if err != nil {
			return scoreCardResult{
				scoreCard:   nil,
				needsUpdate: false,
				err:         fmt.Errorf("failed to fetch artifact %s: %s", extendedPattern.String(), err),
			}
		}

		// needsUpdate tells the calling function if the ScoreCard artifact needs to be updated
		// This condition is required to avoid the scenario mentioned here: https://github.com/apigee/registry/issues/641
		needsUpdate = needsUpdate || takeAction || artifact.GetUpdateTime().AsTime().Add(patterns.ResourceUpdateThresholdSeconds).After(scoreCardArtifact.GetUpdateTime().AsTime())
		// Extract Score from the fetched artifact
		score := &rpc.Score{}
		if err := proto.Unmarshal(artifact.GetContents(), score); err != nil {
			return scoreCardResult{
				scoreCard:   nil,
				needsUpdate: false,
				err:         fmt.Errorf("failed unmarshaling artifact %q into Score proto: %s", artifact.GetName(), err),
			}
		}

		scoreArtifacts = append(scoreArtifacts, score)
	}

	if needsUpdate {
		// Build the final ScoreCard proto
		scoreCard := &rpc.ScoreCard{
			Id:             scoreCardID(definition.GetId()),
			Kind:           "ScoreCard",
			DisplayName:    definition.GetDisplayName(),
			Description:    definition.GetDescription(),
			DefinitionName: fmt.Sprintf("%s/artifacts/%s", project, definition.GetId()),
			Scores:         scoreArtifacts,
		}

		return scoreCardResult{
			scoreCard:   scoreCard,
			needsUpdate: true,
			err:         nil,
		}
	}

	return scoreCardResult{
		scoreCard:   nil,
		needsUpdate: false,
		err:         nil,
	}
}

func uploadScoreCard(ctx context.Context, client artifactClient, resource patterns.ResourceInstance, scoreCard *rpc.ScoreCard) error {
	artifactBytes, err := proto.Marshal(scoreCard)
	if err != nil {
		return err
	}
	artifact := &rpc.Artifact{
		Name:     fmt.Sprintf("%s/artifacts/%s", resource.ResourceName().String(), scoreCard.GetId()),
		Contents: artifactBytes,
		MimeType: patch.MimeTypeForKind("ScoreCard"),
	}
	log.Debugf(ctx, "Uploading %s", artifact.GetName())
	if err = client.SetArtifact(ctx, artifact); err != nil {
		return fmt.Errorf("failed to save artifact %s: %s", artifact.GetName(), err)
	}

	return nil
}
<commit_msg>Return the correct error when ProcessScorePattern fails (#761)<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package scoring

import (
	"context"
	"fmt"

	"github.com/apigee/registry/cmd/registry/core"
	"github.com/apigee/registry/cmd/registry/patch"
	"github.com/apigee/registry/cmd/registry/patterns"
	"github.com/apigee/registry/log"
	"github.com/apigee/registry/rpc"
	"github.com/apigee/registry/server/registry/names"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
)

func scoreCardID(definitionID string) string {
	return fmt.Sprintf("scorecard-%s", definitionID)
}

func FetchScoreCardDefinitions(
	ctx context.Context,
	client artifactClient,
	resource patterns.ResourceName) ([]*rpc.Artifact, error) {
	defArtifacts := make([]*rpc.Artifact, 0)

	project := fmt.Sprintf("%s/locations/global", resource.Project())
	artifact, err := names.ParseArtifact(fmt.Sprintf("%s/artifacts/-", project))
	if err != nil {
		return nil, err
	}
	listFilter := fmt.Sprintf("mime_type == %q", patch.MimeTypeForKind("ScoreCardDefinition"))
	err = client.ListArtifacts(ctx, artifact, listFilter, true,
		func(artifact *rpc.Artifact) error {
			definition := &rpc.ScoreCardDefinition{}
			if err := proto.Unmarshal(artifact.GetContents(), definition); err != nil {
				return err
			}

			// Check if ScoreCardDefinition.TargetResource matches with the supplied resource
			err := matchResourceWithTarget(definition.GetTargetResource(), resource, project)
			if err != nil {
				return err
			}

			defArtifacts = append(defArtifacts, artifact)
			return nil
		})

	if err != nil {
		return nil, err
	}

	return defArtifacts, nil
}

func CalculateScoreCard(
	ctx context.Context,
	client artifactClient,
	defArtifact *rpc.Artifact,
	resource patterns.ResourceInstance,
	dryRun bool) error {
	project := fmt.Sprintf("%s/locations/global", resource.ResourceName().Project())

	// Extract definition
	definition := &rpc.ScoreCardDefinition{}
	if err := proto.Unmarshal(defArtifact.GetContents(), definition); err != nil {
		return err
	}

	var takeAction bool

	// Fetch the to be generated ScoreCard artifact (if present)
	artifactName := fmt.Sprintf("%s/artifacts/%s", resource.ResourceName().String(), scoreCardID(definition.GetId()))
	scoreCardArtifact, err := getArtifact(ctx, client, artifactName, false)
	if err != nil {
		// Generate ScoreCard if the ScoreCard artifact doesn't exist
		if status.Code(err) == codes.NotFound {
			takeAction = true
		} else {
			return fmt.Errorf("failed to fetch artifact %q: %s", artifactName, err)
		}
	}

	// Generate ScoreCard if the definition has been updated
	// This condition is required to avoid the scenario mentioned here: https://github.com/apigee/registry/issues/641
	if scoreCardArtifact != nil && defArtifact.GetUpdateTime().AsTime().Add(patterns.ResourceUpdateThresholdSeconds).After(scoreCardArtifact.GetUpdateTime().AsTime()) {
		takeAction = true
	}

	result := processScorePatterns(ctx, client, definition, resource, scoreCardArtifact, takeAction, project)
	if result.err != nil {
		return result.err
	}

	if result.needsUpdate {
		if dryRun {
			core.PrintMessage(result.scoreCard)
			return nil
		}
		// upload the scoreCard to registry
		return uploadScoreCard(ctx, client, resource, result.scoreCard)
	}

	log.Debugf(ctx, "ScoreCard %s is already up-to-date.", artifactName)
	return nil
}

// Response returned after fetching all the scoreArtifacts to form a ScoreCard.
type scoreCardResult struct {
	// Represents the ScoreCard generated after fetching all the score artifacts from the score_patterns.
	scoreCard *rpc.ScoreCard
	// Represents if the final scoreCardArtifact needs an update
	// This is determined based on the timestamps of the existing scoreCardArtifact and the dependent scoreArtifacts fetched from score_patterns.
	needsUpdate bool
	// Represents the error encountered while generating the ScoreCard.
	err error
}

func processScorePatterns(
	ctx context.Context,
	client artifactClient,
	definition *rpc.ScoreCardDefinition,
	resource patterns.ResourceInstance,
	scoreCardArtifact *rpc.Artifact,
	takeAction bool,
	project string) scoreCardResult {
	var needsUpdate bool
	scoreArtifacts := make([]*rpc.Score, 0)

	for _, scorePattern := range definition.GetScorePatterns() {
		extendedPattern, err := patterns.SubstituteReferenceEntity(scorePattern, resource.ResourceName())
		if err != nil {
			return scoreCardResult{
				scoreCard:   nil,
				needsUpdate: false,
				err:         fmt.Errorf("invalid pattern %q in score_patterns: %s", scorePattern, err),
			}
		}

		// Fetch scoreArtifact
		artifact, err := getArtifact(ctx, client, extendedPattern.String(), true)
		if err != nil {
			return scoreCardResult{
				scoreCard:   nil,
				needsUpdate: false,
				err:         fmt.Errorf("failed to fetch artifact %s: %s", extendedPattern.String(), err),
			}
		}

		// needsUpdate tells the calling function if the ScoreCard artifact needs to be updated
		// This condition is required to avoid the scenario mentioned here: https://github.com/apigee/registry/issues/641
		needsUpdate = needsUpdate || takeAction || artifact.GetUpdateTime().AsTime().Add(patterns.ResourceUpdateThresholdSeconds).After(scoreCardArtifact.GetUpdateTime().AsTime())
		// Extract Score from the fetched artifact
		score := &rpc.Score{}
		if err := proto.Unmarshal(artifact.GetContents(), score); err != nil {
			return scoreCardResult{
				scoreCard:   nil,
				needsUpdate: false,
				err:         fmt.Errorf("failed unmarshaling artifact %q into Score proto: %s", artifact.GetName(), err),
			}
		}

		scoreArtifacts = append(scoreArtifacts, score)
	}

	if needsUpdate {
		// Build the final ScoreCard proto
		scoreCard := &rpc.ScoreCard{
			Id:             scoreCardID(definition.GetId()),
			Kind:           "ScoreCard",
			DisplayName:    definition.GetDisplayName(),
			Description:    definition.GetDescription(),
			DefinitionName: fmt.Sprintf("%s/artifacts/%s", project, definition.GetId()),
			Scores:         scoreArtifacts,
		}

		return scoreCardResult{
			scoreCard:   scoreCard,
			needsUpdate: true,
			err:         nil,
		}
	}

	return scoreCardResult{
		scoreCard:   nil,
		needsUpdate: false,
		err:         nil,
	}
}

func uploadScoreCard(ctx context.Context, client artifactClient, resource patterns.ResourceInstance, scoreCard *rpc.ScoreCard) error {
	artifactBytes, err := proto.Marshal(scoreCard)
	if err != nil {
		return err
	}
	artifact := &rpc.Artifact{
		Name:     fmt.Sprintf("%s/artifacts/%s", resource.ResourceName().String(), scoreCard.GetId()),
		Contents: artifactBytes,
		MimeType: patch.MimeTypeForKind("ScoreCard"),
	}
	log.Debugf(ctx, "Uploading %s", artifact.GetName())
	if err = client.SetArtifact(ctx, artifact); err != nil {
		return fmt.Errorf("failed to save artifact %s: %s", artifact.GetName(), err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Build sets scollector version information and should be run from the
// scollector directory.
package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

func main() {
	const path = "main.go"
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}
	for _, d := range f.Decls {
		switch d := d.(type) {
		case *ast.GenDecl:
			if d.Tok != token.CONST {
				continue
			}
			for _, spec := range d.Specs {
				switch spec := spec.(type) {
				case *ast.ValueSpec:
					if len(spec.Names) != 1 || len(spec.Values) != 1 {
						continue
					}
					switch spec.Names[0].Name {
					case "VersionDate":
						switch value := spec.Values[0].(type) {
						case *ast.BasicLit:
							now := time.Now().UTC().Format("20060102150405")
							value.Value = now
						}
					case "VersionID":
						switch value := spec.Values[0].(type) {
						case *ast.BasicLit:
							rev, err := exec.Command("git", "rev-parse", "HEAD").Output()
							if err != nil {
								log.Fatal(err)
							}
							value.Value = fmt.Sprintf(`"%s"`, strings.TrimSpace(string(rev)))
						}
					}
				}
			}
		}
	}
	file, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}
	info, err := file.Stat()
	if err != nil {
		log.Fatal(err)
	}
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, fset, f); err != nil {
		log.Fatal(err)
	}
	fb, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}
	if err := ioutil.WriteFile(path, fb, info.Mode()); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>cmd/scollector: Output version .yml<commit_after>// Build sets scollector version information and should be run from the
// scollector directory.
package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

func main() {
	const path = "main.go"
	var hash, id string
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}
	for _, d := range f.Decls {
		switch d := d.(type) {
		case *ast.GenDecl:
			if d.Tok != token.CONST {
				continue
			}
			for _, spec := range d.Specs {
				switch spec := spec.(type) {
				case *ast.ValueSpec:
					if len(spec.Names) != 1 || len(spec.Values) != 1 {
						continue
					}
					switch spec.Names[0].Name {
					case "VersionDate":
						switch value := spec.Values[0].(type) {
						case *ast.BasicLit:
							id = time.Now().UTC().Format("20060102150405")
							value.Value = id
						}
					case "VersionID":
						switch value := spec.Values[0].(type) {
						case *ast.BasicLit:
							rev, err := exec.Command("git", "rev-parse", "HEAD").Output()
							if err != nil {
								log.Fatal(err)
							}
							hash = fmt.Sprintf(`"%s"`, strings.TrimSpace(string(rev)))
							value.Value = hash
						}
					}
				}
			}
		}
	}
	file, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}
	info, err := file.Stat()
	if err != nil {
		log.Fatal(err)
	}
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, fset, f); err != nil {
		log.Fatal(err)
	}
	fb, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}
	if err := ioutil.WriteFile(path, fb, info.Mode()); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("version:\n  hash: %s\n  id: %s\n", hash, id)
}
<|endoftext|>