<commit_before>package in

import (
	"errors"
	"fmt"
	"log"
	"path/filepath"
	"strings"

	pivnet "github.com/pivotal-cf/go-pivnet"
	"github.com/pivotal-cf/pivnet-resource/concourse"
	"github.com/pivotal-cf/pivnet-resource/metadata"
	"github.com/pivotal-cf/pivnet-resource/versions"
)

//go:generate counterfeiter --fake-name FakeFilter . filter
type filter interface {
	DownloadLinksByGlobs(downloadLinks map[string]string, glob []string, failOnNoMatch bool) (map[string]string, error)
	DownloadLinks(p []pivnet.ProductFile) map[string]string
}

//go:generate counterfeiter --fake-name FakeDownloader . downloader
type downloader interface {
	Download(downloadLinks map[string]string) ([]string, error)
}

//go:generate counterfeiter --fake-name FakeFileSummer . fileSummer
type fileSummer interface {
	SumFile(filepath string) (string, error)
}

//go:generate counterfeiter --fake-name FakeFileWriter . fileWriter
type fileWriter interface {
	WriteMetadataJSONFile(mdata metadata.Metadata) error
	WriteMetadataYAMLFile(mdata metadata.Metadata) error
	WriteVersionFile(versionWithETag string) error
}

//go:generate counterfeiter --fake-name FakePivnetClient . pivnetClient
type pivnetClient interface {
	GetRelease(productSlug string, productVersion string) (pivnet.Release, error)
	AcceptEULA(productSlug string, releaseID int) error
	GetProductFilesForRelease(productSlug string, releaseID int) ([]pivnet.ProductFile, error)
	GetProductFile(productSlug string, releaseID int, productFileID int) (pivnet.ProductFile, error)
	ReleaseDependencies(productSlug string, releaseID int) ([]pivnet.ReleaseDependency, error)
	ReleaseUpgradePaths(productSlug string, releaseID int) ([]pivnet.ReleaseUpgradePath, error)
}

type InCommand struct {
	logger       *log.Logger
	downloadDir  string
	pivnetClient pivnetClient
	filter       filter
	downloader   downloader
	fileSummer   fileSummer
	fileWriter   fileWriter
}

func NewInCommand(
	logger *log.Logger,
	pivnetClient pivnetClient,
	filter filter,
	downloader downloader,
	fileSummer fileSummer,
	fileWriter fileWriter,
) *InCommand {
	return &InCommand{
		logger:       logger,
		pivnetClient: pivnetClient,
		filter:       filter,
		downloader:   downloader,
		fileSummer:   fileSummer,
		fileWriter:   fileWriter,
	}
}

func (c *InCommand) Run(input concourse.InRequest) (concourse.InResponse, error) {
	productSlug := input.Source.ProductSlug

	productVersion, etag, err := versions.SplitIntoVersionAndETag(input.Version.ProductVersion)
	if err != nil {
		c.logger.Println("Parsing of etag failed; continuing without it")
		productVersion = input.Version.ProductVersion
	}

	c.logger.Printf(
		"Getting release for product_slug %s and product_version %s",
		productSlug,
		productVersion,
	)

	release, err := c.pivnetClient.GetRelease(productSlug, productVersion)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Printf("Accepting EULA for release_id %d", release.ID)

	err = c.pivnetClient.AcceptEULA(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Getting product files")

	productFiles, err := c.getProductFiles(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Getting release dependencies")

	releaseDependencies, err := c.pivnetClient.ReleaseDependencies(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Getting release upgrade paths")

	releaseUpgradePaths, err := c.pivnetClient.ReleaseUpgradePaths(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Downloading files")

	err = c.downloadFiles(input.Params.Globs, productFiles, productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Creating metadata")

	versionWithETag, err := versions.CombineVersionAndETag(productVersion, etag)

	mdata := metadata.Metadata{
		Release: &metadata.Release{
			Version:               release.Version,
			ReleaseType:           string(release.ReleaseType),
			ReleaseDate:           release.ReleaseDate,
			Description:           release.Description,
			ReleaseNotesURL:       release.ReleaseNotesURL,
			Availability:          release.Availability,
			Controlled:            release.Controlled,
			ECCN:                  release.ECCN,
			LicenseException:      release.LicenseException,
			EndOfSupportDate:      release.EndOfSupportDate,
			EndOfGuidanceDate:     release.EndOfGuidanceDate,
			EndOfAvailabilityDate: release.EndOfAvailabilityDate,
		},
	}

	if release.EULA != nil {
		mdata.Release.EULASlug = release.EULA.Slug
	}

	for _, pf := range productFiles {
		mdata.ProductFiles = append(mdata.ProductFiles, metadata.ProductFile{
			ID:           pf.ID,
			File:         pf.Name,
			Description:  pf.Description,
			AWSObjectKey: pf.AWSObjectKey,
			FileType:     pf.FileType,
			FileVersion:  pf.FileVersion,
			MD5:          pf.MD5,
		})
	}

	for _, d := range releaseDependencies {
		mdata.Dependencies = append(mdata.Dependencies, metadata.Dependency{
			Release: metadata.DependentRelease{
				ID:      d.Release.ID,
				Version: d.Release.Version,
				Product: metadata.Product{
					ID:   d.Release.Product.ID,
					Name: d.Release.Product.Name,
				},
			},
		})
	}

	for _, d := range releaseUpgradePaths {
		mdata.UpgradePaths = append(mdata.UpgradePaths, metadata.UpgradePath{
			ID:      d.Release.ID,
			Version: d.Release.Version,
		})
	}

	c.logger.Println("Writing metadata files")

	err = c.fileWriter.WriteVersionFile(versionWithETag)
	if err != nil {
		return concourse.InResponse{}, err
	}

	err = c.fileWriter.WriteMetadataYAMLFile(mdata)
	if err != nil {
		return concourse.InResponse{}, err
	}

	err = c.fileWriter.WriteMetadataJSONFile(mdata)
	if err != nil {
		return concourse.InResponse{}, err
	}

	concourseMetadata := c.addReleaseMetadata([]concourse.Metadata{}, release)

	out := concourse.InResponse{
		Version: concourse.Version{
			ProductVersion: versionWithETag,
		},
		Metadata: concourseMetadata,
	}

	return out, nil
}

func (c InCommand) getProductFiles(
	productSlug string,
	releaseID int,
) ([]pivnet.ProductFile, error) {
	productFiles, err := c.pivnetClient.GetProductFilesForRelease(productSlug, releaseID)
	if err != nil {
		return nil, err
	}

	// Get individual product files to obtain metadata that isn't found
	// in the endpoint for all product files.
	for i, p := range productFiles {
		productFiles[i], err = c.pivnetClient.GetProductFile(productSlug, releaseID, p.ID)
		if err != nil {
			return nil, err
		}
	}

	return productFiles, nil
}

func (c InCommand) downloadFiles(
	globs []string,
	productFiles []pivnet.ProductFile,
	productSlug string,
	releaseID int,
) error {
	c.logger.Println("Getting download links")

	downloadLinks := c.filter.DownloadLinks(productFiles)

	c.logger.Println("Filtering download links by glob")

	// if globs were not provided, do not fail if we do not match anything
	failOnNoMatch := (globs != nil)

	var err error
	downloadLinks, err = c.filter.DownloadLinksByGlobs(
		downloadLinks,
		globs,
		failOnNoMatch,
	)
	if err != nil {
		return err
	}

	c.logger.Println("Downloading filtered files")

	files, err := c.downloader.Download(downloadLinks)
	if err != nil {
		return err
	}

	fileMD5s := map[string]string{}
	for _, p := range productFiles {
		parts := strings.Split(p.AWSObjectKey, "/")

		if len(parts) < 1 {
			panic("not enough components to form filename")
		}

		fileName := parts[len(parts)-1]

		if fileName == "" {
			panic("empty file name")
		}

		if p.FileType == pivnet.FileTypeSoftware {
			fileMD5s[fileName] = p.MD5
		}
	}

	err = c.compareMD5s(files, fileMD5s)
	if err != nil {
		return err
	}

	return nil
}

func (c InCommand) addReleaseMetadata(
	concourseMetadata []concourse.Metadata,
	release pivnet.Release,
) []concourse.Metadata {
	cmdata := append(concourseMetadata,
		concourse.Metadata{Name: "version", Value: release.Version},
		concourse.Metadata{Name: "release_type", Value: string(release.ReleaseType)},
		concourse.Metadata{Name: "release_date", Value: release.ReleaseDate},
		concourse.Metadata{Name: "description", Value: release.Description},
		concourse.Metadata{Name: "release_notes_url", Value: release.ReleaseNotesURL},
		concourse.Metadata{Name: "availability", Value: release.Availability},
		concourse.Metadata{Name: "controlled", Value: fmt.Sprintf("%t", release.Controlled)},
		concourse.Metadata{Name: "eccn", Value: release.ECCN},
		concourse.Metadata{Name: "license_exception", Value: release.LicenseException},
		concourse.Metadata{Name: "end_of_support_date", Value: release.EndOfSupportDate},
		concourse.Metadata{Name: "end_of_guidance_date", Value: release.EndOfGuidanceDate},
		concourse.Metadata{Name: "end_of_availability_date", Value: release.EndOfAvailabilityDate},
	)

	if release.EULA != nil {
		concourseMetadata = append(concourseMetadata,
			concourse.Metadata{Name: "eula_slug", Value: release.EULA.Slug},
		)
	}

	return cmdata
}

func (c InCommand) compareMD5s(filepaths []string, expectedMD5s map[string]string) error {
	c.logger.Println("Calcuating MD5 for downloaded files")

	for _, downloadPath := range filepaths {
		_, f := filepath.Split(downloadPath)

		md5, err := c.fileSummer.SumFile(downloadPath)
		if err != nil {
			return err
		}

		expectedMD5 := expectedMD5s[f]
		if expectedMD5 != "" {
			if md5 != expectedMD5 {
				c.logger.Printf(
					"Failed MD5 comparison for file: %s. Expected %s, got %s\n",
					f,
					expectedMD5,
					md5,
				)
				return errors.New("failed comparison")
			}

			c.logger.Println("MD5 for downloaded file matched")
		}
	}

	return nil
}
<commit_msg>Improve logging for failed MD5 comparisons.<commit_after>package in

import (
	"errors"
	"fmt"
	"log"
	"path/filepath"
	"strings"

	pivnet "github.com/pivotal-cf/go-pivnet"
	"github.com/pivotal-cf/pivnet-resource/concourse"
	"github.com/pivotal-cf/pivnet-resource/metadata"
	"github.com/pivotal-cf/pivnet-resource/versions"
)

//go:generate counterfeiter --fake-name FakeFilter . filter
type filter interface {
	DownloadLinksByGlobs(downloadLinks map[string]string, glob []string, failOnNoMatch bool) (map[string]string, error)
	DownloadLinks(p []pivnet.ProductFile) map[string]string
}

//go:generate counterfeiter --fake-name FakeDownloader . downloader
type downloader interface {
	Download(downloadLinks map[string]string) ([]string, error)
}

//go:generate counterfeiter --fake-name FakeFileSummer . fileSummer
type fileSummer interface {
	SumFile(filepath string) (string, error)
}

//go:generate counterfeiter --fake-name FakeFileWriter . fileWriter
type fileWriter interface {
	WriteMetadataJSONFile(mdata metadata.Metadata) error
	WriteMetadataYAMLFile(mdata metadata.Metadata) error
	WriteVersionFile(versionWithETag string) error
}

//go:generate counterfeiter --fake-name FakePivnetClient . pivnetClient
type pivnetClient interface {
	GetRelease(productSlug string, productVersion string) (pivnet.Release, error)
	AcceptEULA(productSlug string, releaseID int) error
	GetProductFilesForRelease(productSlug string, releaseID int) ([]pivnet.ProductFile, error)
	GetProductFile(productSlug string, releaseID int, productFileID int) (pivnet.ProductFile, error)
	ReleaseDependencies(productSlug string, releaseID int) ([]pivnet.ReleaseDependency, error)
	ReleaseUpgradePaths(productSlug string, releaseID int) ([]pivnet.ReleaseUpgradePath, error)
}

type InCommand struct {
	logger       *log.Logger
	downloadDir  string
	pivnetClient pivnetClient
	filter       filter
	downloader   downloader
	fileSummer   fileSummer
	fileWriter   fileWriter
}

func NewInCommand(
	logger *log.Logger,
	pivnetClient pivnetClient,
	filter filter,
	downloader downloader,
	fileSummer fileSummer,
	fileWriter fileWriter,
) *InCommand {
	return &InCommand{
		logger:       logger,
		pivnetClient: pivnetClient,
		filter:       filter,
		downloader:   downloader,
		fileSummer:   fileSummer,
		fileWriter:   fileWriter,
	}
}

func (c *InCommand) Run(input concourse.InRequest) (concourse.InResponse, error) {
	productSlug := input.Source.ProductSlug

	productVersion, etag, err := versions.SplitIntoVersionAndETag(input.Version.ProductVersion)
	if err != nil {
		c.logger.Println("Parsing of etag failed; continuing without it")
		productVersion = input.Version.ProductVersion
	}

	c.logger.Printf(
		"Getting release for product_slug %s and product_version %s",
		productSlug,
		productVersion,
	)

	release, err := c.pivnetClient.GetRelease(productSlug, productVersion)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Printf("Accepting EULA for release_id %d", release.ID)

	err = c.pivnetClient.AcceptEULA(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Getting product files")

	productFiles, err := c.getProductFiles(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Getting release dependencies")

	releaseDependencies, err := c.pivnetClient.ReleaseDependencies(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Getting release upgrade paths")

	releaseUpgradePaths, err := c.pivnetClient.ReleaseUpgradePaths(productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Downloading files")

	err = c.downloadFiles(input.Params.Globs, productFiles, productSlug, release.ID)
	if err != nil {
		return concourse.InResponse{}, err
	}

	c.logger.Println("Creating metadata")

	versionWithETag, err := versions.CombineVersionAndETag(productVersion, etag)

	mdata := metadata.Metadata{
		Release: &metadata.Release{
			Version:               release.Version,
			ReleaseType:           string(release.ReleaseType),
			ReleaseDate:           release.ReleaseDate,
			Description:           release.Description,
			ReleaseNotesURL:       release.ReleaseNotesURL,
			Availability:          release.Availability,
			Controlled:            release.Controlled,
			ECCN:                  release.ECCN,
			LicenseException:      release.LicenseException,
			EndOfSupportDate:      release.EndOfSupportDate,
			EndOfGuidanceDate:     release.EndOfGuidanceDate,
			EndOfAvailabilityDate: release.EndOfAvailabilityDate,
		},
	}

	if release.EULA != nil {
		mdata.Release.EULASlug = release.EULA.Slug
	}

	for _, pf := range productFiles {
		mdata.ProductFiles = append(mdata.ProductFiles, metadata.ProductFile{
			ID:           pf.ID,
			File:         pf.Name,
			Description:  pf.Description,
			AWSObjectKey: pf.AWSObjectKey,
			FileType:     pf.FileType,
			FileVersion:  pf.FileVersion,
			MD5:          pf.MD5,
		})
	}

	for _, d := range releaseDependencies {
		mdata.Dependencies = append(mdata.Dependencies, metadata.Dependency{
			Release: metadata.DependentRelease{
				ID:      d.Release.ID,
				Version: d.Release.Version,
				Product: metadata.Product{
					ID:   d.Release.Product.ID,
					Name: d.Release.Product.Name,
				},
			},
		})
	}

	for _, d := range releaseUpgradePaths {
		mdata.UpgradePaths = append(mdata.UpgradePaths, metadata.UpgradePath{
			ID:      d.Release.ID,
			Version: d.Release.Version,
		})
	}

	c.logger.Println("Writing metadata files")

	err = c.fileWriter.WriteVersionFile(versionWithETag)
	if err != nil {
		return concourse.InResponse{}, err
	}

	err = c.fileWriter.WriteMetadataYAMLFile(mdata)
	if err != nil {
		return concourse.InResponse{}, err
	}

	err = c.fileWriter.WriteMetadataJSONFile(mdata)
	if err != nil {
		return concourse.InResponse{}, err
	}

	concourseMetadata := c.addReleaseMetadata([]concourse.Metadata{}, release)

	out := concourse.InResponse{
		Version: concourse.Version{
			ProductVersion: versionWithETag,
		},
		Metadata: concourseMetadata,
	}

	return out, nil
}

func (c InCommand) getProductFiles(
	productSlug string,
	releaseID int,
) ([]pivnet.ProductFile, error) {
	productFiles, err := c.pivnetClient.GetProductFilesForRelease(productSlug, releaseID)
	if err != nil {
		return nil, err
	}

	// Get individual product files to obtain metadata that isn't found
	// in the endpoint for all product files.
	for i, p := range productFiles {
		productFiles[i], err = c.pivnetClient.GetProductFile(productSlug, releaseID, p.ID)
		if err != nil {
			return nil, err
		}
	}

	return productFiles, nil
}

func (c InCommand) downloadFiles(
	globs []string,
	productFiles []pivnet.ProductFile,
	productSlug string,
	releaseID int,
) error {
	c.logger.Println("Getting download links")

	downloadLinks := c.filter.DownloadLinks(productFiles)

	c.logger.Println("Filtering download links by glob")

	// It is acceptable to match nothing if globs were not provided.
	// This is the use-case when there are no files on pivnet and the pipeline
	// does not specify anything.
	failOnNoMatch := (globs != nil)

	var err error
	downloadLinks, err = c.filter.DownloadLinksByGlobs(
		downloadLinks,
		globs,
		failOnNoMatch,
	)
	if err != nil {
		return err
	}

	c.logger.Println("Downloading filtered files")

	files, err := c.downloader.Download(downloadLinks)
	if err != nil {
		return err
	}

	fileMD5s := map[string]string{}
	for _, p := range productFiles {
		parts := strings.Split(p.AWSObjectKey, "/")

		if len(parts) < 1 {
			panic("not enough components to form filename")
		}

		fileName := parts[len(parts)-1]

		if fileName == "" {
			panic("empty file name")
		}

		if p.FileType == pivnet.FileTypeSoftware {
			fileMD5s[fileName] = p.MD5
		}
	}

	err = c.compareMD5s(files, fileMD5s)
	if err != nil {
		return err
	}

	return nil
}

func (c InCommand) addReleaseMetadata(
	concourseMetadata []concourse.Metadata,
	release pivnet.Release,
) []concourse.Metadata {
	cmdata := append(concourseMetadata,
		concourse.Metadata{Name: "version", Value: release.Version},
		concourse.Metadata{Name: "release_type", Value: string(release.ReleaseType)},
		concourse.Metadata{Name: "release_date", Value: release.ReleaseDate},
		concourse.Metadata{Name: "description", Value: release.Description},
		concourse.Metadata{Name: "release_notes_url", Value: release.ReleaseNotesURL},
		concourse.Metadata{Name: "availability", Value: release.Availability},
		concourse.Metadata{Name: "controlled", Value: fmt.Sprintf("%t", release.Controlled)},
		concourse.Metadata{Name: "eccn", Value: release.ECCN},
		concourse.Metadata{Name: "license_exception", Value: release.LicenseException},
		concourse.Metadata{Name: "end_of_support_date", Value: release.EndOfSupportDate},
		concourse.Metadata{Name: "end_of_guidance_date", Value: release.EndOfGuidanceDate},
		concourse.Metadata{Name: "end_of_availability_date", Value: release.EndOfAvailabilityDate},
	)

	if release.EULA != nil {
		concourseMetadata = append(concourseMetadata,
			concourse.Metadata{Name: "eula_slug", Value: release.EULA.Slug},
		)
	}

	return cmdata
}

func (c InCommand) compareMD5s(filepaths []string, expectedMD5s map[string]string) error {
	c.logger.Println("Calcuating MD5 for downloaded files")

	for _, downloadPath := range filepaths {
		_, f := filepath.Split(downloadPath)

		actualMD5, err := c.fileSummer.SumFile(downloadPath)
		if err != nil {
			return err
		}

		expectedMD5 := expectedMD5s[f]
		if expectedMD5 != "" && expectedMD5 != actualMD5 {
			c.logger.Println(
				fmt.Sprintf(
					"MD5 comparison failed for downloaded file: '%s'. Expected (from pivnet): '%s' - actual (from file): '%s'\n",
					downloadPath,
					expectedMD5,
					actualMD5,
				),
			)
			return errors.New("failed MD5 comparison")
		}
	}

	c.logger.Println("MD5 matched for all downloaded files")

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"
	"os"
	"time"

	"github.com/ChristopherRabotin/smd"
	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

const (
	r2d = 180 / math.Pi
	d2r = 1 / r2d
)

var (
	minRadius = 300 + smd.Earth.Radius // km
	launch    = time.Date(1989, 10, 8, 0, 0, 0, 0, time.UTC)
	vga1      = time.Date(1990, 2, 10, 0, 0, 0, 0, time.UTC)
	ega1      = time.Date(1990, 12, 10, 0, 0, 0, 0, time.UTC)
	ega2      = time.Date(1992, 12, 9, 12, 0, 0, 0, time.UTC)
	joi       = time.Date(1996, 3, 21, 12, 0, 0, 0, time.UTC)
)

func main() {
	resonance := ega2.Sub(ega1).Hours() / (365.242189 * 24)
	fmt.Printf("%s\t~%f orbits\n", ega2.Sub(ega1), resonance)
	var ViGA2, VfGA1 *mat64.Vector

	fmt.Println("==== QUESTION 1 ====")
	// hwQ 1
	vga1R := mat64.NewVector(3, smd.Venus.HelioOrbit(vga1).R())
	earthAtEGA1 := smd.Earth.HelioOrbit(ega1)
	ega1R := mat64.NewVector(3, earthAtEGA1.R())
	_, VfGA1, _, _ = smd.Lambert(vga1R, ega1R, ega1.Sub(vga1), smd.TTypeAuto, smd.Sun)
	vfloats1 := []float64{VfGA1.At(0, 0), VfGA1.At(1, 0), VfGA1.At(2, 0)}
	ega1Orbit := smd.NewOrbitFromRV(earthAtEGA1.R(), vfloats1, smd.Sun)
	ega1Orbit.ToXCentric(smd.Earth, ega1)
	vInfInGA1 := ega1Orbit.V()
	vInfOutGA1Norm := ega1Orbit.VNorm() // Called OutGA1 because we suppose there was no maneuver during the flyby
	fmt.Printf("%+v\n%f km/s\n", vInfInGA1, vInfOutGA1Norm)
	fmt.Println("==== QUESTION 2 ====")
	// hwQ 2
	earthAtEGA2 := smd.Earth.HelioOrbit(ega2)
	ega2R := mat64.NewVector(3, earthAtEGA2.R())
	joiR := mat64.NewVector(3, smd.Jupiter.HelioOrbit(joi).R())
	ViGA2, _, _, _ = smd.Lambert(ega2R, joiR, joi.Sub(ega2), smd.TTypeAuto, smd.Sun)
	vfloats2 := []float64{ViGA2.At(0, 0), ViGA2.At(1, 0), ViGA2.At(2, 0)}
	ega2Orbit := smd.NewOrbitFromRV(earthAtEGA2.R(), vfloats2, smd.Sun)
	ega2Orbit.ToXCentric(smd.Earth, ega2)
	vInfOutGA2 := ega2Orbit.V()
	vInfOutGA2Norm := ega2Orbit.VNorm()
	fmt.Printf("%+v\n%f km/s\n", vInfOutGA2, vInfOutGA2Norm)

	fmt.Println("==== QUESTION 3 ====")
	aResonance := math.Pow(smd.Sun.GM()*math.Pow(resonance*earthAtEGA1.Period().Seconds()/(2*math.Pi), 2), 1/3.)
	VScSunNorm := math.Sqrt(smd.Sun.GM() * ((2 / earthAtEGA1.RNorm()) - 1/aResonance))
	// Compute angle theta for EGA1
	theta := math.Acos((math.Pow(VScSunNorm, 2) - math.Pow(vInfOutGA1Norm, 2) - math.Pow(earthAtEGA1.VNorm(), 2)) / (-2 * vInfOutGA1Norm * earthAtEGA1.VNorm()))
	fmt.Printf("theta = %f\n", theta*r2d)
	// Compute the VNC2ECI DCMs for EGA1 and EGA2.
	V := unit(earthAtEGA1.V())
	N := unit(earthAtEGA1.H())
	C := cross(V, N)
	dcmVal := make([]float64, 9)
	for i := 0; i < 3; i++ {
		dcmVal[i] = V[i]
		dcmVal[i+3] = N[i]
		dcmVal[i+6] = C[i]
	}
	DCM := mat64.NewDense(3, 3, dcmVal)
	data := "psi\trP1\trP2\n"
	step := (2 * math.Pi) / 10000
	// Print when both become higher than minRadius.
	rpsOkay := false
	for ψ := step; ψ < 2*math.Pi; ψ += step {
		sψ, cψ := math.Sincos(ψ)
		vInfOutGA1VNC := []float64{vInfOutGA1Norm * math.Cos(math.Pi-theta), vInfOutGA1Norm * math.Sin(math.Pi-theta) * cψ, -vInfOutGA1Norm * math.Sin(math.Pi-theta) * sψ}
		vInfOutGA1 := MxV33(DCM, vInfOutGA1VNC)
		_, rP1, _, _, _, _ := smd.GAFromVinf(vInfInGA1, vInfOutGA1, smd.Earth)

		vInfInGA2 := make([]float64, 3)
		for i := 0; i < 3; i++ {
			vInfInGA2[i] = vInfOutGA1[i] + ega1Orbit.V()[i] - ega2Orbit.V()[i]
		}
		_, rP2, _, _, _, _ := smd.GAFromVinf(vInfOutGA1, vInfInGA2, smd.Earth)
		data += fmt.Sprintf("%f\t%f\t%f\n", ψ*r2d, rP1, rP2)
		if !rpsOkay && rP1 > minRadius && rP2 > minRadius {
			rpsOkay = true
			fmt.Printf("[OK ] ψ=%.6f\trP1=%.3f km\trP2=%.3f km\n", ψ*r2d, rP1, rP2)
		}
		if rpsOkay && (rP1 < minRadius || rP2 < minRadius) {
			rpsOkay = false
			fmt.Printf("[NOK] ψ=%.6f\trP1=%.3f km\trP2=%.3f km\n", ψ*r2d, rP1, rP2)
		}
	}
	// Export data
	f, err := os.Create("./q3.tsv")
	if err != nil {
		panic(err)
	}
	f.WriteString(data)
	f.Close()
}

// Unshamefully copied from smd/math.go
func cross(a, b []float64) []float64 {
	return []float64{a[1]*b[2] - a[2]*b[1],
		a[2]*b[0] - a[0]*b[2],
		a[0]*b[1] - a[1]*b[0]} // Cross product R x V.
}

// norm returns the norm of a given vector which is supposed to be 3x1.
func norm(v []float64) float64 {
	return math.Sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])
}

// unit returns the unit vector of a given vector.
func unit(a []float64) (b []float64) {
	n := norm(a)
	if floats.EqualWithinAbs(n, 0, 1e-12) {
		return []float64{0, 0, 0}
	}
	b = make([]float64, len(a))
	for i, val := range a {
		b[i] = val / n
	}
	return
}

// MxV33 multiplies a matrix with a vector. Note that there is no dimension check!
func MxV33(m *mat64.Dense, v []float64) (o []float64) {
	vVec := mat64.NewVector(len(v), v)
	var rVec mat64.Vector
	rVec.MulVec(m, vVec)
	return []float64{rVec.At(0, 0), rVec.At(1, 0), rVec.At(2, 0)}
}
<commit_msg>Continued galileo b plane<commit_after>package main

import (
	"fmt"
	"math"
	"time"

	"github.com/ChristopherRabotin/smd"
	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

const (
	r2d = 180 / math.Pi
	d2r = 1 / r2d
)

var (
	minRadius = 300 + smd.Earth.Radius // km
	launch    = time.Date(1989, 10, 8, 0, 0, 0, 0, time.UTC)
	vga1      = time.Date(1990, 2, 10, 0, 0, 0, 0, time.UTC)
	ega1      = time.Date(1990, 12, 10, 0, 0, 0, 0, time.UTC)
	ega2      = time.Date(1992, 12, 9, 12, 0, 0, 0, time.UTC)
	joi       = time.Date(1996, 3, 21, 12, 0, 0, 0, time.UTC)
)

func main() {
	resonance := ega2.Sub(ega1).Hours() / (365.242189 * 24)
	fmt.Printf("%s\t~%f orbits\n", ega2.Sub(ega1), resonance)

	fmt.Println("==== Launch -> VGA ====")
	// hwQ 1
	earthAtLaunch := smd.Earth.HelioOrbit(launch)
	eLaunchR := mat64.NewVector(3, earthAtLaunch.R())
	venusAtVGA := smd.Venus.HelioOrbit(vga1)
	vga1R := mat64.NewVector(3, venusAtVGA.R())
	_, VfVGA, _, _ := smd.Lambert(eLaunchR, vga1R, vga1.Sub(launch), smd.TTypeAuto, smd.Sun)
	vfloatsVGA := []float64{VfVGA.At(0, 0), VfVGA.At(1, 0), VfVGA.At(2, 0)}
	scOrbitAtVGAIn := smd.NewOrbitFromRV(venusAtVGA.R(), vfloatsVGA, smd.Sun)
	scOrbitAtVGAIn.ToXCentric(smd.Venus, vga1)

	fmt.Println("==== VGA -> EGA1 ====")
	earthAtEGA1 := smd.Earth.HelioOrbit(ega1)
	ega1R := mat64.NewVector(3, earthAtEGA1.R())
	ViVGA, VfGA1, _, _ := smd.Lambert(vga1R, ega1R, ega1.Sub(vga1), smd.TTypeAuto, smd.Sun)
	scOrbitAtVGAOut := smd.NewOrbitFromRV(venusAtVGA.R(), []float64{ViVGA.At(0, 0), ViVGA.At(1, 0), ViVGA.At(2, 0)}, smd.Sun)
	scOrbitAtVGAOut.ToXCentric(smd.Venus, vga1)
	// Okay, we have all the info for Venus, let's compute stuff.
	_, rPVenus, bTVenus, bRVenus, _, _ := smd.GAFromVinf(scOrbitAtVGAIn.V(), scOrbitAtVGAOut.V(), smd.Venus)
	fmt.Printf("==== VENUS INFO ====\nrP=%f km\tBt=%f km\tBr=%f\nVin=%f\tVout=%f\nDelta=%f\n\n", rPVenus, bTVenus, bRVenus, scOrbitAtVGAIn.VNorm(), scOrbitAtVGAOut.VNorm(), scOrbitAtVGAOut.VNorm()-scOrbitAtVGAIn.VNorm())

	scOrbitAtEGA1 := smd.NewOrbitFromRV(earthAtEGA1.R(), []float64{VfGA1.At(0, 0), VfGA1.At(1, 0), VfGA1.At(2, 0)}, smd.Sun)
	scOrbitAtEGA1.ToXCentric(smd.Earth, ega1)
	vInfInGA1 := scOrbitAtEGA1.V()
	vInfOutGA1Norm := scOrbitAtEGA1.VNorm() // Called OutGA1 because we suppose there was no maneuver during the flyby

	fmt.Println("==== EGA2 -> JOI ====")
	// hwQ 2
	earthAtEGA2 := smd.Earth.HelioOrbit(ega2)
	ega2R := mat64.NewVector(3, earthAtEGA2.R())
	joiR := mat64.NewVector(3, smd.Jupiter.HelioOrbit(joi).R())
	ViGA2, _, _, _ := smd.Lambert(ega2R, joiR, joi.Sub(ega2), smd.TTypeAuto, smd.Sun)
	scOrbitAtEGA2 := smd.NewOrbitFromRV(earthAtEGA2.R(), []float64{ViGA2.At(0, 0), ViGA2.At(1, 0), ViGA2.At(2, 0)}, smd.Sun)
	scOrbitAtEGA2.ToXCentric(smd.Earth, ega2)
	vInfOutGA2 := scOrbitAtEGA2.V()
	vInfOutGA2Norm := scOrbitAtEGA2.VNorm()
	fmt.Printf("%+v\n%f km/s\n", vInfOutGA2, vInfOutGA2Norm)

	fmt.Println("==== Earth resonnance ====")
	aResonance := math.Pow(smd.Sun.GM()*math.Pow(resonance*earthAtEGA1.Period().Seconds()/(2*math.Pi), 2), 1/3.)
	VScSunNorm := math.Sqrt(smd.Sun.GM() * ((2 / earthAtEGA1.RNorm()) - 1/aResonance))
	// Compute angle theta for EGA1
	theta := math.Acos((math.Pow(VScSunNorm, 2) - math.Pow(vInfOutGA1Norm, 2) - math.Pow(earthAtEGA1.VNorm(), 2)) / (-2 * vInfOutGA1Norm * earthAtEGA1.VNorm()))
	fmt.Printf("theta = %f\n", theta*r2d)
	// Compute the VNC2ECI DCMs for EGA1 and EGA2.
	V := unit(earthAtEGA1.V())
	N := unit(earthAtEGA1.H())
	C := cross(V, N)
	dcmVal := make([]float64, 9)
	for i := 0; i < 3; i++ {
		dcmVal[i] = V[i]
		dcmVal[i+3] = N[i]
		dcmVal[i+6] = C[i]
	}
	DCM := mat64.NewDense(3, 3, dcmVal)
	// Print when both become higher than minRadius.

	ψ := 165.924 * d2r

	sψ, cψ := math.Sincos(ψ)
	vInfOutGA1VNC := []float64{vInfOutGA1Norm * math.Cos(math.Pi-theta), vInfOutGA1Norm * math.Sin(math.Pi-theta) * cψ, -vInfOutGA1Norm * math.Sin(math.Pi-theta) * sψ}
	vInfOutGA1 := MxV33(DCM, vInfOutGA1VNC)
	_, rPEGA1, bTEGA1, bREGA1, _, _ := smd.GAFromVinf(vInfInGA1, vInfOutGA1, smd.Earth)
	fmt.Printf("==== EGA1 INFO ====\nrP=%f km\tBt=%f km\tBr=%f\nVin=%f\tVout=%f\nDelta=%f\n\n", rPEGA1, bTEGA1, bREGA1, norm(vInfInGA1), norm(vInfOutGA1), norm(vInfOutGA1)-norm(vInfInGA1))

	vInfInGA2 := make([]float64, 3)
	for i := 0; i < 3; i++ {
		vInfInGA2[i] = vInfOutGA1[i] + scOrbitAtEGA1.V()[i] - scOrbitAtEGA2.V()[i]
	}
	_, rPEGA2, bTEGA2, bREGA2, _, _ := smd.GAFromVinf(vInfOutGA1, vInfInGA2, smd.Earth)
	fmt.Printf("==== EGA2 INFO ====\nrP=%f km\tBt=%f km\tBr=%f\nVin=%f\tVout=%f\nDelta=%f\n\n", rPEGA2, bTEGA2, bREGA2, norm(vInfOutGA1), norm(vInfInGA2), norm(vInfInGA2)-norm(vInfOutGA1))

}

// Unshamefully copied from smd/math.go
func cross(a, b []float64) []float64 {
	return []float64{a[1]*b[2] - a[2]*b[1],
		a[2]*b[0] - a[0]*b[2],
		a[0]*b[1] - a[1]*b[0]} // Cross product R x V.
}

// norm returns the norm of a given vector which is supposed to be 3x1.
func norm(v []float64) float64 {
	return math.Sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])
}

// unit returns the unit vector of a given vector.
func unit(a []float64) (b []float64) {
	n := norm(a)
	if floats.EqualWithinAbs(n, 0, 1e-12) {
		return []float64{0, 0, 0}
	}
	b = make([]float64, len(a))
	for i, val := range a {
		b[i] = val / n
	}
	return
}

// MxV33 multiplies a matrix with a vector. Note that there is no dimension check!
func MxV33(m *mat64.Dense, v []float64) (o []float64) {
	vVec := mat64.NewVector(len(v), v)
	var rVec mat64.Vector
	rVec.MulVec(m, vVec)
	return []float64{rVec.At(0, 0), rVec.At(1, 0), rVec.At(2, 0)}
}
<|endoftext|>


<commit_before>package goinsta

import (
	"encoding/json"
	"fmt"
)

// Item represents media items
type Item struct {
	TakenAt                      int          `json:"taken_at"`
	ID                           int64        `json:"pk"`
	IDStr                        string       `json:"id"`
	DeviceTimestamp              int64        `json:"device_timestamp"`
	MediaType                    int          `json:"media_type"`
	Code                         string       `json:"code"`
	ClientCacheKey               string       `json:"client_cache_key"`
	FilterType                   int          `json:"filter_type"`
	CarouselParentID             string       `json:"carousel_parent_id"`
	CarouselMedia                []Item       `json:"carousel_media,omitempty"`
	User                         User         `json:"user"`
	CanViewerReshare             bool         `json:"can_viewer_reshare"`
	Caption                      Caption      `json:"caption"`
	CaptionIsEdited              bool         `json:"caption_is_edited"`
	Likes                        int          `json:"like_count"`
	HasLiked                     bool         `json:"has_liked"`
	TopLikers                    []User       `json:"top_likers"`
	CommentLikesEnabled          bool         `json:"comment_likes_enabled"`
	CommentThreadingEnabled      bool         `json:"comment_threading_enabled"`
	HasMoreComments              bool         `json:"has_more_comments"`
	MaxNumVisiblePreviewComments int          `json:"max_num_visible_preview_comments"`
	PreviewComments              []string     `json:"preview_comments"`
	CommentCount                 int          `json:"comment_count"`
	PhotoOfYou                   bool         `json:"photo_of_you"`
	Usertags                     Tag          `json:"usertags,omitempty"`
	FbUserTags                   Tag          `json:"fb_user_tags"`
	CanViewerSave                bool         `json:"can_viewer_save"`
	OrganicTrackingToken         string       `json:"organic_tracking_token"`
	ImageVersions                ImageVersion `json:"image_versions2,omitempty"`
	OriginalWidth                int          `json:"original_width,omitempty"`
	OriginalHeight               int          `json:"original_height,omitempty"`
	CaptionIsEdited              bool         `json:"caption_is_edited"`
	ImportedTakenAt              int          `json:"imported_taken_at,omitempty"`

	// Only for stories
	StoryEvents              []interface{} `json:"story_events"`
	StoryHashtags            []interface{} `json:"story_hashtags"`
	StoryPolls               []interface{} `json:"story_polls"`
	StoryFeedMedia           []interface{} `json:"story_feed_media"`
	StorySoundOn             []interface{} `json:"story_sound_on"`
	CreativeConfig           interface{}   `json:"creative_config"`
	StoryLocations           []interface{} `json:"story_locations"`
	StorySliders             []interface{} `json:"story_sliders"`
	StoryQuestions           []interface{} `json:"story_questions"`
	StoryProductItems        []interface{} `json:"story_product_items"`
	SupportsReelReactions    bool          `json:"supports_reel_reactions"`
	ShowOneTapFbShareTooltip bool          `json:"show_one_tap_fb_share_tooltip"`
	HasSharedToFb            int           `json:"has_shared_to_fb"`
	Mentions                 []Mentions
	VideoVersions            []VideoVersion `json:"video_versions,omitempty"`
	HasAudio                 bool           `json:"has_audio,omitempty"`
	VideoDuration            float64        `json:"video_duration,omitempty"`
	IsDashEligible           int            `json:"is_dash_eligible,omitempty"`
	VideoDashManifest        string         `json:"video_dash_manifest,omitempty"`
	NumberOfQualities        int            `json:"number_of_qualities,omitempty"`
}

type Media interface {
	Next() error
}

type StoryMedia struct {
	inst     *Instagram
	endpoint string
	uid      int64

	ID int `json:"id"`

	LatestReelMedia int      `json:"latest_reel_media"`
	ExpiringAt      int      `json:"expiring_at"`
	Seen            float64  `json:"seen"`
	CanReply        bool     `json:"can_reply"`
	CanReshare      bool     `json:"can_reshare"`
	ReelType        string   `json:"reel_type"`
	User            User     `json:"user"`
	Items           []Item   `json:"items"`
	ReelMentions    []string `json:"reel_mentions"`
	PrefetchCount   int      `json:"prefetch_count"`
	HasBestiesMedia int      `json:"has_besties_media"`
	Status          string   `json:"status"`
}

// Next allows to paginate after calling:
// User.Stories
func (media *StoryMedia) Next() (err error) {
	var body []byte
	insta := media.inst
	endpoint := media.endpoint

	body, err = insta.sendSimpleRequest(
		endpoint, media.uid,
	)
	if err == nil {
		m := StoryMedia{}
		err = json.Unmarshal(body, &m)
		if err == nil {
			*media = m
			media.inst = insta
			media.endpoint = endpoint
			// TODO check NextID media
		}
	}
	return err
}

// Media represent a set of media items
type FeedMedia struct {
	inst *Instagram

	uid      int64
	endpoint string

	Items               []Item `json:"items"`
	NumResults          int    `json:"num_results"`
	MoreAvailable       bool   `json:"more_available"`
	AutoLoadMoreEnabled bool   `json:"auto_load_more_enabled"`
	Status              string `json:"status"`
	NextID              string `json:"next_max_id"`
}

// Next allows to paginate after calling:
// User.Feed
//
// returns ErrNoMore when list reach the end.
func (media *FeedMedia) Next() (err error) {
	var body []byte
	insta := media.inst
	endpoint := media.endpoint

	body, err = insta.sendRequest(
		&reqOptions{
			Endpoint: fmt.Sprintf(endpoint, media.uid),
			Query: map[string]string{
				"max_id":         media.NextID,
				"rank_token":     insta.rankToken,
				"min_timestamp":  "",
				"ranked_content": "true",
			},
		},
	)
	if err == nil {
		m := FeedMedia{}
		err = json.Unmarshal(body, &m)
		if err == nil {
			*media = m
			media.inst = insta
			media.endpoint = endpoint
			if m.NextID == "" || m.MoreAvailable {
				err = ErrNoMore
			}
		}
	}
	return err
}
<commit_msg>Added default error (because story media does not have pagination)<commit_after>package goinsta

import (
	"encoding/json"
	"fmt"
)

// Item represents media items
type Item struct {
	TakenAt                      int          `json:"taken_at"`
	ID                           int64        `json:"pk"`
	IDStr                        string       `json:"id"`
	DeviceTimestamp              int64        `json:"device_timestamp"`
	MediaType                    int          `json:"media_type"`
	Code                         string       `json:"code"`
	ClientCacheKey               string       `json:"client_cache_key"`
	FilterType                   int          `json:"filter_type"`
	CarouselParentID             string       `json:"carousel_parent_id"`
	CarouselMedia                []Item       `json:"carousel_media,omitempty"`
	User                         User         `json:"user"`
	CanViewerReshare             bool         `json:"can_viewer_reshare"`
	Caption                      Caption      `json:"caption"`
	CaptionIsEdited              bool         `json:"caption_is_edited"`
	Likes                        int          `json:"like_count"`
	HasLiked                     bool         `json:"has_liked"`
	TopLikers                    []User       `json:"top_likers"`
	CommentLikesEnabled          bool         `json:"comment_likes_enabled"`
	CommentThreadingEnabled      bool         `json:"comment_threading_enabled"`
	HasMoreComments              bool         `json:"has_more_comments"`
	MaxNumVisiblePreviewComments int          `json:"max_num_visible_preview_comments"`
	PreviewComments              []string     `json:"preview_comments"`
	CommentCount                 int          `json:"comment_count"`
	PhotoOfYou                   bool         `json:"photo_of_you"`
	Usertags                     Tag          `json:"usertags,omitempty"`
	FbUserTags                   Tag          `json:"fb_user_tags"`
	CanViewerSave                bool         `json:"can_viewer_save"`
	OrganicTrackingToken         string       `json:"organic_tracking_token"`
	ImageVersions                ImageVersion `json:"image_versions2,omitempty"`
	OriginalWidth                int          `json:"original_width,omitempty"`
	OriginalHeight               int          `json:"original_height,omitempty"`
	CaptionIsEdited              bool         `json:"caption_is_edited"`
	ImportedTakenAt              int          `json:"imported_taken_at,omitempty"`

	// Only for stories
	StoryEvents              []interface{} `json:"story_events"`
	StoryHashtags            []interface{} `json:"story_hashtags"`
	StoryPolls               []interface{} `json:"story_polls"`
	StoryFeedMedia           []interface{} `json:"story_feed_media"`
	StorySoundOn             []interface{} `json:"story_sound_on"`
	CreativeConfig           interface{}   `json:"creative_config"`
	StoryLocations           []interface{} `json:"story_locations"`
	StorySliders             []interface{} `json:"story_sliders"`
	StoryQuestions           []interface{} `json:"story_questions"`
	StoryProductItems        []interface{} `json:"story_product_items"`
	SupportsReelReactions    bool          `json:"supports_reel_reactions"`
	ShowOneTapFbShareTooltip bool          `json:"show_one_tap_fb_share_tooltip"`
	HasSharedToFb            int           `json:"has_shared_to_fb"`
	Mentions                 []Mentions
	VideoVersions            []VideoVersion `json:"video_versions,omitempty"`
	HasAudio                 bool           `json:"has_audio,omitempty"`
	VideoDuration            float64        `json:"video_duration,omitempty"`
	IsDashEligible           int            `json:"is_dash_eligible,omitempty"`
	VideoDashManifest        string         `json:"video_dash_manifest,omitempty"`
	NumberOfQualities        int            `json:"number_of_qualities,omitempty"`
}

type Media interface {
	Next() error
}

type StoryMedia struct {
	inst     *Instagram
	endpoint string
	uid      int64

	ID int `json:"id"`

	LatestReelMedia int      `json:"latest_reel_media"`
	ExpiringAt      int      `json:"expiring_at"`
	Seen            float64  `json:"seen"`
	CanReply        bool     `json:"can_reply"`
	CanReshare      bool     `json:"can_reshare"`
	ReelType        string   `json:"reel_type"`
	User            User     `json:"user"`
	Items           []Item   `json:"items"`
	ReelMentions    []string `json:"reel_mentions"`
	PrefetchCount   int      `json:"prefetch_count"`
	HasBestiesMedia int      `json:"has_besties_media"`
	Status          string   `json:"status"`
}

// Next allows to paginate after calling:
// User.Stories
func (media *StoryMedia) Next() (err error) {
	var body []byte
	insta := media.inst
	endpoint := media.endpoint

	body, err = insta.sendSimpleRequest(
		endpoint, media.uid,
	)
	if err == nil {
		m := StoryMedia{}
		err = json.Unmarshal(body, &m)
		if err == nil {
			err = ErrNoMore
			*media = m
			media.inst = insta
			media.endpoint = endpoint
			// TODO check NextID media
		}
	}
	return err
}

// Media represent a set of media items
type FeedMedia struct {
	inst *Instagram

	uid      int64
	endpoint string

	Items               []Item `json:"items"`
	NumResults          int    `json:"num_results"`
	MoreAvailable       bool   `json:"more_available"`
	AutoLoadMoreEnabled bool   `json:"auto_load_more_enabled"`
	Status              string `json:"status"`
	NextID              string `json:"next_max_id"`
}

// Next allows to paginate after calling:
// User.Feed
//
// returns ErrNoMore when list reach the end.
func (media *FeedMedia) Next() (err error) {
	var body []byte
	insta := media.inst
	endpoint := media.endpoint

	body, err = insta.sendRequest(
		&reqOptions{
			Endpoint: fmt.Sprintf(endpoint, media.uid),
			Query: map[string]string{
				"max_id":         media.NextID,
				"rank_token":     insta.rankToken,
				"min_timestamp":  "",
				"ranked_content": "true",
			},
		},
	)
	if err == nil {
		m := FeedMedia{}
		err = json.Unmarshal(body, &m)
		if err == nil {
			*media = m
			media.inst = insta
			media.endpoint = endpoint
			if m.NextID == "" || m.MoreAvailable {
				err = ErrNoMore
			}
		}
	}
	return err
}
<|endoftext|>


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

package updater

import (
	"time"

	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/client/go/updater/sources"
)

type UpdateChecker struct {
	updater *Updater
	ui      UI
	ticker  *time.Ticker
	log     logger.Logger
}

func NewUpdateChecker(updater *Updater, ui UI, log logger.Logger) UpdateChecker {
	return UpdateChecker{
		updater: updater,
		ui:      ui,
		log:     log,
	}
}

// Check checks for an update. If not requested (by user) and not forced it will
// exit early if check has already been applied within checkDuration().
func (u *UpdateChecker) Check(force bool, requested bool) error {
	if !requested && !force {
		if lastCheckedPTime := u.updater.config.GetUpdateLastChecked(); lastCheckedPTime > 0 {
			lastChecked := keybase1.FromTime(lastCheckedPTime)
			if time.Now().Before(lastChecked.Add(checkDuration())) {
				u.log.Debug("Already checked: %s", lastChecked)
				return nil
			}
		}
	}

	checkTime := time.Now()
	_, err := u.updater.Update(u.ui, force, requested)
	if err != nil {
		return err
	}

	u.log.Debug("Saving updater last checked: %s", checkTime)
	u.updater.config.SetUpdateLastChecked(keybase1.ToTime(checkTime))
	return nil
}

func (u *UpdateChecker) Start() {
	if u.ticker != nil {
		return
	}
	u.ticker = time.NewTicker(tickDuration())
	go func() {
		for _ = range u.ticker.C {
			u.log.Debug("Checking for update (ticker)")
			err := u.Check(false, false)
			if err != nil {
				u.log.Errorf("Error in update: %s", err)
			}
		}
	}()
}

func (u *UpdateChecker) Stop() {
	u.ticker.Stop()
	u.ticker = nil
}

// checkDuration is how often to check for updates
func checkDuration() time.Duration {
	if sources.IsPrerelease {
		return time.Hour
	}
	return 24 * time.Hour
}

// tickDuration is how often to call check (should be less than checkDuration or snooze min)
func tickDuration() time.Duration {
	if sources.IsPrerelease {
		return 15 * time.Minute
	}
	return time.Hour
}
<commit_msg>Simplify some code<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package updater

import (
	"time"

	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/client/go/updater/sources"
)

type UpdateChecker struct {
	updater *Updater
	ui      UI
	ticker  *time.Ticker
	log     logger.Logger
}

func NewUpdateChecker(updater *Updater, ui UI, log logger.Logger) UpdateChecker {
	return UpdateChecker{
		updater: updater,
		ui:      ui,
		log:     log,
	}
}

// Check checks for an update. If not requested (by user) and not forced it will
// exit early if check has already been applied within checkDuration().
func (u *UpdateChecker) Check(force bool, requested bool) error {
	if !requested && !force {
		if lastCheckedPTime := u.updater.config.GetUpdateLastChecked(); lastCheckedPTime > 0 {
			lastChecked := keybase1.FromTime(lastCheckedPTime)
			if time.Now().Before(lastChecked.Add(checkDuration())) {
				u.log.Debug("Already checked: %s", lastChecked)
				return nil
			}
		}
	}

	checkTime := time.Now()
	_, err := u.updater.Update(u.ui, force, requested)
	if err != nil {
		return err
	}

	u.log.Debug("Saving updater last checked: %s", checkTime)
	u.updater.config.SetUpdateLastChecked(keybase1.ToTime(checkTime))
	return nil
}

func (u *UpdateChecker) Start() {
	if u.ticker != nil {
		return
	}
	u.ticker = time.NewTicker(tickDuration())
	go func() {
		for range u.ticker.C {
			u.log.Debug("Checking for update (ticker)")
			err := u.Check(false, false)
			if err != nil {
				u.log.Errorf("Error in update: %s", err)
			}
		}
	}()
}

func (u *UpdateChecker) Stop() {
	u.ticker.Stop()
	u.ticker = nil
}

// checkDuration is how often to check for updates
func checkDuration() time.Duration {
	if sources.IsPrerelease {
		return time.Hour
	}
	return 24 * time.Hour
}

// tickDuration is how often to call check (should be less than checkDuration or snooze min)
func tickDuration() time.Duration {
	if sources.IsPrerelease {
		return 15 * time.Minute
	}
	return time.Hour
}
<|endoftext|>


<commit_before>package main

import (
  "errors"
  "fmt"
  "strconv"

  "github.com/hyperledger/fabric/core/chaincode/shim"
)
type ShipmentChaincode struct{}

func (t *ShipmentChaincode) Init(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {

  // Create shiptment table
	err := stub.CreateTable("Shipment", []*shim.ColumnDefinition{
		&shim.ColumnDefinition{Name: "id", Type: shim.ColumnDefinition_STRING, Key: true},
		&shim.ColumnDefinition{Name: "status", Type: shim.ColumnDefinition_STRING, Key: false},
	})
	if err != nil {
		return nil, errors.New("Failed creating Shipment table.")
	}

  return nil, nil
}

func (t *ShipmentChaincode) update(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error)
  id := args[0]
  status := args[1]
  ok, err = stub.InsertRow("Shipment", shim.Row{
		Columns: []*shim.Column{
			&shim.Column{Value: &shim.Column_String_{String_: id}},
			&shim.Column{Value: &shim.Column_Bytes{String: status}}},
	})

	if !ok && err == nil {
		return nil, errors.New("Asset was already assigned.")
	}
}

func (t *ShipmentChaincode) Invoke(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
  if function == "update" {
    return t.update(stub,args)
  }
  return nil, errors.New("Received unknown function invocation")

}

func (t *SimpleChaincode) Query(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
  var err error

	if len(args) != 1 {
		myLogger.Debug("Incorrect number of arguments. Expecting name of an asset to query")
		return nil, errors.New("Incorrect number of arguments. Expecting name of an asset to query")
	}

	// Who is the owner of the asset?
	id := args[0]

  var columns []shim.Column
	col1 := shim.Column{Value: &shim.Column_String_{String_: id}}
	columns = append(columns, col1)

	row, err := stub.GetRow("Shipment", columns)
	if err != nil {
		myLogger.Debugf("Failed retriving asset [%s]: [%s]", string(id), err)
		return nil, fmt.Errorf("Failed retriving asset [%s]: [%s]", string(id), err)
	}

	return row.Columns[1], nil
}

func main() {
	err := shim.Start(new(ShipmentChaincode))
	if err != nil {
		fmt.Printf("Error starting Shipment chaincode: %s", err)
	}
}
<commit_msg>added functions to update status<commit_after>package main

import (
  "errors"
  "fmt"
  "strconv"

  "github.com/hyperledger/fabric/core/chaincode/shim"
)
type ShipmentChaincode struct{}

func (t *ShipmentChaincode) Init(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {

  // Create shiptment table
	err := stub.CreateTable("Shipment", []*shim.ColumnDefinition{
		&shim.ColumnDefinition{Name: "id", Type: shim.ColumnDefinition_STRING, Key: true},
		&shim.ColumnDefinition{Name: "status", Type: shim.ColumnDefinition_STRING, Key: false},
	})
	if err != nil {
		return nil, errors.New("Failed creating Shipment table.")
	}

  return nil, nil
}

func (t *ShipmentChaincode) assign(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error){
  id := args[0]
  status := args[1]
  ok, err = stub.InsertRow("Shipment", shim.Row{
		Columns: []*shim.Column{
			&shim.Column{Value: &shim.Column_String_{String_: id}},
			&shim.Column{Value: &shim.Column_Bytes{String: status}}},
	})

	if !ok && err == nil {
		return nil, errors.New("Asset was already assigned.")
	}
}

func (t *ShipmentChaincode) update(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error){
  id := args[0]
  status := args[1]
  err = stub.DeleteRow(
		"Shipment",
		[]shim.Column{shim.Column{Value: &shim.Column_String_{String_: id}}},
	)
	if err != nil {
		return nil, errors.New("Failed deliting row.")
	}

	_, err = stub.InsertRow(
		"Shipment",
		shim.Row{
			Columns: []*shim.Column{
				&shim.Column{Value: &shim.Column_String_{String_: id}},
				&shim.Column{Value: &shim.Column_String{String: status}},
			},
		})
	if err != nil {
		return nil, errors.New("Failed inserting row.")
	}
}

func (t *ShipmentChaincode) Invoke(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
  if function == "update" {
    return t.update(stub,args)
  }
  return nil, errors.New("Received unknown function invocation")

}

func (t *SimpleChaincode) Query(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
  var err error

	if len(args) != 1 {
		myLogger.Debug("Incorrect number of arguments. Expecting name of an asset to query")
		return nil, errors.New("Incorrect number of arguments. Expecting name of an asset to query")
	}

	// Who is the owner of the asset?
	id := args[0]

  var columns []shim.Column
	col1 := shim.Column{Value: &shim.Column_String_{String_: id}}
	columns = append(columns, col1)

	row, err := stub.GetRow("Shipment", columns)
	if err != nil {
		myLogger.Debugf("Failed retriving asset [%s]: [%s]", string(id), err)
		return nil, fmt.Errorf("Failed retriving asset [%s]: [%s]", string(id), err)
	}

	return row.Columns[1], nil
}

func main() {
	err := shim.Start(new(ShipmentChaincode))
	if err != nil {
		fmt.Printf("Error starting Shipment chaincode: %s", err)
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2020 The Syncthing Authors.
//
// 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 https://mozilla.org/MPL/2.0/.

package ur

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/syncthing/syncthing/lib/build"
	"github.com/syncthing/syncthing/lib/config"
	"github.com/syncthing/syncthing/lib/dialer"
	"github.com/syncthing/syncthing/lib/events"
	"github.com/syncthing/syncthing/lib/svcutil"

	"github.com/thejerf/suture/v4"
)

var (
	// When a specific failure first occurs, it is delayed by minDelay. If
	// more of the same failures occurs those are further delayed and
	// aggregated for maxDelay.
	minDelay             = 10 * time.Second
	maxDelay             = time.Minute
	sendTimeout          = time.Minute
	finalSendTimeout     = svcutil.ServiceTimeout / 2
	evChanClosed         = "failure event channel closed"
	invalidEventDataType = "failure event data is not a string"
)

type FailureReport struct {
	FailureData
	Count   int
	Version string
}

type FailureData struct {
	Description string
	Goroutines  string
	Extra       map[string]string
}

func FailureDataWithGoroutines(description string) FailureData {
	var buf *strings.Builder
	pprof.NewProfile("goroutine").WriteTo(buf, 1)
	return FailureData{
		Description: description,
		Goroutines:  buf.String(),
	}
}

type FailureHandler interface {
	suture.Service
	config.Committer
}

func NewFailureHandler(cfg config.Wrapper, evLogger events.Logger) FailureHandler {
	return &failureHandler{
		cfg:      cfg,
		evLogger: evLogger,
		optsChan: make(chan config.OptionsConfiguration),
		buf:      make(map[string]*failureStat),
	}
}

type failureHandler struct {
	cfg      config.Wrapper
	evLogger events.Logger
	optsChan chan config.OptionsConfiguration
	buf      map[string]*failureStat
}

type failureStat struct {
	first, last time.Time
	count       int
	data        FailureData
}

func (h *failureHandler) Serve(ctx context.Context) error {
	cfg := h.cfg.Subscribe(h)
	defer h.cfg.Unsubscribe(h)
	url, sub, evChan := h.applyOpts(cfg.Options, nil)

	var err error
	timer := time.NewTimer(minDelay)
	resetTimer := make(chan struct{})
	for err == nil {
		select {
		case opts := <-h.optsChan:
			url, sub, evChan = h.applyOpts(opts, sub)
		case e, ok := <-evChan:
			if !ok {
				// Just to be safe - shouldn't ever happen, as
				// evChan is set to nil when unsubscribing.
				h.addReport(FailureData{Description: evChanClosed}, time.Now())
				evChan = nil
				continue
			}
			var data FailureData
			switch d := e.Data.(type) {
			case string:
				data.Description = d
			case FailureData:
				data = d
			default:
				// Same here, shouldn't ever happen.
				h.addReport(FailureData{Description: invalidEventDataType}, time.Now())
				continue
			}
			h.addReport(data, e.Time)
		case <-timer.C:
			reports := make([]FailureReport, 0, len(h.buf))
			now := time.Now()
			for descr, stat := range h.buf {
				if now.Sub(stat.last) > minDelay || now.Sub(stat.first) > maxDelay {
					reports = append(reports, newFailureReport(stat))
					delete(h.buf, descr)
				}
			}
			if len(reports) > 0 {
				// Lets keep process events/configs while it might be timing out for a while
				go func() {
					sendFailureReports(ctx, reports, url)
					select {
					case resetTimer <- struct{}{}:
					case <-ctx.Done():
					}
				}()
			} else {
				timer.Reset(minDelay)
			}
		case <-resetTimer:
			timer.Reset(minDelay)
		case <-ctx.Done():
			err = ctx.Err()
		}
	}

	if sub != nil {
		sub.Unsubscribe()
		if len(h.buf) > 0 {
			reports := make([]FailureReport, 0, len(h.buf))
			for _, stat := range h.buf {
				reports = append(reports, newFailureReport(stat))
			}
			timeout, cancel := context.WithTimeout(context.Background(), finalSendTimeout)
			defer cancel()
			sendFailureReports(timeout, reports, url)
		}
	}
	return err
}

func (h *failureHandler) applyOpts(opts config.OptionsConfiguration, sub events.Subscription) (string, events.Subscription, <-chan events.Event) {
	// Sub nil checks just for safety - config updates can be racy.
	url := opts.CRURL + "/failure"
	if opts.URAccepted > 0 {
		if sub == nil {
			sub = h.evLogger.Subscribe(events.Failure)
		}
		return url, sub, sub.C()
	}
	if sub != nil {
		sub.Unsubscribe()
	}
	return url, nil, nil
}

func (h *failureHandler) addReport(data FailureData, evTime time.Time) {
	if stat, ok := h.buf[data.Description]; ok {
		stat.last = evTime
		stat.count++
		return
	}
	h.buf[data.Description] = &failureStat{
		first: evTime,
		last:  evTime,
		count: 1,
		data:  data,
	}
}

func (h *failureHandler) VerifyConfiguration(_, _ config.Configuration) error {
	return nil
}

func (h *failureHandler) CommitConfiguration(from, to config.Configuration) bool {
	if from.Options.CREnabled != to.Options.CREnabled || from.Options.CRURL != to.Options.CRURL {
		h.optsChan <- to.Options
	}
	return true
}

func (h *failureHandler) String() string {
	return "FailureHandler"
}

func sendFailureReports(ctx context.Context, reports []FailureReport, url string) {
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(reports); err != nil {
		panic(err)
	}

	client := &http.Client{
		Transport: &http.Transport{
			DialContext: dialer.DialContext,
			Proxy:       http.ProxyFromEnvironment,
		},
	}

	reqCtx, reqCancel := context.WithTimeout(ctx, sendTimeout)
	defer reqCancel()
	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, url, &b)
	if err != nil {
		l.Infoln("Failed to send failure report:", err)
		return
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		l.Infoln("Failed to send failure report:", err)
		return
	}
	resp.Body.Close()
	return
}

func newFailureReport(stat *failureStat) FailureReport {
	return FailureReport{
		FailureData: stat.data,
		Count:       stat.count,
		Version:     build.LongVersion,
	}
}
<commit_msg>lib/ur: Fix panic getting goroutines for failures (ref #7785) (#7890)<commit_after>// Copyright (C) 2020 The Syncthing Authors.
//
// 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 https://mozilla.org/MPL/2.0/.

package ur

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/syncthing/syncthing/lib/build"
	"github.com/syncthing/syncthing/lib/config"
	"github.com/syncthing/syncthing/lib/dialer"
	"github.com/syncthing/syncthing/lib/events"
	"github.com/syncthing/syncthing/lib/svcutil"

	"github.com/thejerf/suture/v4"
)

var (
	// When a specific failure first occurs, it is delayed by minDelay. If
	// more of the same failures occurs those are further delayed and
	// aggregated for maxDelay.
	minDelay             = 10 * time.Second
	maxDelay             = time.Minute
	sendTimeout          = time.Minute
	finalSendTimeout     = svcutil.ServiceTimeout / 2
	evChanClosed         = "failure event channel closed"
	invalidEventDataType = "failure event data is not a string"
)

type FailureReport struct {
	FailureData
	Count   int
	Version string
}

type FailureData struct {
	Description string
	Goroutines  string
	Extra       map[string]string
}

func FailureDataWithGoroutines(description string) FailureData {
	var buf *strings.Builder
	pprof.Lookup("goroutine").WriteTo(buf, 1)
	return FailureData{
		Description: description,
		Goroutines:  buf.String(),
	}
}

type FailureHandler interface {
	suture.Service
	config.Committer
}

func NewFailureHandler(cfg config.Wrapper, evLogger events.Logger) FailureHandler {
	return &failureHandler{
		cfg:      cfg,
		evLogger: evLogger,
		optsChan: make(chan config.OptionsConfiguration),
		buf:      make(map[string]*failureStat),
	}
}

type failureHandler struct {
	cfg      config.Wrapper
	evLogger events.Logger
	optsChan chan config.OptionsConfiguration
	buf      map[string]*failureStat
}

type failureStat struct {
	first, last time.Time
	count       int
	data        FailureData
}

func (h *failureHandler) Serve(ctx context.Context) error {
	cfg := h.cfg.Subscribe(h)
	defer h.cfg.Unsubscribe(h)
	url, sub, evChan := h.applyOpts(cfg.Options, nil)

	var err error
	timer := time.NewTimer(minDelay)
	resetTimer := make(chan struct{})
	for err == nil {
		select {
		case opts := <-h.optsChan:
			url, sub, evChan = h.applyOpts(opts, sub)
		case e, ok := <-evChan:
			if !ok {
				// Just to be safe - shouldn't ever happen, as
				// evChan is set to nil when unsubscribing.
				h.addReport(FailureData{Description: evChanClosed}, time.Now())
				evChan = nil
				continue
			}
			var data FailureData
			switch d := e.Data.(type) {
			case string:
				data.Description = d
			case FailureData:
				data = d
			default:
				// Same here, shouldn't ever happen.
				h.addReport(FailureData{Description: invalidEventDataType}, time.Now())
				continue
			}
			h.addReport(data, e.Time)
		case <-timer.C:
			reports := make([]FailureReport, 0, len(h.buf))
			now := time.Now()
			for descr, stat := range h.buf {
				if now.Sub(stat.last) > minDelay || now.Sub(stat.first) > maxDelay {
					reports = append(reports, newFailureReport(stat))
					delete(h.buf, descr)
				}
			}
			if len(reports) > 0 {
				// Lets keep process events/configs while it might be timing out for a while
				go func() {
					sendFailureReports(ctx, reports, url)
					select {
					case resetTimer <- struct{}{}:
					case <-ctx.Done():
					}
				}()
			} else {
				timer.Reset(minDelay)
			}
		case <-resetTimer:
			timer.Reset(minDelay)
		case <-ctx.Done():
			err = ctx.Err()
		}
	}

	if sub != nil {
		sub.Unsubscribe()
		if len(h.buf) > 0 {
			reports := make([]FailureReport, 0, len(h.buf))
			for _, stat := range h.buf {
				reports = append(reports, newFailureReport(stat))
			}
			timeout, cancel := context.WithTimeout(context.Background(), finalSendTimeout)
			defer cancel()
			sendFailureReports(timeout, reports, url)
		}
	}
	return err
}

func (h *failureHandler) applyOpts(opts config.OptionsConfiguration, sub events.Subscription) (string, events.Subscription, <-chan events.Event) {
	// Sub nil checks just for safety - config updates can be racy.
	url := opts.CRURL + "/failure"
	if opts.URAccepted > 0 {
		if sub == nil {
			sub = h.evLogger.Subscribe(events.Failure)
		}
		return url, sub, sub.C()
	}
	if sub != nil {
		sub.Unsubscribe()
	}
	return url, nil, nil
}

func (h *failureHandler) addReport(data FailureData, evTime time.Time) {
	if stat, ok := h.buf[data.Description]; ok {
		stat.last = evTime
		stat.count++
		return
	}
	h.buf[data.Description] = &failureStat{
		first: evTime,
		last:  evTime,
		count: 1,
		data:  data,
	}
}

func (h *failureHandler) VerifyConfiguration(_, _ config.Configuration) error {
	return nil
}

func (h *failureHandler) CommitConfiguration(from, to config.Configuration) bool {
	if from.Options.CREnabled != to.Options.CREnabled || from.Options.CRURL != to.Options.CRURL {
		h.optsChan <- to.Options
	}
	return true
}

func (h *failureHandler) String() string {
	return "FailureHandler"
}

func sendFailureReports(ctx context.Context, reports []FailureReport, url string) {
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(reports); err != nil {
		panic(err)
	}

	client := &http.Client{
		Transport: &http.Transport{
			DialContext: dialer.DialContext,
			Proxy:       http.ProxyFromEnvironment,
		},
	}

	reqCtx, reqCancel := context.WithTimeout(ctx, sendTimeout)
	defer reqCancel()
	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, url, &b)
	if err != nil {
		l.Infoln("Failed to send failure report:", err)
		return
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		l.Infoln("Failed to send failure report:", err)
		return
	}
	resp.Body.Close()
	return
}

func newFailureReport(stat *failureStat) FailureReport {
	return FailureReport{
		FailureData: stat.data,
		Count:       stat.count,
		Version:     build.LongVersion,
	}
}
<|endoftext|>


<commit_before>package generator

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"text/template"

	yaml "gopkg.in/yaml.v2"

	"github.com/granateio/granate/generator/utils"
	"github.com/graphql-go/graphql/language/ast"
	"github.com/graphql-go/graphql/language/parser"
	"github.com/graphql-go/graphql/language/source"
)

// Generator represents the code generator main object
type Generator struct {
	Code     string
	Schema   string
	Template *template.Template
	Ast      *ast.Document
	Config   ProjectConfig
	LangConf LanguageConfig
	Nodes    astNodes

	TmplConf map[string]string
}

// ProjectConfig contains the granate.yaml information
type ProjectConfig struct {
	// TODO Support a globbing system
	Schemas  []string
	Language string
	Output   map[string]string
}

// LanguageConfig defines the language specific
// implementation information
type LanguageConfig struct {

	// Language specific syntax config
	Language struct {
		Scalars map[string]string
		Root    []string
	}

	// This is passed to the generators Cfg variable
	Config map[string]string

	// Main templates, each template in this list is executed in it's own
	// go routine
	Templates []string

	// Program/command used for formatting the output code
	Formatter struct {
		CMD  string
		Args []string
	}
}

type OutputFileBuffer struct {
	Path   string
	Buffer *bytes.Buffer
}

func (out *OutputFileBuffer) GetBuffer() *bytes.Buffer {
	return out.Buffer
}

type TemplateFileFuncs struct {
	BufferStack   *utils.Lifo
	SwapBuffer    *utils.SwapBuffer
	LocalTemplate *template.Template
	linenumber    int
}

func (tmpl *TemplateFileFuncs) LineNumbers() int {
	return tmpl.linenumber
}

func (tmpl *TemplateFileFuncs) Start(path string) string {
	// Push current buffer on the stack
	tmpl.BufferStack.Push(tmpl.SwapBuffer.GetBuffer())

	// Create a new OpaqueBuffer
	output := &OutputFileBuffer{
		Path:   path,
		Buffer: &bytes.Buffer{},
	}

	tmpl.SwapBuffer.SetBuffer(output)

	return ""
}

func (tmpl *TemplateFileFuncs) End() string {

	output, ok := tmpl.SwapBuffer.GetBuffer().(*OutputFileBuffer)

	if ok == false {
		panic("GetBuffer() does not return a pointer to OutputFileBuffer")
	}

	if output.Path == "" {
		return ""
	}

	// fmt.Println(output.GetBuffer().String())
	// Unnecessary:
	// tmpl.FileBuffers = append(tmpl.FileBuffers, output)

	dir := path.Join(".", path.Dir(output.Path))
	err := os.MkdirAll(dir, os.ModePerm)

	if err != nil {
		fmt.Println(err)
	}

	// TODO: Read the fmt command from config
	// cmd := exec.Command("gofmt")
	cmd := exec.Command("goimports")
	stdin, err := cmd.StdinPipe()
	check(err)

	// fmt.Println(output.GetBuffer().String())

	go func() {
		defer stdin.Close()
		io.WriteString(stdin, output.GetBuffer().String())
	}()

	out, err := cmd.CombinedOutput()

	ln, _ := utils.LineCounter(bytes.NewReader(out))
	tmpl.linenumber += ln

	err = ioutil.WriteFile(output.Path, out, 0644)
	// err = ioutil.WriteFile(output.Path, output.GetBuffer().Bytes(), 0644)
	check(err)

	prevBuffer, ok := tmpl.BufferStack.Pop().(utils.OpaqueBytesBuffer)
	if ok == false {
		panic("Found wrong type in BufferStack")
	}

	tmpl.SwapBuffer.SetBuffer(prevBuffer)

	return ""
}

func (lang LanguageConfig) IsRoot(val string) bool {
	for _, root := range lang.Language.Root {
		if root == val {
			return true
		}
	}

	return false
}

// New creates a new Generator instance
func New(config string) (*Generator, error) {

	confFile, err := ioutil.ReadFile(config)
	check(err)

	genCfg := ProjectConfig{}
	err = yaml.Unmarshal(confFile, &genCfg)
	check(err)

	// Combine all .graphql files into one schema
	var schema bytes.Buffer
	for _, scm := range genCfg.Schemas {
		file, err := ioutil.ReadFile(scm)
		check(err)
		schema.Write(file)
	}

	// Create the generated package directory
	// Ignore error for now
	// err = os.Mkdir(genCfg.Package, 0766)

	src := source.NewSource(&source.Source{
		Body: schema.Bytes(),
		Name: "Schema",
	})

	AST, err := parser.Parse(parser.ParseParams{
		Source: src,
	})

	check(err)

	gopath := os.Getenv("GOPATH")
	projectpath := gopath + "/src/github.com/granateio/granate/"
	langpath := projectpath + "language/" + genCfg.Language + "/"

	langConfigFile, err := ioutil.ReadFile(langpath + "config.yaml")
	check(err)

	langConfig := LanguageConfig{}
	err = yaml.Unmarshal(langConfigFile, &langConfig)
	check(err)

	gen := &Generator{
		Schema:   schema.String(),
		Ast:      AST,
		TmplConf: langConfig.Config,
		Config:   genCfg,
		LangConf: langConfig,
	}

	// gen.Nodes.Connection = make(map[string]ast.Node)

	gen.Template, err = template.New("main").
		Funcs(gen.funcMap()).
		ParseGlob(langpath + "*.tmpl")

	check(err)

	return gen, nil
}

type namedDefinition interface {
	GetName() *ast.Name
	GetKind() string
}

// TODO: Find a better name for the NamedLookup function
func (gen *Generator) NamedLookup(name string) ast.Node {
	return NodeByName(gen.Nodes.Definition, name)
}

// TODO: Much the same as the NamedLookup function
func NodeByName(nodes []ast.Node, name string) ast.Node {
	for _, node := range nodes {
		named, ok := node.(namedDefinition)
		if ok == false {
			continue
		}
		if named.GetName().Value == name {
			return node
		}
	}

	log.Fatalf("Type with name '%s' is not defined", name)
	return nil
}

type generatorPass struct {
	Name string
	File string
}

func (gen generatorPass) template(name string) string {
	return gen.Name + "_" + name
}

type astNodes struct {
	Root       []ast.Node
	Definition []ast.Node
	Object     []ast.Node
	Relay      []ast.Node
}

type ConnectionDefinition struct {
	Name     *ast.Name
	Loc      *ast.Location
	NodeType ast.Node
}

func (con ConnectionDefinition) GetKind() string {
	return "ConnectionDefinition"
}

func (con ConnectionDefinition) GetLoc() *ast.Location {
	return con.Loc
}

func (con ConnectionDefinition) GetName() *ast.Name {
	return con.Name
}

// Generate starts the code generation process
func (gen *Generator) Generate() {
	definitions := gen.Ast.Definitions

	tmpl := gen.Template
	mainTemplates := gen.LangConf.Templates

	var wait sync.WaitGroup
	var nodes astNodes
	connections := make(map[string]bool)

	// Gather usefull definitions
	for _, def := range definitions {
		namedef, ok := def.(namedDefinition)

		if ok == false {
			continue
		}

		nodes.Definition = append(nodes.Definition, def)

		if gen.LangConf.IsRoot(namedef.GetName().Value) {
			nodes.Root = append(nodes.Root, def)
		}

		objectDef, ok := def.(*ast.ObjectDefinition)
		if ok == false {
			continue
		}

		nodes.Object = append(nodes.Object, def)

		// Find and add relay connections
		for _, connection := range objectDef.Fields {
			conloc := connection.Type.GetLoc()
			contype := string(conloc.Source.Body[conloc.Start:conloc.End])
			if strings.HasSuffix(contype, "Connection") {
				// if _, ok := nodes.Connection[contype]; ok == true {
				if _, ok := connections[contype]; ok == true {
					continue
				}
				con := ConnectionDefinition{
					Name: ast.NewName(&ast.Name{
						Value: contype,
						Loc:   conloc,
					}),
					Loc:      conloc,
					NodeType: NodeByName(gen.Ast.Definitions, strings.TrimSuffix(contype, "Connection")),
				}
				nodes.Definition = append(nodes.Definition, con)
				connections[contype] = true
			}
		}

		for _, iface := range objectDef.Interfaces {
			body := string(iface.Loc.Source.Body)
			name := body[iface.Loc.Start:iface.Loc.End]
			if name == "Node" {
				nodes.Relay = append(nodes.Relay, def)
			}
		}
	}

	gen.Nodes = nodes

	linecounter := make(chan int)
	quit := make(chan bool)

	go func(quit chan bool, counter chan int) {
		sum := 0
		for {
			select {
			case number := <-counter:
				sum += number
			case <-quit:
				fmt.Println("Generated", sum, "lines of code")
				return
			}
		}
	}(quit, linecounter)

	for _, mainTmpl := range mainTemplates {
		wait.Add(1)

		go func(mainTmpl string, counter chan int) {
			defer wait.Done()

			localTemplate, err := tmpl.Clone()
			if err != nil {
				panic(err)
			}

			codebuffer := &utils.SwapBuffer{}
			codebuffer.SetBuffer(&OutputFileBuffer{
				Buffer: &bytes.Buffer{},
			})

			localFileFuncs := TemplateFileFuncs{
				BufferStack:   &utils.Lifo{},
				SwapBuffer:    codebuffer,
				LocalTemplate: localTemplate,
			}

			partialfunc := func(name string, data interface{}) string {
				localbuffer := bytes.Buffer{}
				localTemplate.ExecuteTemplate(&localbuffer, name, data)
				return localbuffer.String()
			}

			fileFuncsMap := template.FuncMap{
				"startfile": localFileFuncs.Start,
				"endfile":   localFileFuncs.End,
				"partial":   partialfunc,
			}

			localTemplate = localTemplate.Funcs(fileFuncsMap)

			err = localTemplate.ExecuteTemplate(codebuffer, mainTmpl, nil)
			if err != nil {
				panic(err)
			}

			counter <- localFileFuncs.LineNumbers()

		}(mainTmpl, linecounter)
	}

	wait.Wait()

	quit <- true

	// fmt.Printf("Generated %d lines of code\n", lines)

}

func check(e error) {
	if e != nil {
		panic(e)
	}
}
<commit_msg>More appropriate error handling/messages<commit_after>package generator

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"text/template"

	yaml "gopkg.in/yaml.v2"

	"github.com/granateio/granate/generator/utils"
	"github.com/graphql-go/graphql/language/ast"
	"github.com/graphql-go/graphql/language/parser"
	"github.com/graphql-go/graphql/language/source"
)

// Generator represents the code generator main object
type Generator struct {
	Code     string
	Schema   string
	Template *template.Template
	Ast      *ast.Document
	Config   ProjectConfig
	LangConf LanguageConfig
	Nodes    astNodes

	TmplConf map[string]string
}

// ProjectConfig contains the granate.yaml information
type ProjectConfig struct {
	// TODO Support a globbing system
	Schemas  []string
	Language string
	Output   map[string]string
}

// LanguageConfig defines the language specific
// implementation information
type LanguageConfig struct {

	// Language specific syntax config
	Language struct {
		Scalars map[string]string
		Root    []string
	}

	// This is passed to the generators Cfg variable
	Config map[string]string

	// Main templates, each template in this list is executed in it's own
	// go routine
	Templates []string

	// Program/command used for formatting the output code
	Formatter struct {
		CMD  string
		Args []string
	}
}

type OutputFileBuffer struct {
	Path   string
	Buffer *bytes.Buffer
}

func (out *OutputFileBuffer) GetBuffer() *bytes.Buffer {
	return out.Buffer
}

type TemplateFileFuncs struct {
	BufferStack   *utils.Lifo
	SwapBuffer    *utils.SwapBuffer
	LocalTemplate *template.Template
	linenumber    int
}

func (tmpl *TemplateFileFuncs) LineNumbers() int {
	return tmpl.linenumber
}

func (tmpl *TemplateFileFuncs) Start(path string) string {
	// Push current buffer on the stack
	tmpl.BufferStack.Push(tmpl.SwapBuffer.GetBuffer())

	// Create a new OpaqueBuffer
	output := &OutputFileBuffer{
		Path:   path,
		Buffer: &bytes.Buffer{},
	}

	tmpl.SwapBuffer.SetBuffer(output)

	return ""
}

func (tmpl *TemplateFileFuncs) End() string {

	output, ok := tmpl.SwapBuffer.GetBuffer().(*OutputFileBuffer)

	if ok == false {
		panic("GetBuffer() does not return a pointer to OutputFileBuffer")
	}

	if output.Path == "" {
		return ""
	}

	// fmt.Println(output.GetBuffer().String())
	// Unnecessary:
	// tmpl.FileBuffers = append(tmpl.FileBuffers, output)

	dir := path.Join(".", path.Dir(output.Path))
	err := os.MkdirAll(dir, os.ModePerm)

	if err != nil {
		fmt.Println(err)
	}

	// TODO: Read the fmt command from config
	// cmd := exec.Command("gofmt")
	cmd := exec.Command("goimports")
	stdin, err := cmd.StdinPipe()
	check(err)

	go func() {
		defer stdin.Close()
		io.WriteString(stdin, output.GetBuffer().String())
	}()

	out, err := cmd.CombinedOutput()

	ln, _ := utils.LineCounter(bytes.NewReader(out))
	tmpl.linenumber += ln

	err = ioutil.WriteFile(output.Path, out, 0644)
	// err = ioutil.WriteFile(output.Path, output.GetBuffer().Bytes(), 0644)

	check(err)

	prevBuffer, ok := tmpl.BufferStack.Pop().(utils.OpaqueBytesBuffer)
	if ok == false {
		panic("Found wrong type in BufferStack")
	}

	tmpl.SwapBuffer.SetBuffer(prevBuffer)

	return ""
}

func (lang LanguageConfig) IsRoot(val string) bool {
	for _, root := range lang.Language.Root {
		if root == val {
			return true
		}
	}

	return false
}

// New creates a new Generator instance
func New(config string) (*Generator, error) {

	confFile, err := ioutil.ReadFile(config)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	genCfg := ProjectConfig{}
	err = yaml.Unmarshal(confFile, &genCfg)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	// Combine all .graphql files into one schema
	var schema bytes.Buffer
	for _, scm := range genCfg.Schemas {
		file, err := ioutil.ReadFile(scm)
		check(err)
		schema.Write(file)
	}

	// Create the generated package directory
	// Ignore error for now
	// err = os.Mkdir(genCfg.Package, 0766)

	src := source.NewSource(&source.Source{
		Body: schema.Bytes(),
		Name: "Schema",
	})

	AST, err := parser.Parse(parser.ParseParams{
		Source: src,
	})

	check(err)

	gopath := os.Getenv("GOPATH")
	projectpath := gopath + "/src/github.com/granateio/granate/"
	langpath := projectpath + "language/" + genCfg.Language + "/"

	langConfigFile, err := ioutil.ReadFile(langpath + "config.yaml")
	check(err)

	langConfig := LanguageConfig{}
	err = yaml.Unmarshal(langConfigFile, &langConfig)
	check(err)

	gen := &Generator{
		Schema:   schema.String(),
		Ast:      AST,
		TmplConf: langConfig.Config,
		Config:   genCfg,
		LangConf: langConfig,
	}

	// gen.Nodes.Connection = make(map[string]ast.Node)

	gen.Template, err = template.New("main").
		Funcs(gen.funcMap()).
		ParseGlob(langpath + "*.tmpl")

	check(err)

	return gen, nil
}

type namedDefinition interface {
	GetName() *ast.Name
	GetKind() string
}

// TODO: Find a better name for the NamedLookup function
func (gen *Generator) NamedLookup(name string) ast.Node {
	return NodeByName(gen.Nodes.Definition, name)
}

// TODO: Much the same as the NamedLookup function
func NodeByName(nodes []ast.Node, name string) ast.Node {
	for _, node := range nodes {
		named, ok := node.(namedDefinition)
		if ok == false {
			continue
		}
		if named.GetName().Value == name {
			return node
		}
	}

	log.Fatalf("Type with name '%s' is not defined", name)
	return nil
}

type generatorPass struct {
	Name string
	File string
}

func (gen generatorPass) template(name string) string {
	return gen.Name + "_" + name
}

type astNodes struct {
	Root       []ast.Node
	Definition []ast.Node
	Object     []ast.Node
	Relay      []ast.Node
}

type ConnectionDefinition struct {
	Name     *ast.Name
	Loc      *ast.Location
	NodeType ast.Node
}

func (con ConnectionDefinition) GetKind() string {
	return "ConnectionDefinition"
}

func (con ConnectionDefinition) GetLoc() *ast.Location {
	return con.Loc
}

func (con ConnectionDefinition) GetName() *ast.Name {
	return con.Name
}

// Generate starts the code generation process
func (gen *Generator) Generate() {
	definitions := gen.Ast.Definitions

	tmpl := gen.Template
	mainTemplates := gen.LangConf.Templates

	var wait sync.WaitGroup
	var nodes astNodes
	connections := make(map[string]bool)

	// Gather usefull definitions
	for _, def := range definitions {
		namedef, ok := def.(namedDefinition)

		if ok == false {
			continue
		}

		nodes.Definition = append(nodes.Definition, def)

		if gen.LangConf.IsRoot(namedef.GetName().Value) {
			nodes.Root = append(nodes.Root, def)
		}

		objectDef, ok := def.(*ast.ObjectDefinition)
		if ok == false {
			continue
		}

		nodes.Object = append(nodes.Object, def)

		// Find and add relay connections
		for _, connection := range objectDef.Fields {
			conloc := connection.Type.GetLoc()
			contype := string(conloc.Source.Body[conloc.Start:conloc.End])
			if strings.HasSuffix(contype, "Connection") {
				// if _, ok := nodes.Connection[contype]; ok == true {
				if _, ok := connections[contype]; ok == true {
					continue
				}
				con := ConnectionDefinition{
					Name: ast.NewName(&ast.Name{
						Value: contype,
						Loc:   conloc,
					}),
					Loc:      conloc,
					NodeType: NodeByName(gen.Ast.Definitions, strings.TrimSuffix(contype, "Connection")),
				}
				nodes.Definition = append(nodes.Definition, con)
				connections[contype] = true
			}
		}

		for _, iface := range objectDef.Interfaces {
			body := string(iface.Loc.Source.Body)
			name := body[iface.Loc.Start:iface.Loc.End]
			if name == "Node" {
				nodes.Relay = append(nodes.Relay, def)
			}
		}
	}

	gen.Nodes = nodes

	linecounter := make(chan int)
	quit := make(chan bool)

	go func(quit chan bool, counter chan int) {
		sum := 0
		for {
			select {
			case number := <-counter:
				sum += number
			case <-quit:
				fmt.Println("Generated", sum, "lines of code")
				return
			}
		}
	}(quit, linecounter)

	for _, mainTmpl := range mainTemplates {
		wait.Add(1)

		go func(mainTmpl string, counter chan int) {
			defer wait.Done()

			localTemplate, err := tmpl.Clone()
			if err != nil {
				panic(err)
			}

			codebuffer := &utils.SwapBuffer{}
			codebuffer.SetBuffer(&OutputFileBuffer{
				Buffer: &bytes.Buffer{},
			})

			localFileFuncs := TemplateFileFuncs{
				BufferStack:   &utils.Lifo{},
				SwapBuffer:    codebuffer,
				LocalTemplate: localTemplate,
			}

			partialfunc := func(name string, data interface{}) string {
				localbuffer := bytes.Buffer{}
				localTemplate.ExecuteTemplate(&localbuffer, name, data)
				return localbuffer.String()
			}

			fileFuncsMap := template.FuncMap{
				"startfile": localFileFuncs.Start,
				"endfile":   localFileFuncs.End,
				"partial":   partialfunc,
			}

			localTemplate = localTemplate.Funcs(fileFuncsMap)

			err = localTemplate.ExecuteTemplate(codebuffer, mainTmpl, nil)
			if err != nil {
				panic(err)
			}

			counter <- localFileFuncs.LineNumbers()

		}(mainTmpl, linecounter)
	}

	wait.Wait()

	quit <- true

	// fmt.Printf("Generated %d lines of code\n", lines)

}

func check(e error) {
	if e != nil {
		panic(e)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"time"

	"github.com/sam701/awstools/config"
	"github.com/sam701/awstools/cred"
	"github.com/sam701/awstools/sess"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/urfave/cli"
)

var scriptOutput io.Writer = os.Stdout

func actionAssumeRole(c *cli.Context) error {
	if len(c.Args()) == 2 {
		var account, role string
		account = c.Args().Get(0)
		role = c.Args().Get(1)
		exportScriptPath := c.String("export")
		if exportScriptPath != "" {
			f, err := os.OpenFile(exportScriptPath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0600)
			if err != nil {
				log.Fatalln("ERROR:", err)
			}
			defer f.Close()
			scriptOutput = f
		}
		assumeRole(account, role)
	} else {
		cli.ShowCommandHelp(c, "assume")
	}
	return nil
}

func assumeRole(account, role string) {
	account = adjustAccountName(account)
	role = adjustRoleName(role)

	err := tryToAssumeRole(account, role)
	if err != nil {

		getMainAccountMfaSessionToken()
		err = tryToAssumeRole(account, role)
		if err != nil {
			log.Fatalln(err)
		}

		if needRotateKey() {
			rotateMainAccountKey(iam.New(sess.New(config.Current.Profiles.MainAccountMfaSession)))
		}
	}
}

func needRotateKey() bool {
	session := sess.New(config.Current.Profiles.MainAccountMfaSession)
	cl := iam.New(session)

	keyId := cred.GetMainAccountKeyId(config.Current.Profiles.MainAccount)
	out, err := cl.ListAccessKeys(&iam.ListAccessKeysInput{})
	if err != nil {
		log.Fatalln("ERROR", err)
	}

	var creationTime time.Time
	for _, md := range out.AccessKeyMetadata {
		if keyId == *md.AccessKeyId {
			creationTime = *md.CreateDate
			break
		}
	}
	if creationTime.IsZero() {
		log.Fatalln("Cannot get creation time for key", keyId)
	}

	log.Println("Interval:", config.Current.KeyRotationIntervalMinutes)
	log.Println("Lifetime:", int(time.Now().Sub(creationTime).Minutes()))

	return int(time.Now().Sub(creationTime).Minutes()) >= config.Current.KeyRotationIntervalMinutes
}

func adjustAccountName(account string) string {
	var candidate string
	for k, _ := range config.Current.Accounts {
		if strings.Contains(k, account) {
			if candidate == "" {
				candidate = k
			} else {
				log.Fatalf("Ambiguous account name. Possible matches: %s, %s\n", candidate, k)
			}
		}
	}
	if candidate == "" {
		log.Fatalln("No such account:", account)
	}
	return candidate
}

func adjustRoleName(role string) string {
	if role[0] == 'r' {
		return "ReadOnlyAccess"
	} else if role[0] == 'w' {
		return "PowerUserAccess"
	} else {
		return role
	}
}

var userName string

func getUserName() string {
	if userName == "" {
		client := iam.New(sess.New(config.Current.Profiles.MainAccount))
		data, err := client.GetUser(&iam.GetUserInput{})
		if err != nil {
			log.Fatalln("ERROR:", err)
		}
		userName = *data.User.UserName
	}
	return userName
}

func readMfaToken() string {
	scanner := bufio.NewScanner(os.Stdin)
	fmt.Print("MFA token: ")
	scanner.Scan()
	return scanner.Text()
}

func getMainAccountMfaSessionToken() {
	token := readMfaToken()
	session := sess.New(config.Current.Profiles.MainAccount)
	stsClient := sts.New(session)
	data, err := stsClient.GetSessionToken(&sts.GetSessionTokenInput{
		SerialNumber: aws.String(fmt.Sprintf("arn:aws:iam::%s:mfa/%s",
			accountId(config.Current.Profiles.MainAccount),
			getUserName())),
		TokenCode: aws.String(token),
	})
	if err != nil {
		log.Fatalln("ERROR:", err)
	}

	persistSharedCredentials(data.Credentials, config.Current.Profiles.MainAccountMfaSession)
}

func accountId(accountName string) string {
	accountId := config.Current.Accounts[accountName]
	if accountId == "" {
		log.Fatalln("Unknown account name:", accountName)
	}
	return accountId
}

func tryToAssumeRole(accountName, role string) error {
	session := sess.New(config.Current.Profiles.MainAccountMfaSession)
	accountId := accountId(accountName)

	stsClient := sts.New(session)
	assumeData, err := stsClient.AssumeRole(&sts.AssumeRoleInput{
		RoleArn:         aws.String(fmt.Sprintf("arn:aws:iam::%s:role/%s", accountId, role)),
		RoleSessionName: aws.String(getUserName()),
	})
	if err != nil {
		return err
	}

	profile := fmt.Sprintf("%s %s", accountName, role)
	persistSharedCredentials(assumeData.Credentials, profile)
	printExport(assumeData.Credentials)

	cred.SetProfileRegion(profile, config.Current.DefaultRegion)
	return nil
}

func printExport(cred *sts.Credentials) {
	shell := path.Base(os.Getenv("SHELL"))
	var pattern string
	if shell == "fish" {
		pattern = "set -xg %s \"%s\"\n"
	} else {
		pattern = "export %s=\"%s\"\n"
	}

	exp := func(key, value string) {
		fmt.Fprintf(scriptOutput, pattern, key, value)
	}

	exp("AWS_ACCESS_KEY_ID", *cred.AccessKeyId)
	exp("AWS_SECRET_ACCESS_KEY", *cred.SecretAccessKey)
	exp("AWS_SESSION_TOKEN", *cred.SessionToken)
}

func persistSharedCredentials(c *sts.Credentials, profile string) {
	cred.SaveCredentials(profile, *c.AccessKeyId, *c.SecretAccessKey, *c.SessionToken)
}
<commit_msg>Remove debugging output<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"time"

	"github.com/sam701/awstools/config"
	"github.com/sam701/awstools/cred"
	"github.com/sam701/awstools/sess"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/urfave/cli"
)

var scriptOutput io.Writer = os.Stdout

func actionAssumeRole(c *cli.Context) error {
	if len(c.Args()) == 2 {
		var account, role string
		account = c.Args().Get(0)
		role = c.Args().Get(1)
		exportScriptPath := c.String("export")
		if exportScriptPath != "" {
			f, err := os.OpenFile(exportScriptPath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0600)
			if err != nil {
				log.Fatalln("ERROR:", err)
			}
			defer f.Close()
			scriptOutput = f
		}
		assumeRole(account, role)
	} else {
		cli.ShowCommandHelp(c, "assume")
	}
	return nil
}

func assumeRole(account, role string) {
	account = adjustAccountName(account)
	role = adjustRoleName(role)

	err := tryToAssumeRole(account, role)
	if err != nil {

		getMainAccountMfaSessionToken()
		err = tryToAssumeRole(account, role)
		if err != nil {
			log.Fatalln(err)
		}

		if needRotateKey() {
			rotateMainAccountKey(iam.New(sess.New(config.Current.Profiles.MainAccountMfaSession)))
		}
	}
}

func needRotateKey() bool {
	session := sess.New(config.Current.Profiles.MainAccountMfaSession)
	cl := iam.New(session)

	keyId := cred.GetMainAccountKeyId(config.Current.Profiles.MainAccount)
	out, err := cl.ListAccessKeys(&iam.ListAccessKeysInput{})
	if err != nil {
		log.Fatalln("ERROR", err)
	}

	var creationTime time.Time
	for _, md := range out.AccessKeyMetadata {
		if keyId == *md.AccessKeyId {
			creationTime = *md.CreateDate
			break
		}
	}
	if creationTime.IsZero() {
		log.Fatalln("Cannot get creation time for key", keyId)
	}

	return int(time.Now().Sub(creationTime).Minutes()) >= config.Current.KeyRotationIntervalMinutes
}

func adjustAccountName(account string) string {
	var candidate string
	for k, _ := range config.Current.Accounts {
		if strings.Contains(k, account) {
			if candidate == "" {
				candidate = k
			} else {
				log.Fatalf("Ambiguous account name. Possible matches: %s, %s\n", candidate, k)
			}
		}
	}
	if candidate == "" {
		log.Fatalln("No such account:", account)
	}
	return candidate
}

func adjustRoleName(role string) string {
	if role[0] == 'r' {
		return "ReadOnlyAccess"
	} else if role[0] == 'w' {
		return "PowerUserAccess"
	} else {
		return role
	}
}

var userName string

func getUserName() string {
	if userName == "" {
		client := iam.New(sess.New(config.Current.Profiles.MainAccount))
		data, err := client.GetUser(&iam.GetUserInput{})
		if err != nil {
			log.Fatalln("ERROR:", err)
		}
		userName = *data.User.UserName
	}
	return userName
}

func readMfaToken() string {
	scanner := bufio.NewScanner(os.Stdin)
	fmt.Print("MFA token: ")
	scanner.Scan()
	return scanner.Text()
}

func getMainAccountMfaSessionToken() {
	token := readMfaToken()
	session := sess.New(config.Current.Profiles.MainAccount)
	stsClient := sts.New(session)
	data, err := stsClient.GetSessionToken(&sts.GetSessionTokenInput{
		SerialNumber: aws.String(fmt.Sprintf("arn:aws:iam::%s:mfa/%s",
			accountId(config.Current.Profiles.MainAccount),
			getUserName())),
		TokenCode: aws.String(token),
	})
	if err != nil {
		log.Fatalln("ERROR:", err)
	}

	persistSharedCredentials(data.Credentials, config.Current.Profiles.MainAccountMfaSession)
}

func accountId(accountName string) string {
	accountId := config.Current.Accounts[accountName]
	if accountId == "" {
		log.Fatalln("Unknown account name:", accountName)
	}
	return accountId
}

func tryToAssumeRole(accountName, role string) error {
	session := sess.New(config.Current.Profiles.MainAccountMfaSession)
	accountId := accountId(accountName)

	stsClient := sts.New(session)
	assumeData, err := stsClient.AssumeRole(&sts.AssumeRoleInput{
		RoleArn:         aws.String(fmt.Sprintf("arn:aws:iam::%s:role/%s", accountId, role)),
		RoleSessionName: aws.String(getUserName()),
	})
	if err != nil {
		return err
	}

	profile := fmt.Sprintf("%s %s", accountName, role)
	persistSharedCredentials(assumeData.Credentials, profile)
	printExport(assumeData.Credentials)

	cred.SetProfileRegion(profile, config.Current.DefaultRegion)
	return nil
}

func printExport(cred *sts.Credentials) {
	shell := path.Base(os.Getenv("SHELL"))
	var pattern string
	if shell == "fish" {
		pattern = "set -xg %s \"%s\"\n"
	} else {
		pattern = "export %s=\"%s\"\n"
	}

	exp := func(key, value string) {
		fmt.Fprintf(scriptOutput, pattern, key, value)
	}

	exp("AWS_ACCESS_KEY_ID", *cred.AccessKeyId)
	exp("AWS_SECRET_ACCESS_KEY", *cred.SecretAccessKey)
	exp("AWS_SESSION_TOKEN", *cred.SessionToken)
}

func persistSharedCredentials(c *sts.Credentials, profile string) {
	cred.SaveCredentials(profile, *c.AccessKeyId, *c.SecretAccessKey, *c.SessionToken)
}
<|endoftext|>


<commit_before>// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package hack

import (
	"reflect"
	"unsafe"
)

// MutableString can be used as string via string(MutableString) without performance loss.
type MutableString string

// String converts slice to MutableString without copy.
// The MutableString can be converts to string without copy.
// Use it at your own risk.
func String(b []byte) (s MutableString) {
	if len(b) == 0 {
		return ""
	}
	pbytes := (*reflect.SliceHeader)(unsafe.Pointer(&b))
	pstring := (*reflect.StringHeader)(unsafe.Pointer(&s))
	pstring.Data = pbytes.Data
	pstring.Len = pbytes.Len
	return
}

// Slice converts string to slice without copy.
// Use at your own risk.
func Slice(s string) (b []byte) {
	pbytes := (*reflect.SliceHeader)(unsafe.Pointer(&b))
	pstring := *(*reflect.StringHeader)(unsafe.Pointer(&s))
	pbytes.Data = pstring.Data
	pbytes.Len = pstring.Len
	pbytes.Cap = pstring.Len
	return
}
<commit_msg>util: fix vet (#23042)<commit_after>// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package hack

import (
	"reflect"
	"unsafe"
)

// MutableString can be used as string via string(MutableString) without performance loss.
type MutableString string

// String converts slice to MutableString without copy.
// The MutableString can be converts to string without copy.
// Use it at your own risk.
func String(b []byte) (s MutableString) {
	if len(b) == 0 {
		return ""
	}
	pbytes := (*reflect.SliceHeader)(unsafe.Pointer(&b))
	pstring := (*reflect.StringHeader)(unsafe.Pointer(&s))
	pstring.Data = pbytes.Data
	pstring.Len = pbytes.Len
	return
}

// Slice converts string to slice without copy.
// Use at your own risk.
func Slice(s string) (b []byte) {
	pbytes := (*reflect.SliceHeader)(unsafe.Pointer(&b))
	pstring := (*reflect.StringHeader)(unsafe.Pointer(&s))
	pbytes.Data = pstring.Data
	pbytes.Len = pstring.Len
	pbytes.Cap = pstring.Len
	return
}
<|endoftext|>