<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"log"
	"math/big"
	"strings"
	"sync"
	"time"

	"github.com/go-mangos/mangos"
	"github.com/go-mangos/mangos/protocol/bus"
	"github.com/go-mangos/mangos/transport/tlstcp"
	"github.com/nictuku/dht"
)

type BusMessage struct {
	From    string
	To      string
	Seq     int64
	Cmd     string
	Payload string
}

func NewBusMessage(data []byte, err error) (*BusMessage, error) {
	if err != nil {
		return nil, err
	}

	rv := &BusMessage{}
	if err = json.Unmarshal(data, rv); err != nil {
		return nil, err
	}
	return rv, nil
}

func (b *BusMessage) String() string {
	return fmt.Sprintf("[%s: %v (ID: %d)]", b.Cmd, b.Payload, b.Seq)
}

func (b *BusMessage) Bytes() []byte {
	m, err := json.Marshal(b)
	if err != nil {
		log.Fatalf(`can't marshall bus message to json: %s`, err)
	}
	return m
}

type Friend struct {
	name     string
	msg      *BusMessage
	lastSeen time.Time
	mtx      *sync.Mutex
	pm       *PeerManager
}

func NewFriend(name string, pm *PeerManager) *Friend {
	return &Friend{name: name, mtx: &sync.Mutex{}, lastSeen: time.Now(), pm: pm}
}

func (f *Friend) String() string {
	f.mtx.Lock()
	defer f.mtx.Unlock()

	if f.msg == nil {
		return fmt.Sprintf(`%s: no message, last seen: %s`, f.name, f.lastSeen)
	}

	return fmt.Sprintf("%s: last message: %s, last seen: %s ago", f.name, f.msg, time.Now().Sub(f.lastSeen))
}

func (f *Friend) HandleMessage(m *BusMessage) {
	f.mtx.Lock()
	defer f.mtx.Unlock()

	if f.msg != nil && f.msg.Seq >= m.Seq {
		/* We've already seen this message */
		return
	}

	f.msg = m
	f.lastSeen = time.Now()

	switch m.Cmd {
	case "event":
		ev := Event{}
		if err := json.Unmarshal([]byte(m.Payload), &ev); err != nil {
			log.Printf(`can't unmarshal event payload: %s`, err)
		}
		log.Printf(`got an event: %v`, ev)

		if err := f.pm.db.MergeEvent(ev); err != nil {
			log.Printf(`can't merge event: %s`, err)
		}
	case "hello":
		f.pm.mtx.Lock()
		defer f.pm.mtx.Unlock()

		f.pm.need_full_sync = true
	case "i'm alive":
		/* nothing */
	default:
		log.Printf(`unhandled message: tgt: %s, src: %s, cmd: %s, payload: %v`, m.To, m.From, m.Cmd, m.Payload)
	}
}

type PeerManager struct {
	d              *dht.DHT
	db             *DB
	bus            mangos.Socket
	nick           string
	friends        map[string]*Friend
	oldfriends     map[string]*Friend
	venue          string
	sequence       int64
	need_full_sync bool
	opts map[string]interface{}
	mtx            *sync.RWMutex
}

func NewPeerManager(db *DB) *PeerManager {
	conf := dht.NewConfig()
	conf.Port = 55000

	pm := PeerManager{db: db}

	pm.opts = map[string]interface{}{ mangos.OptionTLSConfig: pm.setupTLS() }

	var err error
	pm.bus, err = bus.NewSocket()
	if err != nil {
		log.Fatalf(`can't create BUS socket: %s`, err)
	}
	pm.bus.AddTransport(tlstcp.NewTransport())

	if pm.d, err = dht.New(conf); err != nil {
		log.Fatalf(`can't create DHT: %s`, err)
	}

	sum := sha1.Sum([]byte("LetsMeetHere"))
	pm.venue = hex.EncodeToString(sum[:])

	buf := make([]byte, 4)
	rand.Read(buf)
	pm.nick = hex.EncodeToString(buf)

	log.Printf(`My nickname is %s`, pm.nick)
	log.Printf(`I will meet my friends at %s`, pm.venue)

	pm.friends = make(map[string]*Friend)
	pm.oldfriends = make(map[string]*Friend)

	pm.mtx = &sync.RWMutex{}

	return &pm
}

func (pm *PeerManager) setupTLS() *tls.Config {
	cert := pm.setupCertificate()

	return &tls.Config{
		InsecureSkipVerify: true,
		Certificates:       []tls.Certificate{cert},
	}
}

func (pm *PeerManager) setupCertificate() tls.Certificate {
	template := x509.Certificate{
		SerialNumber: big.NewInt(4711),
		Subject: pkix.Name{
			Organization: []string{"unobtanium"},
			CommonName:   "*",
		},
		DNSNames:              []string{`*`},
		NotAfter:              time.Now().AddDate(0, 0, 10),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		SignatureAlgorithm:    x509.SHA512WithRSA,
		PublicKeyAlgorithm:    x509.ECDSA,
		SubjectKeyId:          []byte{1, 2, 3, 4, 5},
		BasicConstraintsValid: true,
		IsCA: true,
	}

	log.Printf(`generating 2048 bit RSA private key`)
	priv, _ := rsa.GenerateKey(rand.Reader, 2048)
	pub := &priv.PublicKey
	log.Printf(`building self-signed certificate`)
	cert_raw, err := x509.CreateCertificate(rand.Reader, &template, &template, pub, priv)
	if err != nil {
		log.Fatalf(`can't generate TLS certificate: %s`, err)
	}
	log.Printf(`created a cert of length %d bytes`, len(cert_raw))

	privbuf := &bytes.Buffer{}
	pem.Encode(privbuf, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	certbuf := &bytes.Buffer{}
	pem.Encode(certbuf, &pem.Block{Type: "CERTIFICATE", Bytes: cert_raw})

	cert, err := tls.X509KeyPair(certbuf.Bytes(), privbuf.Bytes())
	if err != nil {
		log.Fatalf(`can't build key pair: %s`, err)
	}

	return cert
}

func (pm *PeerManager) String() string {
	pm.mtx.RLock()
	defer pm.mtx.RUnlock()

	s := []string{
		fmt.Sprintf("DHT: Port: %d", pm.d.Port()),
		fmt.Sprintf("\r\nI am %s, my sequence ID is %d", pm.nick, pm.sequence),
		fmt.Sprintf("I have %d friend(s)", len(pm.friends)),
		fmt.Sprintf("We're meeting at '%s'", pm.venue),
		"\r\nThese are my friends:\r\n",
	}

	for _, f := range pm.friends {
		s = append(s, f.String())
	}

	s = append(s, "\r\nI haven't heard from these guys in a while:\r\n")
	for _, f := range pm.oldfriends {
		s = append(s, f.String())
	}

	return strings.Join(s, "\r\n")
}

func (pm *PeerManager) Publish(dst, cmd, payload string) error {
	pm.mtx.Lock()
	pm.sequence++
	m := &BusMessage{From: pm.nick, To: dst, Seq: pm.sequence, Cmd: cmd, Payload: payload}
	pm.mtx.Unlock()

	return pm.bus.Send(m.Bytes())
}

func (pm *PeerManager) Loop() {
	ih, err := dht.DecodeInfoHash(pm.venue)
	if err != nil {
		log.Printf(`can't decode infohash: %s`, err)
		return
	}

	if err := pm.d.Start(); err != nil {
		log.Printf(`can't start DHT: %s`, err)
		return
	}

	log.Printf(`DHT bound to port %d`, pm.d.Port())

	if err := pm.bus.ListenOptions(fmt.Sprintf("tls+tcp://*:%d", pm.d.Port()), pm.opts); err != nil {
		log.Fatalf(`can't listen on BUS socket: %s`, err)
	}

	go pm.drainPeers()

	go func() {
		for {
			time.Sleep(2 * time.Second)
			pm.mtx.Lock()
			for name, f := range pm.friends {
				f.mtx.Lock()
				if f.lastSeen.Add(10 * time.Second).Before(time.Now()) {
					log.Printf(`haven't heard from %s in a while`, name)
					pm.oldfriends[name] = f
					delete(pm.friends, name)
				}
				f.mtx.Unlock()
			}
			pm.mtx.Unlock()
		}
	}()

	go func() {
		for {
			/* Receive message */
			m, err := NewBusMessage(pm.bus.Recv())
			if err != nil {
				log.Printf(`can't receive message from bus: %s`, err)
				continue
			}

			/* Ignore messages from ourselves */
			if m.From == pm.nick {
				continue
			}

			/* Ignore messages not to 'everyone' or us */
			if m.To != pm.nick && m.To != "*" {
				continue
			}

			pm.mtx.Lock()
			f, ok := pm.friends[m.From]
			if !ok {
				f = NewFriend(m.From, pm)
				pm.friends[m.From] = f
				pm.need_full_sync = true
			}
			pm.mtx.Unlock()

			f.HandleMessage(m)

			if !ok {
				log.Printf(`%s is a new friend!`, m.From)
				pm.Publish(m.From, "hello", "")
			}
		}
	}()

	go func() {
		i := 0
		for {
			pm.Publish("*", "i'm alive", fmt.Sprintf("%d", i))
			i++
			time.Sleep(5 * time.Second)
		}
	}()

	go func() {
		for {
			buf, err := json.Marshal(<-pm.db.Events)
			if err != nil {
				log.Fatalf(`can't marshal event: %s`, err)
			}
			pm.Publish("*", "event", string(buf))
		}
	}()

	go func() {
		for {
			time.Sleep(5 * time.Second)
			pm.mtx.Lock()
			if !pm.need_full_sync {
				pm.mtx.Unlock()
				continue
			}
			pm.need_full_sync = false
			pm.mtx.Unlock()

			log.Println(`doing a full sync`)

			for _, env := range pm.db.AllEnvelopes() {
				_, events, err := pm.db.EnvelopeWithHistory(env.Id)
				if err != nil {
					log.Fatalf(`can't get events for envelope %s: %s`, env.Id, err)
				}
				for _, e := range events {
					pm.db.Events <- e
				}
			}
		}
	}()

	for {
		pm.d.PeersRequest(string(ih), true)
		time.Sleep(5 * time.Second)
	}
}

func (pm *PeerManager) drainPeers() {
	log.Printf(`draining DHT`)
	seen := make(map[string]struct{})

	for r := range pm.d.PeersRequestResults {
		for _, peers := range r {
			for _, x := range peers {
				addr := dht.DecodePeerAddress(x)
				if _, ok := seen[addr]; !ok {
					pm.connectToPeer(addr)
				}
				seen[addr] = struct{}{}
			}
		}
	}
}

func (pm *PeerManager) connectToPeer(addr string) {
	if err := pm.bus.DialOptions(fmt.Sprintf("tls+tcp://%s", addr), pm.opts); err != nil {
		log.Printf(`can't connect SUB to %s: %s`, addr, err)
	}
}
<commit_msg>Gofmt<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"log"
	"math/big"
	"strings"
	"sync"
	"time"

	"github.com/go-mangos/mangos"
	"github.com/go-mangos/mangos/protocol/bus"
	"github.com/go-mangos/mangos/transport/tlstcp"
	"github.com/nictuku/dht"
)

type BusMessage struct {
	From    string
	To      string
	Seq     int64
	Cmd     string
	Payload string
}

func NewBusMessage(data []byte, err error) (*BusMessage, error) {
	if err != nil {
		return nil, err
	}

	rv := &BusMessage{}
	if err = json.Unmarshal(data, rv); err != nil {
		return nil, err
	}
	return rv, nil
}

func (b *BusMessage) String() string {
	return fmt.Sprintf("[%s: %v (ID: %d)]", b.Cmd, b.Payload, b.Seq)
}

func (b *BusMessage) Bytes() []byte {
	m, err := json.Marshal(b)
	if err != nil {
		log.Fatalf(`can't marshall bus message to json: %s`, err)
	}
	return m
}

type Friend struct {
	name     string
	msg      *BusMessage
	lastSeen time.Time
	mtx      *sync.Mutex
	pm       *PeerManager
}

func NewFriend(name string, pm *PeerManager) *Friend {
	return &Friend{name: name, mtx: &sync.Mutex{}, lastSeen: time.Now(), pm: pm}
}

func (f *Friend) String() string {
	f.mtx.Lock()
	defer f.mtx.Unlock()

	if f.msg == nil {
		return fmt.Sprintf(`%s: no message, last seen: %s`, f.name, f.lastSeen)
	}

	return fmt.Sprintf("%s: last message: %s, last seen: %s ago", f.name, f.msg, time.Now().Sub(f.lastSeen))
}

func (f *Friend) HandleMessage(m *BusMessage) {
	f.mtx.Lock()
	defer f.mtx.Unlock()

	if f.msg != nil && f.msg.Seq >= m.Seq {
		/* We've already seen this message */
		return
	}

	f.msg = m
	f.lastSeen = time.Now()

	switch m.Cmd {
	case "event":
		ev := Event{}
		if err := json.Unmarshal([]byte(m.Payload), &ev); err != nil {
			log.Printf(`can't unmarshal event payload: %s`, err)
		}
		log.Printf(`got an event: %v`, ev)

		if err := f.pm.db.MergeEvent(ev); err != nil {
			log.Printf(`can't merge event: %s`, err)
		}
	case "hello":
		f.pm.mtx.Lock()
		defer f.pm.mtx.Unlock()

		f.pm.need_full_sync = true
	case "i'm alive":
		/* nothing */
	default:
		log.Printf(`unhandled message: tgt: %s, src: %s, cmd: %s, payload: %v`, m.To, m.From, m.Cmd, m.Payload)
	}
}

type PeerManager struct {
	d              *dht.DHT
	db             *DB
	bus            mangos.Socket
	nick           string
	friends        map[string]*Friend
	oldfriends     map[string]*Friend
	venue          string
	sequence       int64
	need_full_sync bool
	opts           map[string]interface{}
	mtx            *sync.RWMutex
}

func NewPeerManager(db *DB) *PeerManager {
	conf := dht.NewConfig()
	conf.Port = 55000

	pm := PeerManager{db: db}

	pm.opts = map[string]interface{}{mangos.OptionTLSConfig: pm.setupTLS()}

	var err error
	pm.bus, err = bus.NewSocket()
	if err != nil {
		log.Fatalf(`can't create BUS socket: %s`, err)
	}
	pm.bus.AddTransport(tlstcp.NewTransport())

	if pm.d, err = dht.New(conf); err != nil {
		log.Fatalf(`can't create DHT: %s`, err)
	}

	sum := sha1.Sum([]byte("LetsMeetHere"))
	pm.venue = hex.EncodeToString(sum[:])

	buf := make([]byte, 4)
	rand.Read(buf)
	pm.nick = hex.EncodeToString(buf)

	log.Printf(`My nickname is %s`, pm.nick)
	log.Printf(`I will meet my friends at %s`, pm.venue)

	pm.friends = make(map[string]*Friend)
	pm.oldfriends = make(map[string]*Friend)

	pm.mtx = &sync.RWMutex{}

	return &pm
}

func (pm *PeerManager) setupTLS() *tls.Config {
	cert := pm.setupCertificate()

	return &tls.Config{
		InsecureSkipVerify: true,
		Certificates:       []tls.Certificate{cert},
	}
}

func (pm *PeerManager) setupCertificate() tls.Certificate {
	template := x509.Certificate{
		SerialNumber: big.NewInt(4711),
		Subject: pkix.Name{
			Organization: []string{"unobtanium"},
			CommonName:   "*",
		},
		DNSNames:              []string{`*`},
		NotAfter:              time.Now().AddDate(0, 0, 10),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		SignatureAlgorithm:    x509.SHA512WithRSA,
		PublicKeyAlgorithm:    x509.ECDSA,
		SubjectKeyId:          []byte{1, 2, 3, 4, 5},
		BasicConstraintsValid: true,
		IsCA: true,
	}

	log.Printf(`generating 2048 bit RSA private key`)
	priv, _ := rsa.GenerateKey(rand.Reader, 2048)
	pub := &priv.PublicKey
	log.Printf(`building self-signed certificate`)
	cert_raw, err := x509.CreateCertificate(rand.Reader, &template, &template, pub, priv)
	if err != nil {
		log.Fatalf(`can't generate TLS certificate: %s`, err)
	}
	log.Printf(`created a cert of length %d bytes`, len(cert_raw))

	privbuf := &bytes.Buffer{}
	pem.Encode(privbuf, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	certbuf := &bytes.Buffer{}
	pem.Encode(certbuf, &pem.Block{Type: "CERTIFICATE", Bytes: cert_raw})

	cert, err := tls.X509KeyPair(certbuf.Bytes(), privbuf.Bytes())
	if err != nil {
		log.Fatalf(`can't build key pair: %s`, err)
	}

	return cert
}

func (pm *PeerManager) String() string {
	pm.mtx.RLock()
	defer pm.mtx.RUnlock()

	s := []string{
		fmt.Sprintf("DHT: Port: %d", pm.d.Port()),
		fmt.Sprintf("\r\nI am %s, my sequence ID is %d", pm.nick, pm.sequence),
		fmt.Sprintf("I have %d friend(s)", len(pm.friends)),
		fmt.Sprintf("We're meeting at '%s'", pm.venue),
		"\r\nThese are my friends:\r\n",
	}

	for _, f := range pm.friends {
		s = append(s, f.String())
	}

	s = append(s, "\r\nI haven't heard from these guys in a while:\r\n")
	for _, f := range pm.oldfriends {
		s = append(s, f.String())
	}

	return strings.Join(s, "\r\n")
}

func (pm *PeerManager) Publish(dst, cmd, payload string) error {
	pm.mtx.Lock()
	pm.sequence++
	m := &BusMessage{From: pm.nick, To: dst, Seq: pm.sequence, Cmd: cmd, Payload: payload}
	pm.mtx.Unlock()

	return pm.bus.Send(m.Bytes())
}

func (pm *PeerManager) Loop() {
	ih, err := dht.DecodeInfoHash(pm.venue)
	if err != nil {
		log.Printf(`can't decode infohash: %s`, err)
		return
	}

	if err := pm.d.Start(); err != nil {
		log.Printf(`can't start DHT: %s`, err)
		return
	}

	log.Printf(`DHT bound to port %d`, pm.d.Port())

	if err := pm.bus.ListenOptions(fmt.Sprintf("tls+tcp://*:%d", pm.d.Port()), pm.opts); err != nil {
		log.Fatalf(`can't listen on BUS socket: %s`, err)
	}

	go pm.drainPeers()

	go func() {
		for {
			time.Sleep(2 * time.Second)
			pm.mtx.Lock()
			for name, f := range pm.friends {
				f.mtx.Lock()
				if f.lastSeen.Add(10 * time.Second).Before(time.Now()) {
					log.Printf(`haven't heard from %s in a while`, name)
					pm.oldfriends[name] = f
					delete(pm.friends, name)
				}
				f.mtx.Unlock()
			}
			pm.mtx.Unlock()
		}
	}()

	go func() {
		for {
			/* Receive message */
			m, err := NewBusMessage(pm.bus.Recv())
			if err != nil {
				log.Printf(`can't receive message from bus: %s`, err)
				continue
			}

			/* Ignore messages from ourselves */
			if m.From == pm.nick {
				continue
			}

			/* Ignore messages not to 'everyone' or us */
			if m.To != pm.nick && m.To != "*" {
				continue
			}

			pm.mtx.Lock()
			f, ok := pm.friends[m.From]
			if !ok {
				f = NewFriend(m.From, pm)
				pm.friends[m.From] = f
				pm.need_full_sync = true
			}
			pm.mtx.Unlock()

			f.HandleMessage(m)

			if !ok {
				log.Printf(`%s is a new friend!`, m.From)
				pm.Publish(m.From, "hello", "")
			}
		}
	}()

	go func() {
		i := 0
		for {
			pm.Publish("*", "i'm alive", fmt.Sprintf("%d", i))
			i++
			time.Sleep(5 * time.Second)
		}
	}()

	go func() {
		for {
			buf, err := json.Marshal(<-pm.db.Events)
			if err != nil {
				log.Fatalf(`can't marshal event: %s`, err)
			}
			pm.Publish("*", "event", string(buf))
		}
	}()

	go func() {
		for {
			time.Sleep(5 * time.Second)
			pm.mtx.Lock()
			if !pm.need_full_sync {
				pm.mtx.Unlock()
				continue
			}
			pm.need_full_sync = false
			pm.mtx.Unlock()

			log.Println(`doing a full sync`)

			for _, env := range pm.db.AllEnvelopes() {
				_, events, err := pm.db.EnvelopeWithHistory(env.Id)
				if err != nil {
					log.Fatalf(`can't get events for envelope %s: %s`, env.Id, err)
				}
				for _, e := range events {
					pm.db.Events <- e
				}
			}
		}
	}()

	for {
		pm.d.PeersRequest(string(ih), true)
		time.Sleep(5 * time.Second)
	}
}

func (pm *PeerManager) drainPeers() {
	log.Printf(`draining DHT`)
	seen := make(map[string]struct{})

	for r := range pm.d.PeersRequestResults {
		for _, peers := range r {
			for _, x := range peers {
				addr := dht.DecodePeerAddress(x)
				if _, ok := seen[addr]; !ok {
					pm.connectToPeer(addr)
				}
				seen[addr] = struct{}{}
			}
		}
	}
}

func (pm *PeerManager) connectToPeer(addr string) {
	if err := pm.bus.DialOptions(fmt.Sprintf("tls+tcp://%s", addr), pm.opts); err != nil {
		log.Printf(`can't connect SUB to %s: %s`, addr, err)
	}
}
<|endoftext|>


<commit_before>package bagman

import (
    "fmt"
    "io"
    "os"
    "strings"
    "crypto/md5"
	"github.com/diamondap/goamz/aws"
	"github.com/diamondap/goamz/s3"
)

type S3Client struct {
    S3    *s3.S3
}

// Returns an S3Client for the specified region, using
// AWS credentials from the environment. Please keep your AWS
// keys out of the source code repos! Store them somewhere
// else and load them into environment variables AWS_ACCESS_KEY_ID
// and AWS_SECRET_ACCESS_KEY.
func NewS3Client(region aws.Region) (*S3Client, error) {
    auth, err := aws.EnvAuth()
    if err != nil {
        return nil, err
    }
    S3 := s3.New(auth, region)
    return &S3Client{S3}, nil
}

// Returns a list of keys in the specified bucket.
// If limit is zero, this will return all the keys in the bucket;
// otherwise, it will return only the number of keys specifed.
// Note that listing all keys may result in the underlying client
// issuing multiple requests.
func (client *S3Client)ListBucket(bucketName string, limit int) (keys []s3.Key, err error) {
    bucket := client.S3.Bucket(bucketName)
    if bucket == nil {
        err = fmt.Errorf("Cannot retrieve bucket: %s", bucketName)
        return nil, err
    }
    actualLimit := limit
    if limit == 0 {
        actualLimit = 1000
    }
    bucketList, err := bucket.List("", "/", "", actualLimit)
    if err != nil {
        return nil, err
    }
    contents := bucketList.Contents
    if len(contents) == 0 {
        return contents, nil
    }
    for limit == 0 {
        lastKey := contents[len(contents) - 1].Key
        bucketList, err := bucket.List("", "/", lastKey, actualLimit)
        if err != nil {
            return nil, err
        }
        contents = append(contents, bucketList.Contents ...)
        if !bucketList.IsTruncated {
            break
        }
    }
    return contents, nil
}

// Fetches key from bucket and saves it to path.
// This validates the md5 sum of the byte stream before
// saving to disk. If the md5 sum of the downloaded bytes
// does not match the md5 sum in the key, this will not
// save the file. It will just return an error.
func (client *S3Client)FetchToFile(bucketName string, key s3.Key, path string) (fetchResult *FetchResult) {
    bucket := client.S3.Bucket(bucketName)
    result := new(FetchResult)
    result.BucketName = bucketName
    result.Key = key.Key
    result.LocalTarFile = path

    // In general, we want to retry if the fetch operation
    // fails. We will override this in certain cases below.
    result.Retry = true

    // S3 etag is md5 hex string enclosed in quotes,
    // unless file was a multipart upload. See below for that.
    result.RemoteMd5 = strings.Replace(key.ETag, "\"", "", -1)

    // Fetch the file into a reader instead of using the usual bucket.Get().
    // Files may be up to 250GB, so we want to process them as streams.
    // If we get an error here, it's typically a network error, and we
    // will want to retry later.
    readCloser, err := bucket.GetReader(key.Key)
    if readCloser != nil {
        defer readCloser.Close()
    }
    if err != nil {
        result.ErrorMessage = fmt.Sprintf("Error retrieving file from receiving bucket: %v", err)
        if strings.Contains(err.Error(), "key does not exist") {
            result.Retry = false
        }
        return result
    }

    // Write the contents of the stream into both our md5 hasher
    // and the file.
    md5Hash := md5.New()
    outputFile, err := os.Create(path)
    if outputFile != nil {
        defer outputFile.Close()
    }
    if err != nil {
        result.ErrorMessage = fmt.Sprintf("Could not create local file %s: %v", path, err)
        return result
    }

    multiWriter := io.MultiWriter(outputFile, md5Hash)
    bytesWritten, err := io.Copy(multiWriter, readCloser)
    if err != nil {
        result.ErrorMessage = fmt.Sprintf("Error copying file from receiving bucket: %v", err)
        return result
    }
    if bytesWritten != key.Size {
        result.ErrorMessage = fmt.Sprintf("While downloading from receiving bucket, " +
            "copied only %d of %d bytes for %s", bytesWritten, key.Size, key.Key)
        return result
    }

    result.LocalMd5 = fmt.Sprintf("%x", md5Hash.Sum(nil))

    // ETag for S3 multi-part upload is not an accurate md5 sum.
    // If the ETag ends with a dash and some number, it's a
    // multi-part upload.
    if strings.Contains(result.RemoteMd5, "-") {
        result.Warning = fmt.Sprintf("Skipping md5 check on %s: this was a multi-part upload", key.Key)
        result.Md5Verified = false
        result.Md5Verifiable = false
    } else {
        result.Md5Verifiable = true
        result.Md5Verified = true
        if result.LocalMd5 != result.RemoteMd5 {
            os.Remove(path)
            result.ErrorMessage = fmt.Sprintf("Our md5 sum '%x' does not match the S3 md5 sum '%s'",
                result.LocalMd5, result.RemoteMd5)
            result.Md5Verified = false
            // Don't bother reprocessing this item.
            result.Retry = false
        }
    }
    return result
}

// Collects info about all of the buckets listed in buckets.
// TODO: Write unit test
func (client *S3Client) CheckAllBuckets(buckets []string) (bucketSummaries []*BucketSummary, err error) {
    bucketSummaries = make([]*BucketSummary, 0)
    for _, bucketName := range(buckets) {
        bucketSummary, err := client.CheckBucket(bucketName)
        if err != nil {
            return bucketSummaries, err
        }
        bucketSummaries = append(bucketSummaries, bucketSummary)
    }
    return bucketSummaries, nil
}

// Returns info about the contents of the bucket named bucketName.
// BucketSummary contains the bucket name, a list of keys, and the
// size of the largest file in the bucket.
// TODO: Write unit test
func (client *S3Client) CheckBucket(bucketName string) (bucketSummary *BucketSummary, err error) {
    bucket := client.S3.Bucket(bucketName)
    if bucket == nil {
        err = fmt.Errorf("Cannot retrieve bucket: %s", bucketName)
        return nil, err
    }
    bucketSummary = new(BucketSummary)
    bucketSummary.BucketName = bucketName
    bucketSummary.Keys, err = client.ListBucket(bucketName, 0)
    if err != nil {
        return nil, err
    }
    return bucketSummary, nil
}

// Creates an options struct that adds metadata headers to the S3 put.
func (client *S3Client) MakeOptions(md5sum string, metadata map[string][]string) (s3.Options) {
	if md5sum != "" {
		return s3.Options{
			ContentMD5: md5sum,
			Meta: metadata,
		}
	} else {
		return s3.Options{
			Meta: metadata,
		}
	}
}

// Saves a file to S3 with default access of Private.
// The underlying S3 client does not return the md5 checksum
// from s3, but we already have this info elsewhere. If the
// PUT produces no error, we assume the copy worked and the
// files md5 sum is the same on S3 as here.
func (client *S3Client) SaveToS3(bucketName, fileName, contentType string, reader io.Reader, byteCount int64, options s3.Options) (url string, err error) {
    bucket := client.S3.Bucket(bucketName)
    putErr := bucket.PutReader(fileName, reader, byteCount,
		contentType, s3.Private, options)
    if putErr != nil {
        err = fmt.Errorf("Error saving file '%s' to bucket '%s': %v",
            fileName, bucketName, putErr)
		return "", err
    }
	url = fmt.Sprintf("https://s3.amazonaws.com/%s/%s", bucketName, fileName)
    return url, nil
}

// Returns an S3 key object for the specified file in the
// specified bucket. The key object has the ETag, last mod
// date, size and other useful info.
func (client *S3Client) GetKey(bucketName, fileName string) (*s3.Key, error) {
    bucket := client.S3.Bucket(bucketName)
    listResp, err := bucket.List(fileName, "", "", 1)
    if err != nil {
        err = fmt.Errorf("Error checking key '%s' in bucket '%s': '%v'",
            fileName, bucketName, err)
        return nil, err
    }
    if listResp == nil || len(listResp.Contents) < 1 {
        err = fmt.Errorf("Key '%s' not found in bucket '%s'",
            fileName, bucketName)
        return nil, err
    }
    return &listResp.Contents[0], nil
}

// Deletes an item from S3
func (client *S3Client) Delete(bucketName, fileName string) (error) {
    bucket := client.S3.Bucket(bucketName)
    return bucket.Del(fileName)
}

// Sends a large file (>= 5GB) to S3 in 200MB chunks. This operation
// may take several minutes to complete. Note that os.File satisfies
// the s3.ReaderAtSeeker interface.
func (client *S3Client) SaveLargeFileToS3(bucketName, fileName, contentType string,
	reader s3.ReaderAtSeeker, byteCount int64, options s3.Options, chunkSize int64) (url string, err error) {

    bucket := client.S3.Bucket(bucketName)
	multipartPut, err := bucket.InitMulti(fileName, contentType, s3.Private, options)
	if err != nil {
		return "", err
	}

	// Send all of the individual parts to S3 in chunks
	parts, err := multipartPut.PutAll(reader, chunkSize)
	if err != nil {
		return "", err
	}

	// This command tells S3 to stitch all the parts into a single file.
	err = multipartPut.Complete(parts)
	if err != nil {
		return "", err
	}

	resp, err := bucket.Head(fileName, nil)
	if err != nil {
		return "", fmt.Errorf("Files were uploaded to S3, but attempt to " +
			"confirm metadata returned this error: %v", err)
	}

	// Make sure all the meta data made it there.
	// Var metadata is the metadata we sent to S3.
	metadata := options.Meta
	notVerified := ""

	if !metadataMatches(metadata, "institution", resp.Header, "X-Amz-Meta-Institution") {
		notVerified += "institution, "
	}
	if !metadataMatches(metadata, "bag", resp.Header, "X-Amz-Meta-Bag") {
		notVerified += "bag, "
	}
	if !metadataMatches(metadata, "bagpath", resp.Header, "X-Amz-Meta-Bagpath") {
		notVerified += "bagpath, "
	}
	if !metadataMatches(metadata, "md5", resp.Header, "X-Amz-Meta-Md5") {
		notVerified += "md5"
	}
	if len(notVerified) > 0 {
		return "", fmt.Errorf("Multi-part upload succeeded, but S3 does not return " +
			"the following metadata: %s", notVerified)
	}

	url = fmt.Sprintf("https://s3.amazonaws.com/%s/%s", bucketName, fileName)
    return url, nil
}

func metadataMatches(metadata map[string][]string, key string, s3headers map[string][]string, headerName string)(bool) {
	metaValue, keyExists := metadata[key]
	headerValue, headerExists := s3headers[headerName]

	// If we didn't send this metadata in the first place, we
	// don't care if S3 has it.
	if !keyExists {
		return true
	}

	// If we sent the metadata, test whether S3 returned
	// what we sent.
	if keyExists && len(metaValue) > 0 && headerExists && len(headerValue) > 0 {
		return metaValue[0] == headerValue[0]
	}

	// If we get here, the key exists in the metadata we
	// sent, but not in the S3 headers.
	return false
}
<commit_msg>Explicitly failed abort multi-part puts<commit_after>package bagman

import (
    "fmt"
    "io"
    "os"
    "strings"
    "crypto/md5"
	"github.com/diamondap/goamz/aws"
	"github.com/diamondap/goamz/s3"
)

type S3Client struct {
    S3    *s3.S3
}

// Returns an S3Client for the specified region, using
// AWS credentials from the environment. Please keep your AWS
// keys out of the source code repos! Store them somewhere
// else and load them into environment variables AWS_ACCESS_KEY_ID
// and AWS_SECRET_ACCESS_KEY.
func NewS3Client(region aws.Region) (*S3Client, error) {
    auth, err := aws.EnvAuth()
    if err != nil {
        return nil, err
    }
    S3 := s3.New(auth, region)
    return &S3Client{S3}, nil
}

// Returns a list of keys in the specified bucket.
// If limit is zero, this will return all the keys in the bucket;
// otherwise, it will return only the number of keys specifed.
// Note that listing all keys may result in the underlying client
// issuing multiple requests.
func (client *S3Client)ListBucket(bucketName string, limit int) (keys []s3.Key, err error) {
    bucket := client.S3.Bucket(bucketName)
    if bucket == nil {
        err = fmt.Errorf("Cannot retrieve bucket: %s", bucketName)
        return nil, err
    }
    actualLimit := limit
    if limit == 0 {
        actualLimit = 1000
    }
    bucketList, err := bucket.List("", "/", "", actualLimit)
    if err != nil {
        return nil, err
    }
    contents := bucketList.Contents
    if len(contents) == 0 {
        return contents, nil
    }
    for limit == 0 {
        lastKey := contents[len(contents) - 1].Key
        bucketList, err := bucket.List("", "/", lastKey, actualLimit)
        if err != nil {
            return nil, err
        }
        contents = append(contents, bucketList.Contents ...)
        if !bucketList.IsTruncated {
            break
        }
    }
    return contents, nil
}

// Fetches key from bucket and saves it to path.
// This validates the md5 sum of the byte stream before
// saving to disk. If the md5 sum of the downloaded bytes
// does not match the md5 sum in the key, this will not
// save the file. It will just return an error.
func (client *S3Client)FetchToFile(bucketName string, key s3.Key, path string) (fetchResult *FetchResult) {
    bucket := client.S3.Bucket(bucketName)
    result := new(FetchResult)
    result.BucketName = bucketName
    result.Key = key.Key
    result.LocalTarFile = path

    // In general, we want to retry if the fetch operation
    // fails. We will override this in certain cases below.
    result.Retry = true

    // S3 etag is md5 hex string enclosed in quotes,
    // unless file was a multipart upload. See below for that.
    result.RemoteMd5 = strings.Replace(key.ETag, "\"", "", -1)

    // Fetch the file into a reader instead of using the usual bucket.Get().
    // Files may be up to 250GB, so we want to process them as streams.
    // If we get an error here, it's typically a network error, and we
    // will want to retry later.
    readCloser, err := bucket.GetReader(key.Key)
    if readCloser != nil {
        defer readCloser.Close()
    }
    if err != nil {
        result.ErrorMessage = fmt.Sprintf("Error retrieving file from receiving bucket: %v", err)
        if strings.Contains(err.Error(), "key does not exist") {
            result.Retry = false
        }
        return result
    }

    // Write the contents of the stream into both our md5 hasher
    // and the file.
    md5Hash := md5.New()
    outputFile, err := os.Create(path)
    if outputFile != nil {
        defer outputFile.Close()
    }
    if err != nil {
        result.ErrorMessage = fmt.Sprintf("Could not create local file %s: %v", path, err)
        return result
    }

    multiWriter := io.MultiWriter(outputFile, md5Hash)
    bytesWritten, err := io.Copy(multiWriter, readCloser)
    if err != nil {
        result.ErrorMessage = fmt.Sprintf("Error copying file from receiving bucket: %v", err)
        return result
    }
    if bytesWritten != key.Size {
        result.ErrorMessage = fmt.Sprintf("While downloading from receiving bucket, " +
            "copied only %d of %d bytes for %s", bytesWritten, key.Size, key.Key)
        return result
    }

    result.LocalMd5 = fmt.Sprintf("%x", md5Hash.Sum(nil))

    // ETag for S3 multi-part upload is not an accurate md5 sum.
    // If the ETag ends with a dash and some number, it's a
    // multi-part upload.
    if strings.Contains(result.RemoteMd5, "-") {
        result.Warning = fmt.Sprintf("Skipping md5 check on %s: this was a multi-part upload", key.Key)
        result.Md5Verified = false
        result.Md5Verifiable = false
    } else {
        result.Md5Verifiable = true
        result.Md5Verified = true
        if result.LocalMd5 != result.RemoteMd5 {
            os.Remove(path)
            result.ErrorMessage = fmt.Sprintf("Our md5 sum '%x' does not match the S3 md5 sum '%s'",
                result.LocalMd5, result.RemoteMd5)
            result.Md5Verified = false
            // Don't bother reprocessing this item.
            result.Retry = false
        }
    }
    return result
}

// Collects info about all of the buckets listed in buckets.
// TODO: Write unit test
func (client *S3Client) CheckAllBuckets(buckets []string) (bucketSummaries []*BucketSummary, err error) {
    bucketSummaries = make([]*BucketSummary, 0)
    for _, bucketName := range(buckets) {
        bucketSummary, err := client.CheckBucket(bucketName)
        if err != nil {
            return bucketSummaries, err
        }
        bucketSummaries = append(bucketSummaries, bucketSummary)
    }
    return bucketSummaries, nil
}

// Returns info about the contents of the bucket named bucketName.
// BucketSummary contains the bucket name, a list of keys, and the
// size of the largest file in the bucket.
// TODO: Write unit test
func (client *S3Client) CheckBucket(bucketName string) (bucketSummary *BucketSummary, err error) {
    bucket := client.S3.Bucket(bucketName)
    if bucket == nil {
        err = fmt.Errorf("Cannot retrieve bucket: %s", bucketName)
        return nil, err
    }
    bucketSummary = new(BucketSummary)
    bucketSummary.BucketName = bucketName
    bucketSummary.Keys, err = client.ListBucket(bucketName, 0)
    if err != nil {
        return nil, err
    }
    return bucketSummary, nil
}

// Creates an options struct that adds metadata headers to the S3 put.
func (client *S3Client) MakeOptions(md5sum string, metadata map[string][]string) (s3.Options) {
	if md5sum != "" {
		return s3.Options{
			ContentMD5: md5sum,
			Meta: metadata,
		}
	} else {
		return s3.Options{
			Meta: metadata,
		}
	}
}

// Saves a file to S3 with default access of Private.
// The underlying S3 client does not return the md5 checksum
// from s3, but we already have this info elsewhere. If the
// PUT produces no error, we assume the copy worked and the
// files md5 sum is the same on S3 as here.
func (client *S3Client) SaveToS3(bucketName, fileName, contentType string, reader io.Reader, byteCount int64, options s3.Options) (url string, err error) {
    bucket := client.S3.Bucket(bucketName)
    putErr := bucket.PutReader(fileName, reader, byteCount,
		contentType, s3.Private, options)
    if putErr != nil {
        err = fmt.Errorf("Error saving file '%s' to bucket '%s': %v",
            fileName, bucketName, putErr)
		return "", err
    }
	url = fmt.Sprintf("https://s3.amazonaws.com/%s/%s", bucketName, fileName)
    return url, nil
}

// Returns an S3 key object for the specified file in the
// specified bucket. The key object has the ETag, last mod
// date, size and other useful info.
func (client *S3Client) GetKey(bucketName, fileName string) (*s3.Key, error) {
    bucket := client.S3.Bucket(bucketName)
    listResp, err := bucket.List(fileName, "", "", 1)
    if err != nil {
        err = fmt.Errorf("Error checking key '%s' in bucket '%s': '%v'",
            fileName, bucketName, err)
        return nil, err
    }
    if listResp == nil || len(listResp.Contents) < 1 {
        err = fmt.Errorf("Key '%s' not found in bucket '%s'",
            fileName, bucketName)
        return nil, err
    }
    return &listResp.Contents[0], nil
}

// Deletes an item from S3
func (client *S3Client) Delete(bucketName, fileName string) (error) {
    bucket := client.S3.Bucket(bucketName)
    return bucket.Del(fileName)
}

// Sends a large file (>= 5GB) to S3 in 200MB chunks. This operation
// may take several minutes to complete. Note that os.File satisfies
// the s3.ReaderAtSeeker interface.
func (client *S3Client) SaveLargeFileToS3(bucketName, fileName, contentType string,
	reader s3.ReaderAtSeeker, byteCount int64, options s3.Options, chunkSize int64) (url string, err error) {

    bucket := client.S3.Bucket(bucketName)
	multipartPut, err := bucket.InitMulti(fileName, contentType, s3.Private, options)
	if err != nil {
		return "", err
	}

	// Send all of the individual parts to S3 in chunks
	parts, err := multipartPut.PutAll(reader, chunkSize)
	if err != nil {
		abortErr := multipartPut.Abort()
		if abortErr != nil {
			return "", fmt.Errorf("Multipart put failed with error %v " +
				"while uploading a part and abort failed with error %v",
				err, abortErr)
		}
		return "", err
	}

	// This command tells S3 to stitch all the parts into a single file.
	err = multipartPut.Complete(parts)
	if err != nil {
		abortErr := multipartPut.Abort()
		if abortErr != nil {
			return "", fmt.Errorf("Multipart put failed in 'complete' stage " +
				"with error %v and abort failed with error %v",
				err, abortErr)
		}
		return "", err
	}

	resp, err := bucket.Head(fileName, nil)
	if err != nil {
		return "", fmt.Errorf("Files were uploaded to S3, but attempt to " +
			"confirm metadata returned this error: %v", err)
	}

	// Make sure all the meta data made it there.
	// Var metadata is the metadata we sent to S3.
	metadata := options.Meta
	notVerified := ""

	if !metadataMatches(metadata, "institution", resp.Header, "X-Amz-Meta-Institution") {
		notVerified += "institution, "
	}
	if !metadataMatches(metadata, "bag", resp.Header, "X-Amz-Meta-Bag") {
		notVerified += "bag, "
	}
	if !metadataMatches(metadata, "bagpath", resp.Header, "X-Amz-Meta-Bagpath") {
		notVerified += "bagpath, "
	}
	if !metadataMatches(metadata, "md5", resp.Header, "X-Amz-Meta-Md5") {
		notVerified += "md5"
	}
	if len(notVerified) > 0 {
		return "", fmt.Errorf("Multi-part upload succeeded, but S3 does not return " +
			"the following metadata: %s", notVerified)
	}

	url = fmt.Sprintf("https://s3.amazonaws.com/%s/%s", bucketName, fileName)
    return url, nil
}

func metadataMatches(metadata map[string][]string, key string, s3headers map[string][]string, headerName string)(bool) {
	metaValue, keyExists := metadata[key]
	headerValue, headerExists := s3headers[headerName]

	// If we didn't send this metadata in the first place, we
	// don't care if S3 has it.
	if !keyExists {
		return true
	}

	// If we sent the metadata, test whether S3 returned
	// what we sent.
	if keyExists && len(metaValue) > 0 && headerExists && len(headerValue) > 0 {
		return metaValue[0] == headerValue[0]
	}

	// If we get here, the key exists in the metadata we
	// sent, but not in the S3 headers.
	return false
}
<|endoftext|>


<commit_before>package pg

import (
	"errors"
	"io"
	"time"

	"github.com/go-pg/pg/internal"
	"github.com/go-pg/pg/internal/pool"
	"github.com/go-pg/pg/orm"
)

var errTxDone = errors.New("pg: transaction has already been committed or rolled back")

// Tx is an in-progress database transaction.
//
// A transaction must end with a call to Commit or Rollback.
//
// After a call to Commit or Rollback, all operations on the transaction fail
// with ErrTxDone.
//
// The statements prepared for a transaction by calling the transaction's
// Prepare or Stmt methods are closed by the call to Commit or Rollback.
type Tx struct {
	db *DB
	cn *pool.Conn

	stmts []*Stmt
}

var _ orm.DB = (*Tx)(nil)

// Begin starts a transaction. Most callers should use RunInTransaction instead.
func (db *DB) Begin() (*Tx, error) {
	tx := &Tx{
		db: db,
	}

	cn, err := db.conn()
	if err != nil {
		return nil, err
	}
	tx.cn = cn

	if err := tx.begin(); err != nil {
		return nil, err
	}

	return tx, nil
}

// RunInTransaction runs a function in a transaction. If function
// returns an error transaction is rollbacked, otherwise transaction
// is committed.
func (db *DB) RunInTransaction(fn func(*Tx) error) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	return tx.RunInTransaction(fn)
}

// Begin returns the transaction.
func (tx *Tx) Begin() (*Tx, error) {
	return tx, nil
}

// RunInTransaction runs a function in the transaction. If function
// returns an error transaction is rollbacked, otherwise transaction
// is committed.
func (tx *Tx) RunInTransaction(fn func(*Tx) error) error {
	defer func() {
		if err := recover(); err != nil {
			tx.Rollback()
			panic(err)
		}
	}()
	if err := fn(tx); err != nil {
		tx.Rollback()
		return err
	}
	return tx.Commit()
}

func (tx *Tx) conn() (*pool.Conn, error) {
	if tx.cn == nil {
		return nil, errTxDone
	}

	tx.cn.SetTimeout(tx.db.opt.ReadTimeout, tx.db.opt.WriteTimeout)
	return tx.cn, nil
}

func (tx *Tx) freeConn(cn *pool.Conn, err error) {}

// Stmt returns a transaction-specific prepared statement from an existing statement.
func (tx *Tx) Stmt(stmt *Stmt) *Stmt {
	stmt, err := tx.Prepare(stmt.q)
	if err != nil {
		return &Stmt{stickyErr: err}
	}
	return stmt
}

// Prepare creates a prepared statement for use within a transaction.
//
// The returned statement operates within the transaction and can no longer
// be used once the transaction has been committed or rolled back.
//
// To use an existing prepared statement on this transaction, see Tx.Stmt.
func (tx *Tx) Prepare(q string) (*Stmt, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	stmt, err := prepare(tx.db, cn, q)
	tx.freeConn(cn, err)
	if err != nil {
		return nil, err
	}

	stmt.inTx = true
	tx.stmts = append(tx.stmts, stmt)

	return stmt, nil
}

// Exec is an alias for DB.Exec.
func (tx *Tx) Exec(query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	start := time.Now()
	res, err := tx.db.simpleQuery(cn, query, params...)
	tx.freeConn(cn, err)
	tx.db.queryProcessed(tx, start, query, params, 0, res, err)

	return res, err
}

// ExecOne is an alias for DB.ExecOne.
func (tx *Tx) ExecOne(query interface{}, params ...interface{}) (orm.Result, error) {
	res, err := tx.Exec(query, params...)
	if err != nil {
		return nil, err
	}

	if err := internal.AssertOneRow(res.RowsAffected()); err != nil {
		return nil, err
	}
	return res, nil
}

// Query is an alias for DB.Query.
func (tx *Tx) Query(model interface{}, query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	start := time.Now()
	res, err := tx.db.simpleQueryData(cn, model, query, params...)
	tx.freeConn(cn, err)
	tx.db.queryProcessed(tx, start, query, params, 0, res, err)

	if err != nil {
		return nil, err
	}

	if mod := res.Model(); mod != nil && res.RowsReturned() > 0 {
		if err = mod.AfterQuery(tx); err != nil {
			return res, err
		}
	}

	return res, err
}

// QueryOne is an alias for DB.QueryOne.
func (tx *Tx) QueryOne(model interface{}, query interface{}, params ...interface{}) (orm.Result, error) {
	mod, err := orm.NewModel(model)
	if err != nil {
		return nil, err
	}

	res, err := tx.Query(mod, query, params...)
	if err != nil {
		return nil, err
	}

	if err := internal.AssertOneRow(res.RowsAffected()); err != nil {
		return nil, err
	}
	return res, nil
}

// Model is an alias for DB.Model.
func (tx *Tx) Model(model ...interface{}) *orm.Query {
	return orm.NewQuery(tx, model...)
}

// Select is an alias for DB.Select.
func (tx *Tx) Select(model interface{}) error {
	return orm.Select(tx, model)
}

// Insert is an alias for DB.Insert.
func (tx *Tx) Insert(model ...interface{}) error {
	return orm.Insert(tx, model...)
}

// Update is an alias for DB.Update.
func (tx *Tx) Update(model ...interface{}) error {
	return orm.Update(tx, model...)
}

// Delete is an alias for DB.Delete.
func (tx *Tx) Delete(model interface{}) error {
	return orm.Delete(tx, model)
}

// CreateTable is an alias for DB.CreateTable.
func (tx *Tx) CreateTable(model interface{}, opt *orm.CreateTableOptions) error {
	_, err := orm.CreateTable(tx, model, opt)
	return err
}

// DropTable is an alias for DB.DropTable.
func (tx *Tx) DropTable(model interface{}, opt *orm.DropTableOptions) error {
	_, err := orm.DropTable(tx, model, opt)
	return err
}

// CopyFrom is an alias for DB.CopyFrom.
func (tx *Tx) CopyFrom(r io.Reader, query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	res, err := tx.db.copyFrom(cn, r, query, params...)
	tx.freeConn(cn, err)
	return res, err
}

// CopyTo is an alias for DB.CopyTo.
func (tx *Tx) CopyTo(w io.Writer, query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	res, err := tx.db.copyTo(cn, w, query, params...)
	tx.freeConn(cn, err)
	return res, err
}

func (tx *Tx) FormatQuery(dst []byte, query string, params ...interface{}) []byte {
	return tx.db.FormatQuery(dst, query, params...)
}

func (tx *Tx) begin() error {
	_, err := tx.Exec("BEGIN")
	if err != nil {
		tx.close(err)
	}
	return err
}

// Commit commits the transaction.
func (tx *Tx) Commit() error {
	_, err := tx.Exec("COMMIT")
	tx.close(err)
	return err
}

// Rollback aborts the transaction.
func (tx *Tx) Rollback() error {
	_, err := tx.Exec("ROLLBACK")
	tx.close(err)
	return err
}

func (tx *Tx) close(lastErr error) error {
	if tx.cn == nil {
		return errTxDone
	}

	for _, stmt := range tx.stmts {
		_ = stmt.Close()
	}
	tx.stmts = nil

	err := tx.db.freeConn(tx.cn, lastErr)
	tx.cn = nil

	return err
}
<commit_msg>add context.Context support for transactions<commit_after>package pg

import (
	"context"
	"errors"
	"io"
	"time"

	"github.com/go-pg/pg/internal"
	"github.com/go-pg/pg/internal/pool"
	"github.com/go-pg/pg/orm"
)

var errTxDone = errors.New("pg: transaction has already been committed or rolled back")

// Tx is an in-progress database transaction.
//
// A transaction must end with a call to Commit or Rollback.
//
// After a call to Commit or Rollback, all operations on the transaction fail
// with ErrTxDone.
//
// The statements prepared for a transaction by calling the transaction's
// Prepare or Stmt methods are closed by the call to Commit or Rollback.
type Tx struct {
	db *DB
	cn *pool.Conn

	stmts []*Stmt
}

var _ orm.DB = (*Tx)(nil)

// Begin starts a transaction. Most callers should use RunInTransaction instead.
func (db *DB) Begin() (*Tx, error) {
	tx := &Tx{
		db: db,
	}

	cn, err := db.conn()
	if err != nil {
		return nil, err
	}
	tx.cn = cn

	if err := tx.begin(); err != nil {
		return nil, err
	}

	return tx, nil
}

// RunInTransaction runs a function in a transaction. If function
// returns an error transaction is rollbacked, otherwise transaction
// is committed.
func (db *DB) RunInTransaction(fn func(*Tx) error) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	return tx.RunInTransaction(fn)
}

// Begin returns the transaction.
func (tx *Tx) Begin() (*Tx, error) {
	return tx, nil
}

// RunInTransaction runs a function in the transaction. If function
// returns an error transaction is rollbacked, otherwise transaction
// is committed.
func (tx *Tx) RunInTransaction(fn func(*Tx) error) error {
	defer func() {
		if err := recover(); err != nil {
			tx.Rollback()
			panic(err)
		}
	}()
	if err := fn(tx); err != nil {
		tx.Rollback()
		return err
	}
	return tx.Commit()
}

func (tx *Tx) conn() (*pool.Conn, error) {
	if tx.cn == nil {
		return nil, errTxDone
	}

	tx.cn.SetTimeout(tx.db.opt.ReadTimeout, tx.db.opt.WriteTimeout)
	return tx.cn, nil
}

func (tx *Tx) freeConn(cn *pool.Conn, err error) {}

// Stmt returns a transaction-specific prepared statement from an existing statement.
func (tx *Tx) Stmt(stmt *Stmt) *Stmt {
	stmt, err := tx.Prepare(stmt.q)
	if err != nil {
		return &Stmt{stickyErr: err}
	}
	return stmt
}

// Prepare creates a prepared statement for use within a transaction.
//
// The returned statement operates within the transaction and can no longer
// be used once the transaction has been committed or rolled back.
//
// To use an existing prepared statement on this transaction, see Tx.Stmt.
func (tx *Tx) Prepare(q string) (*Stmt, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	stmt, err := prepare(tx.db, cn, q)
	tx.freeConn(cn, err)
	if err != nil {
		return nil, err
	}

	stmt.inTx = true
	tx.stmts = append(tx.stmts, stmt)

	return stmt, nil
}

// Exec is an alias for DB.Exec.
func (tx *Tx) Exec(query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	start := time.Now()
	res, err := tx.db.simpleQuery(cn, query, params...)
	tx.freeConn(cn, err)
	tx.db.queryProcessed(tx, start, query, params, 0, res, err)

	return res, err
}

// ExecOne is an alias for DB.ExecOne.
func (tx *Tx) ExecOne(query interface{}, params ...interface{}) (orm.Result, error) {
	res, err := tx.Exec(query, params...)
	if err != nil {
		return nil, err
	}

	if err := internal.AssertOneRow(res.RowsAffected()); err != nil {
		return nil, err
	}
	return res, nil
}

// Query is an alias for DB.Query.
func (tx *Tx) Query(model interface{}, query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	start := time.Now()
	res, err := tx.db.simpleQueryData(cn, model, query, params...)
	tx.freeConn(cn, err)
	tx.db.queryProcessed(tx, start, query, params, 0, res, err)

	if err != nil {
		return nil, err
	}

	if mod := res.Model(); mod != nil && res.RowsReturned() > 0 {
		if err = mod.AfterQuery(tx); err != nil {
			return res, err
		}
	}

	return res, err
}

// QueryOne is an alias for DB.QueryOne.
func (tx *Tx) QueryOne(model interface{}, query interface{}, params ...interface{}) (orm.Result, error) {
	mod, err := orm.NewModel(model)
	if err != nil {
		return nil, err
	}

	res, err := tx.Query(mod, query, params...)
	if err != nil {
		return nil, err
	}

	if err := internal.AssertOneRow(res.RowsAffected()); err != nil {
		return nil, err
	}
	return res, nil
}

// Model is an alias for DB.Model.
func (tx *Tx) Model(model ...interface{}) *orm.Query {
	return orm.NewQuery(tx, model...)
}

// Select is an alias for DB.Select.
func (tx *Tx) Select(model interface{}) error {
	return orm.Select(tx, model)
}

// Insert is an alias for DB.Insert.
func (tx *Tx) Insert(model ...interface{}) error {
	return orm.Insert(tx, model...)
}

// Update is an alias for DB.Update.
func (tx *Tx) Update(model ...interface{}) error {
	return orm.Update(tx, model...)
}

// Delete is an alias for DB.Delete.
func (tx *Tx) Delete(model interface{}) error {
	return orm.Delete(tx, model)
}

// CreateTable is an alias for DB.CreateTable.
func (tx *Tx) CreateTable(model interface{}, opt *orm.CreateTableOptions) error {
	_, err := orm.CreateTable(tx, model, opt)
	return err
}

// DropTable is an alias for DB.DropTable.
func (tx *Tx) DropTable(model interface{}, opt *orm.DropTableOptions) error {
	_, err := orm.DropTable(tx, model, opt)
	return err
}

// CopyFrom is an alias for DB.CopyFrom.
func (tx *Tx) CopyFrom(r io.Reader, query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	res, err := tx.db.copyFrom(cn, r, query, params...)
	tx.freeConn(cn, err)
	return res, err
}

// CopyTo is an alias for DB.CopyTo.
func (tx *Tx) CopyTo(w io.Writer, query interface{}, params ...interface{}) (orm.Result, error) {
	cn, err := tx.conn()
	if err != nil {
		return nil, err
	}

	res, err := tx.db.copyTo(cn, w, query, params...)
	tx.freeConn(cn, err)
	return res, err
}

func (tx *Tx) FormatQuery(dst []byte, query string, params ...interface{}) []byte {
	return tx.db.FormatQuery(dst, query, params...)
}

func (tx *Tx) begin() error {
	_, err := tx.Exec("BEGIN")
	if err != nil {
		tx.close(err)
	}
	return err
}

// Commit commits the transaction.
func (tx *Tx) Commit() error {
	_, err := tx.Exec("COMMIT")
	tx.close(err)
	return err
}

// Rollback aborts the transaction.
func (tx *Tx) Rollback() error {
	_, err := tx.Exec("ROLLBACK")
	tx.close(err)
	return err
}

func (tx *Tx) close(lastErr error) error {
	if tx.cn == nil {
		return errTxDone
	}

	for _, stmt := range tx.stmts {
		_ = stmt.Close()
	}
	tx.stmts = nil

	err := tx.db.freeConn(tx.cn, lastErr)
	tx.cn = nil

	return err
}

func (tx *Tx) Context() context.Context {
	return tx.db.Context()
}

func (tx *Tx) WithContext(ctx context.Context) *Tx {
	cp := *tx
	cp.db = cp.db.WithContext(ctx)
	return &cp
}
<|endoftext|>


<commit_before>package deploystack

import (
	"bufio"
	"fmt"
	"os"
	"sort"
	"strings"
)

// ImageManage promps a user to select a disk type.
func ImageManage(project string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine images to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about machine images.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/images%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	colorPrintln("Choose an operating system.", TERMCYANB)
	ImageTypeProject := listSelect(DiskProjects, DefaultImageProject)

	fmt.Printf("Polling for %s images...\n", ImageTypeProject.Value)
	images, err := images(project, ImageTypeProject.Value)
	if err != nil {
		return "", err
	}

	families := getListOfImageFamilies(images)

	colorPrintln("Choose a disk family to use for this application.", TERMCYANB)
	family := listSelect(families, DefaultImageFamily)

	imagesByFam := getListOfImageTypesByFamily(images, ImageTypeProject.Value, family.Value)

	colorPrintln("Choose a disk type to use for this application.", TERMCYANB)
	result := listSelect(imagesByFam, imagesByFam[len(imagesByFam)-1].Value)

	return result.Value, nil
}

func colorPrintln(msg, color string) {
	fmt.Printf("%s%s %s\n", color, msg, TERMCLEAR)
}

func MachineTypeManage(project, zone string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine types to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about Machine types.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/machine-types%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	fmt.Printf("Polling for machine types...\n")
	types, err := machineTypes(project, zone)
	if err != nil {
		return "", fmt.Errorf("error polling for machine types : %s", err)
	}

	typefamilies := getListOfMachineTypeFamily(types)

	fmt.Printf("Choose an Machine Type Family\n")
	familyProject := listSelect(typefamilies, DefaultMachineType)

	filteredtypes := getListOfMachineTypeByFamily(types, familyProject.Value)

	fmt.Printf("%sChoose a machine type to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	result := listSelect(filteredtypes, filteredtypes[0].Value)

	return result.Value, nil
}

type GCEInstanceConfig map[string]string

func (gce GCEInstanceConfig) Print(title string) {
	keys := []string{}
	for i := range gce {
		keys = append(keys, i)
	}

	longest := longestLength(toLabeledValueSlice(keys))

	colorPrintln(title, TERMCYANREV)
	exclude := []string{}

	if s, ok := gce["instance-name"]; ok && len(s) > 0 {
		printSetting("instance-name", s, longest)
		exclude = append(exclude, "instance-name")
	}

	if s, ok := gce["region"]; ok && len(s) > 0 {
		printSetting("region", s, longest)
		exclude = append(exclude, "region")
	}

	if s, ok := gce["zone"]; ok && len(s) > 0 {
		printSetting("zone", s, longest)
		exclude = append(exclude, "zone")
	}

	ordered := []string{}
	for i, v := range gce {
		if strings.Contains(strings.Join(exclude, " "), i) {
			continue
		}
		if len(v) < 1 {
			continue
		}

		ordered = append(ordered, i)
	}
	sort.Strings(ordered)

	for i := range ordered {
		key := ordered[i]
		printSetting(key, gce[key], longest)
	}
}

func GCEInstanceManage(project, basename string) (GCEInstanceConfig, error) {
	var err error
	configs := make(map[string]string)

	defaultImage, err := getLatestImage(project, DefaultImageProject, DefaultImageFamily)
	if err != nil {
		return configs, err
	}

	defaultConfig := GCEInstanceConfig{
		"instance-image":        defaultImage,
		"instance-disksize":     "200",
		"instance-disktype":     "pd-standard",
		"instance-tags":         "[http-server,https-server]",
		"instance-name":         fmt.Sprintf("%s-instance", basename),
		"region":                DefaultRegion,
		"zone":                  fmt.Sprintf("%s-a", DefaultRegion),
		"instance-machine-type": "n1-standard-1",
	}

	ClearScreen()
	fmt.Println(Divider)
	colorPrintln("Configure a Compute Engine Instance", TERMCYANB)
	fmt.Printf("Let's walk through configuring a Compute Engine Instance (Virtual Machine). \n")
	fmt.Printf("you can either accept a default configuration with settings that work for \n")
	fmt.Printf("trying out most use cases, or hand configure key settings. \n")
	fmt.Println(Divider)

	defaultConfig.Print("Default Configuration")

	chooseDefault := Custom{
		Name:        "choosedefault",
		Description: "Do you want to use the default? ('No' means custom)",
		Default:     "yes",
		Validation:  "yesorno",
	}

	if err := chooseDefault.Collect(); err != nil {
		return configs, err
	}

	if chooseDefault.Value == "yes" {
		return defaultConfig, nil
	}

	nameItem := Custom{
		Name:        "name",
		Description: "Enter the name of the instance",
		Default:     fmt.Sprintf("%s-instance", basename),
	}

	if err := nameItem.Collect(); err != nil {
		return configs, err
	}

	configs["instance-name"] = nameItem.Value

	configs["region"], err = RegionManage(project, "compute", DefaultRegion)
	if err != nil {
		return configs, err
	}

	configs["zone"], err = ZoneManage(project, configs["region"])
	if err != nil {
		return configs, err
	}

	configs["instance-machine-type"], err = MachineTypeManage(project, configs["zone"])
	if err != nil {
		return configs, err
	}
	configs["instance-image"], err = ImageManage(project)
	if err != nil {
		return configs, err
	}

	items := Customs{
		{Name: "instance-disksize", Description: "Enter the size of the boot disk you want in GB", Default: "200", Validation: "integer"},
		{Name: "instance-disktype", Description: "Enter the type of the boot disk you want", Default: "pd-standard", Options: []string{"pd-standard", "pd-balanced", "pd-ssd"}},
		{Name: "webserver", Description: "Do you want this to be a webserver (Expose ports 80 & 443)? ", Default: "no", Validation: "yesorno"},
	}

	if err := items.Collect(); err != nil {
		return configs, err
	}

	for _, v := range items {

		if v.Name == "webserver" {
			configs["instance-tags"] = "[]"
			if v.Value == "yes" {
				configs["instance-tags"] = "[http-server,https-server]"
			}
			continue
		}

		configs[v.Name] = v.Value

	}

	return configs, nil
}

// BillingAccountManage either grabs the users only BillingAccount or
// presents a list of BillingAccounts to select from.
func BillingAccountManage() (string, error) {
	accounts, err := billingAccounts()
	if err != nil {
		return "", fmt.Errorf("could not get list of billing accounts: %s", err)
	}

	labeled := []string{}

	for _, v := range accounts {
		labeled = append(labeled, fmt.Sprintf("%s (%s)", v.DisplayName, strings.ReplaceAll(v.Name, "billingAccounts/", "")))
	}

	if len(accounts) == 1 {
		fmt.Printf("\nOnly found one billing account. Using : %s%s%s.\n", TERMCYAN, accounts[0].DisplayName, TERMCLEAR)
		return extractAccount(labeled[0]), nil
	}

	fmt.Printf("\n%sPlease select one of your billing accounts to use with this project%s.\n", TERMCYAN, TERMCLEAR)
	result := listSelect(toLabeledValueSlice(labeled), labeled[0])

	return extractAccount(result.Value), nil
}

func extractAccount(s string) string {
	sl := strings.Split(s, "(")
	return strings.ReplaceAll(sl[1], ")", "")
}

// ProjectManage promps a user to select a project.
func ProjectManage() (string, error) {
	createString := "CREATE NEW PROJECT"
	project, err := ProjectID()
	if err != nil {
		return "", err
	}

	projects, err := projects()
	if err != nil {
		return "", err
	}

	projdis := []string{}

	for _, v := range projects {
		if strings.Contains(v, "Billing Disabled") {
			v = fmt.Sprintf("%s%s%s", TERMGREY, v, TERMCLEAR)
		}
		projdis = append(projdis, v)
	}

	projdis = append([]string{createString}, projdis...)

	fmt.Printf("\n%sChoose a project to use for this application.%s\n\n", TERMCYANB, TERMCLEAR)
	fmt.Printf("%sNOTE:%s This app will make changes to the project. %s\n", TERMCYANREV, TERMCYAN, TERMCLEAR)
	fmt.Printf("While those changes are reverseable, it would be better to put it in a fresh new project. \n")

	project = listSelect(toLabeledValueSlice(projdis), project).Value

	if project == createString {
		project, err = projectPrompt()
		if err != nil {
			return "", err
		}
	}

	if err := ProjectIDSet(project); err != nil {
		return project, fmt.Errorf("error: unable to set project (%s) in environment: %s", project, err)
	}

	return project, nil
}

// projectPrompt manages the interaction of creating a project, including prompts.
func projectPrompt() (string, error) {
	result := ""
	sec1 := NewSection("Creating the project")

	sec1.Open()
	fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)

	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Print("> ")
		text, _ := reader.ReadString('\n')
		text = strings.Replace(text, "\n", "", -1)

		if len(text) == 0 {
			fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)
			continue
		}

		if err := projectCreate(text); err != nil {
			fmt.Printf("%sProject name could not be created %s\n", TERMREDREV, TERMCLEAR)
			fmt.Printf("%sReason: %s %s\n", TERMREDB, err, TERMCLEAR)
			fmt.Printf("%sPlease choose another. %s\n", TERMREDREV, TERMCLEAR)
			continue
		}

		fmt.Printf("Project Created\n")
		result = text
		break

	}
	sec1.Close()

	sec2 := NewSection("Activating Billing for the project")
	sec2.Open()
	account, err := BillingAccountManage()
	if err != nil {
		return "", fmt.Errorf("could not determine proper billing account: %s ", err)
	}

	if err := BillingAccountProjectAttach(result, account); err != nil {
		return "", fmt.Errorf("could not link billing account: %s ", err)
	}
	sec2.Close()
	return result, nil
}

// regions will return a list of regions depending on product type
func regions(project, product string) ([]string, error) {
	switch product {
	case "compute":
		return regionsCompute(project)
	case "functions":
		return regionsFunctions(project)
	case "run":
		return regionsRun(project)
	}

	return []string{}, fmt.Errorf("invalid product requested: %s", product)
}

// RegionManage promps a user to select a region.
func RegionManage(project, product, def string) (string, error) {
	fmt.Printf("Polling for regions...\n")
	regions, err := regions(project, product)
	if err != nil {
		return "", err
	}
	fmt.Printf("%sChoose a valid region to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	region := listSelect(toLabeledValueSlice(regions), def)

	return region.Value, nil
}

// ZoneManage promps a user to select a zone.
func ZoneManage(project, region string) (string, error) {
	fmt.Printf("Polling for zones...\n")
	zones, err := zones(project, region)
	if err != nil {
		return "", err
	}

	fmt.Printf("%sChoose a valid zone to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	zone := listSelect(toLabeledValueSlice(zones), zones[0])
	return zone.Value, nil
}
<commit_msg>Making Project name rules more clear.<commit_after>package deploystack

import (
	"bufio"
	"fmt"
	"os"
	"sort"
	"strings"
)

// ImageManage promps a user to select a disk type.
func ImageManage(project string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine images to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about machine images.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/images%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	colorPrintln("Choose an operating system.", TERMCYANB)
	ImageTypeProject := listSelect(DiskProjects, DefaultImageProject)

	fmt.Printf("Polling for %s images...\n", ImageTypeProject.Value)
	images, err := images(project, ImageTypeProject.Value)
	if err != nil {
		return "", err
	}

	families := getListOfImageFamilies(images)

	colorPrintln("Choose a disk family to use for this application.", TERMCYANB)
	family := listSelect(families, DefaultImageFamily)

	imagesByFam := getListOfImageTypesByFamily(images, ImageTypeProject.Value, family.Value)

	colorPrintln("Choose a disk type to use for this application.", TERMCYANB)
	result := listSelect(imagesByFam, imagesByFam[len(imagesByFam)-1].Value)

	return result.Value, nil
}

func colorPrintln(msg, color string) {
	fmt.Printf("%s%s %s\n", color, msg, TERMCLEAR)
}

func MachineTypeManage(project, zone string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine types to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about Machine types.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/machine-types%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	fmt.Printf("Polling for machine types...\n")
	types, err := machineTypes(project, zone)
	if err != nil {
		return "", fmt.Errorf("error polling for machine types : %s", err)
	}

	typefamilies := getListOfMachineTypeFamily(types)

	fmt.Printf("Choose an Machine Type Family\n")
	familyProject := listSelect(typefamilies, DefaultMachineType)

	filteredtypes := getListOfMachineTypeByFamily(types, familyProject.Value)

	fmt.Printf("%sChoose a machine type to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	result := listSelect(filteredtypes, filteredtypes[0].Value)

	return result.Value, nil
}

type GCEInstanceConfig map[string]string

func (gce GCEInstanceConfig) Print(title string) {
	keys := []string{}
	for i := range gce {
		keys = append(keys, i)
	}

	longest := longestLength(toLabeledValueSlice(keys))

	colorPrintln(title, TERMCYANREV)
	exclude := []string{}

	if s, ok := gce["instance-name"]; ok && len(s) > 0 {
		printSetting("instance-name", s, longest)
		exclude = append(exclude, "instance-name")
	}

	if s, ok := gce["region"]; ok && len(s) > 0 {
		printSetting("region", s, longest)
		exclude = append(exclude, "region")
	}

	if s, ok := gce["zone"]; ok && len(s) > 0 {
		printSetting("zone", s, longest)
		exclude = append(exclude, "zone")
	}

	ordered := []string{}
	for i, v := range gce {
		if strings.Contains(strings.Join(exclude, " "), i) {
			continue
		}
		if len(v) < 1 {
			continue
		}

		ordered = append(ordered, i)
	}
	sort.Strings(ordered)

	for i := range ordered {
		key := ordered[i]
		printSetting(key, gce[key], longest)
	}
}

func GCEInstanceManage(project, basename string) (GCEInstanceConfig, error) {
	var err error
	configs := make(map[string]string)

	defaultImage, err := getLatestImage(project, DefaultImageProject, DefaultImageFamily)
	if err != nil {
		return configs, err
	}

	defaultConfig := GCEInstanceConfig{
		"instance-image":        defaultImage,
		"instance-disksize":     "200",
		"instance-disktype":     "pd-standard",
		"instance-tags":         "[http-server,https-server]",
		"instance-name":         fmt.Sprintf("%s-instance", basename),
		"region":                DefaultRegion,
		"zone":                  fmt.Sprintf("%s-a", DefaultRegion),
		"instance-machine-type": "n1-standard-1",
	}

	ClearScreen()
	fmt.Println(Divider)
	colorPrintln("Configure a Compute Engine Instance", TERMCYANB)
	fmt.Printf("Let's walk through configuring a Compute Engine Instance (Virtual Machine). \n")
	fmt.Printf("you can either accept a default configuration with settings that work for \n")
	fmt.Printf("trying out most use cases, or hand configure key settings. \n")
	fmt.Println(Divider)

	defaultConfig.Print("Default Configuration")

	chooseDefault := Custom{
		Name:        "choosedefault",
		Description: "Do you want to use the default? ('No' means custom)",
		Default:     "yes",
		Validation:  "yesorno",
	}

	if err := chooseDefault.Collect(); err != nil {
		return configs, err
	}

	if chooseDefault.Value == "yes" {
		return defaultConfig, nil
	}

	nameItem := Custom{
		Name:        "name",
		Description: "Enter the name of the instance",
		Default:     fmt.Sprintf("%s-instance", basename),
	}

	if err := nameItem.Collect(); err != nil {
		return configs, err
	}

	configs["instance-name"] = nameItem.Value

	configs["region"], err = RegionManage(project, "compute", DefaultRegion)
	if err != nil {
		return configs, err
	}

	configs["zone"], err = ZoneManage(project, configs["region"])
	if err != nil {
		return configs, err
	}

	configs["instance-machine-type"], err = MachineTypeManage(project, configs["zone"])
	if err != nil {
		return configs, err
	}
	configs["instance-image"], err = ImageManage(project)
	if err != nil {
		return configs, err
	}

	items := Customs{
		{Name: "instance-disksize", Description: "Enter the size of the boot disk you want in GB", Default: "200", Validation: "integer"},
		{Name: "instance-disktype", Description: "Enter the type of the boot disk you want", Default: "pd-standard", Options: []string{"pd-standard", "pd-balanced", "pd-ssd"}},
		{Name: "webserver", Description: "Do you want this to be a webserver (Expose ports 80 & 443)? ", Default: "no", Validation: "yesorno"},
	}

	if err := items.Collect(); err != nil {
		return configs, err
	}

	for _, v := range items {

		if v.Name == "webserver" {
			configs["instance-tags"] = "[]"
			if v.Value == "yes" {
				configs["instance-tags"] = "[http-server,https-server]"
			}
			continue
		}

		configs[v.Name] = v.Value

	}

	return configs, nil
}

// BillingAccountManage either grabs the users only BillingAccount or
// presents a list of BillingAccounts to select from.
func BillingAccountManage() (string, error) {
	accounts, err := billingAccounts()
	if err != nil {
		return "", fmt.Errorf("could not get list of billing accounts: %s", err)
	}

	labeled := []string{}

	for _, v := range accounts {
		labeled = append(labeled, fmt.Sprintf("%s (%s)", v.DisplayName, strings.ReplaceAll(v.Name, "billingAccounts/", "")))
	}

	if len(accounts) == 1 {
		fmt.Printf("\nOnly found one billing account. Using : %s%s%s.\n", TERMCYAN, accounts[0].DisplayName, TERMCLEAR)
		return extractAccount(labeled[0]), nil
	}

	fmt.Printf("\n%sPlease select one of your billing accounts to use with this project%s.\n", TERMCYAN, TERMCLEAR)
	result := listSelect(toLabeledValueSlice(labeled), labeled[0])

	return extractAccount(result.Value), nil
}

func extractAccount(s string) string {
	sl := strings.Split(s, "(")
	return strings.ReplaceAll(sl[1], ")", "")
}

// ProjectManage promps a user to select a project.
func ProjectManage() (string, error) {
	createString := "CREATE NEW PROJECT"
	project, err := ProjectID()
	if err != nil {
		return "", err
	}

	projects, err := projects()
	if err != nil {
		return "", err
	}

	projdis := []string{}

	for _, v := range projects {
		if strings.Contains(v, "Billing Disabled") {
			v = fmt.Sprintf("%s%s%s", TERMGREY, v, TERMCLEAR)
		}
		projdis = append(projdis, v)
	}

	projdis = append([]string{createString}, projdis...)

	fmt.Printf("\n%sChoose a project to use for this application.%s\n\n", TERMCYANB, TERMCLEAR)
	fmt.Printf("%sNOTE:%s This app will make changes to the project. %s\n", TERMCYANREV, TERMCYAN, TERMCLEAR)
	fmt.Printf("While those changes are reverseable, it would be better to put it in a fresh new project. \n")

	project = listSelect(toLabeledValueSlice(projdis), project).Value

	if project == createString {
		project, err = projectPrompt()
		if err != nil {
			return "", err
		}
	}

	if err := ProjectIDSet(project); err != nil {
		return project, fmt.Errorf("error: unable to set project (%s) in environment: %s", project, err)
	}

	return project, nil
}

// projectPrompt manages the interaction of creating a project, including prompts.
func projectPrompt() (string, error) {
	result := ""
	sec1 := NewSection("Creating the project")

	sec1.Open()
	fmt.Printf("Project IDs are immutable and can be set only during project \n")
	fmt.Printf("creation. They must start with a lowercase letter and can have \n")
	fmt.Printf("lowercase ASCII letters, digits or hyphens.  \n")
	fmt.Printf("Project IDs must be between 6 and 30 characters. \n")
	fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)

	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Print("> ")
		text, _ := reader.ReadString('\n')
		text = strings.Replace(text, "\n", "", -1)

		if len(text) == 0 {
			fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)
			continue
		}

		if err := projectCreate(text); err != nil {
			fmt.Printf("%sProject name could not be created %s\n", TERMREDREV, TERMCLEAR)
			fmt.Printf("%sReason: %s %s\n", TERMREDB, err, TERMCLEAR)
			fmt.Printf("%sPlease choose another. %s\n", TERMREDREV, TERMCLEAR)
			continue
		}

		fmt.Printf("Project Created\n")
		result = text
		break

	}
	sec1.Close()

	sec2 := NewSection("Activating Billing for the project")
	sec2.Open()
	account, err := BillingAccountManage()
	if err != nil {
		return "", fmt.Errorf("could not determine proper billing account: %s ", err)
	}

	if err := BillingAccountProjectAttach(result, account); err != nil {
		return "", fmt.Errorf("could not link billing account: %s ", err)
	}
	sec2.Close()
	return result, nil
}

// regions will return a list of regions depending on product type
func regions(project, product string) ([]string, error) {
	switch product {
	case "compute":
		return regionsCompute(project)
	case "functions":
		return regionsFunctions(project)
	case "run":
		return regionsRun(project)
	}

	return []string{}, fmt.Errorf("invalid product requested: %s", product)
}

// RegionManage promps a user to select a region.
func RegionManage(project, product, def string) (string, error) {
	fmt.Printf("Polling for regions...\n")
	regions, err := regions(project, product)
	if err != nil {
		return "", err
	}
	fmt.Printf("%sChoose a valid region to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	region := listSelect(toLabeledValueSlice(regions), def)

	return region.Value, nil
}

// ZoneManage promps a user to select a zone.
func ZoneManage(project, region string) (string, error) {
	fmt.Printf("Polling for zones...\n")
	zones, err := zones(project, region)
	if err != nil {
		return "", err
	}

	fmt.Printf("%sChoose a valid zone to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	zone := listSelect(toLabeledValueSlice(zones), zones[0])
	return zone.Value, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Damjan Cvetko. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package pcapgo

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"time"

	"bufio"
	"compress/gzip"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

// Reader wraps an underlying io.Reader to read packet data in PCAP
// format.  See http://wiki.wireshark.org/Development/LibpcapFileFormat
// for information on the file format.
//
// We currenty read v2.4 file format with nanosecond and microsecdond
// timestamp resolution in little-endian and big-endian encoding.
//
// If the PCAP data is gzip compressed it is transparently uncompressed
// by wrapping the given io.Reader with a gzip.Reader.
type Reader struct {
	r              io.Reader
	byteOrder      binary.ByteOrder
	nanoSecsFactor uint32
	versionMajor   uint16
	versionMinor   uint16
	// timezone
	// sigfigs
	snaplen  uint32
	linkType layers.LinkType
	// reusable buffer
	buf [16]byte
}

const magicNanoseconds = 0xA1B23C4D
const magicMicrosecondsBigendian = 0xD4C3B2A1
const magicNanosecondsBigendian = 0x4D3CB2A1

const magicGzip1 = 0x1f
const magicGzip2 = 0x8b

// NewReader returns a new reader object, for reading packet data from
// the given reader. The reader must be open and header data is
// read from it at this point.
// If the file format is not supported an error is returned
//
//  // Create new reader:
//  f, _ := os.Open("/tmp/file.pcap")
//  defer f.Close()
//  r, err := NewReader(f)
//  data, ci, err := r.ReadPacketData()
func NewReader(r io.Reader) (*Reader, error) {
	ret := Reader{r: r}
	if err := ret.readHeader(); err != nil {
		return nil, err
	}
	return &ret, nil
}

func (r *Reader) readHeader() error {
	br := bufio.NewReader(r.r)
	gzipMagic, err := br.Peek(2)
	if err != nil {
		return err
	}

	if gzipMagic[0] == magicGzip1 && gzipMagic[1] == magicGzip2 {
		if r.r, err = gzip.NewReader(br); err != nil {
			return err
		}
	} else {
		r.r = br
	}

	buf := make([]byte, 24)
	if n, err := io.ReadFull(r.r, buf); err != nil {
		return err
	} else if n < 24 {
		return errors.New("Not enough data for read")
	}
	if magic := binary.LittleEndian.Uint32(buf[0:4]); magic == magicNanoseconds {
		r.byteOrder = binary.LittleEndian
		r.nanoSecsFactor = 1
	} else if magic == magicNanosecondsBigendian {
		r.byteOrder = binary.BigEndian
		r.nanoSecsFactor = 1
	} else if magic == magicMicroseconds {
		r.byteOrder = binary.LittleEndian
		r.nanoSecsFactor = 1000
	} else if magic == magicMicrosecondsBigendian {
		r.byteOrder = binary.BigEndian
		r.nanoSecsFactor = 1000
	} else {
		return fmt.Errorf("Unknown magic %x", magic)
	}
	if r.versionMajor = r.byteOrder.Uint16(buf[4:6]); r.versionMajor != versionMajor {
		return fmt.Errorf("Unknown major version %d", r.versionMajor)
	}
	if r.versionMinor = r.byteOrder.Uint16(buf[6:8]); r.versionMinor != versionMinor {
		return fmt.Errorf("Unknown minor version %d", r.versionMinor)
	}
	// ignore timezone 8:12 and sigfigs 12:16
	r.snaplen = r.byteOrder.Uint32(buf[16:20])
	r.linkType = layers.LinkType(r.byteOrder.Uint32(buf[20:24]))
	return nil
}

// ReadPacketData reads next packet from file.
func (r *Reader) ReadPacketData() (data []byte, ci gopacket.CaptureInfo, err error) {
	if ci, err = r.readPacketHeader(); err != nil {
		return
	}
	if ci.CaptureLength > int(r.snaplen) {
		err = fmt.Errorf("capture length exceeds snap length: %d > %d", 16+ci.CaptureLength, r.snaplen)
		return
	}
	data = make([]byte, ci.CaptureLength)
	_, err = io.ReadFull(r.r, data)
	return data, ci, err
}

func (r *Reader) readPacketHeader() (ci gopacket.CaptureInfo, err error) {
	if _, err = io.ReadFull(r.r, r.buf[:]); err != nil {
		return
	}
	ci.Timestamp = time.Unix(int64(r.byteOrder.Uint32(r.buf[0:4])), int64(r.byteOrder.Uint32(r.buf[4:8])*r.nanoSecsFactor)).UTC()
	ci.CaptureLength = int(r.byteOrder.Uint32(r.buf[8:12]))
	ci.Length = int(r.byteOrder.Uint32(r.buf[12:16]))
	return
}

// LinkType returns network, as a layers.LinkType.
func (r *Reader) LinkType() layers.LinkType {
	return r.linkType
}

// Reader formater
func (r *Reader) String() string {
	return fmt.Sprintf("PcapFile  maj: %x min: %x snaplen: %d linktype: %s", r.versionMajor, r.versionMinor, r.snaplen, r.linkType)
}
<commit_msg>add Snaplen() method to Reader to access private snaplen value<commit_after>// Copyright 2014 Damjan Cvetko. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package pcapgo

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"time"

	"bufio"
	"compress/gzip"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

// Reader wraps an underlying io.Reader to read packet data in PCAP
// format.  See http://wiki.wireshark.org/Development/LibpcapFileFormat
// for information on the file format.
//
// We currenty read v2.4 file format with nanosecond and microsecdond
// timestamp resolution in little-endian and big-endian encoding.
//
// If the PCAP data is gzip compressed it is transparently uncompressed
// by wrapping the given io.Reader with a gzip.Reader.
type Reader struct {
	r              io.Reader
	byteOrder      binary.ByteOrder
	nanoSecsFactor uint32
	versionMajor   uint16
	versionMinor   uint16
	// timezone
	// sigfigs
	snaplen  uint32
	linkType layers.LinkType
	// reusable buffer
	buf [16]byte
}

const magicNanoseconds = 0xA1B23C4D
const magicMicrosecondsBigendian = 0xD4C3B2A1
const magicNanosecondsBigendian = 0x4D3CB2A1

const magicGzip1 = 0x1f
const magicGzip2 = 0x8b

// NewReader returns a new reader object, for reading packet data from
// the given reader. The reader must be open and header data is
// read from it at this point.
// If the file format is not supported an error is returned
//
//  // Create new reader:
//  f, _ := os.Open("/tmp/file.pcap")
//  defer f.Close()
//  r, err := NewReader(f)
//  data, ci, err := r.ReadPacketData()
func NewReader(r io.Reader) (*Reader, error) {
	ret := Reader{r: r}
	if err := ret.readHeader(); err != nil {
		return nil, err
	}
	return &ret, nil
}

func (r *Reader) readHeader() error {
	br := bufio.NewReader(r.r)
	gzipMagic, err := br.Peek(2)
	if err != nil {
		return err
	}

	if gzipMagic[0] == magicGzip1 && gzipMagic[1] == magicGzip2 {
		if r.r, err = gzip.NewReader(br); err != nil {
			return err
		}
	} else {
		r.r = br
	}

	buf := make([]byte, 24)
	if n, err := io.ReadFull(r.r, buf); err != nil {
		return err
	} else if n < 24 {
		return errors.New("Not enough data for read")
	}
	if magic := binary.LittleEndian.Uint32(buf[0:4]); magic == magicNanoseconds {
		r.byteOrder = binary.LittleEndian
		r.nanoSecsFactor = 1
	} else if magic == magicNanosecondsBigendian {
		r.byteOrder = binary.BigEndian
		r.nanoSecsFactor = 1
	} else if magic == magicMicroseconds {
		r.byteOrder = binary.LittleEndian
		r.nanoSecsFactor = 1000
	} else if magic == magicMicrosecondsBigendian {
		r.byteOrder = binary.BigEndian
		r.nanoSecsFactor = 1000
	} else {
		return fmt.Errorf("Unknown magic %x", magic)
	}
	if r.versionMajor = r.byteOrder.Uint16(buf[4:6]); r.versionMajor != versionMajor {
		return fmt.Errorf("Unknown major version %d", r.versionMajor)
	}
	if r.versionMinor = r.byteOrder.Uint16(buf[6:8]); r.versionMinor != versionMinor {
		return fmt.Errorf("Unknown minor version %d", r.versionMinor)
	}
	// ignore timezone 8:12 and sigfigs 12:16
	r.snaplen = r.byteOrder.Uint32(buf[16:20])
	r.linkType = layers.LinkType(r.byteOrder.Uint32(buf[20:24]))
	return nil
}

// ReadPacketData reads next packet from file.
func (r *Reader) ReadPacketData() (data []byte, ci gopacket.CaptureInfo, err error) {
	if ci, err = r.readPacketHeader(); err != nil {
		return
	}
	if ci.CaptureLength > int(r.snaplen) {
		err = fmt.Errorf("capture length exceeds snap length: %d > %d", 16+ci.CaptureLength, r.snaplen)
		return
	}
	data = make([]byte, ci.CaptureLength)
	_, err = io.ReadFull(r.r, data)
	return data, ci, err
}

func (r *Reader) readPacketHeader() (ci gopacket.CaptureInfo, err error) {
	if _, err = io.ReadFull(r.r, r.buf[:]); err != nil {
		return
	}
	ci.Timestamp = time.Unix(int64(r.byteOrder.Uint32(r.buf[0:4])), int64(r.byteOrder.Uint32(r.buf[4:8])*r.nanoSecsFactor)).UTC()
	ci.CaptureLength = int(r.byteOrder.Uint32(r.buf[8:12]))
	ci.Length = int(r.byteOrder.Uint32(r.buf[12:16]))
	return
}

// LinkType returns network, as a layers.LinkType.
func (r *Reader) LinkType() layers.LinkType {
	return r.linkType
}

// Snaplen returns the snapshot length of the capture file.
func (r *Reader) Snaplen() uint32 {
	return r.snaplen
}

// Reader formater
func (r *Reader) String() string {
	return fmt.Sprintf("PcapFile  maj: %x min: %x snaplen: %d linktype: %s", r.versionMajor, r.versionMinor, r.snaplen, r.linkType)
}
<|endoftext|>


<commit_before>package moskus

// locale: en_US

// NOTE: Work in progress!

// Anything that's specific to the US, with the exception of some movable dates which are in movable.go
// This calendar has a corresponding locale code in the NewCalendar function in calendar.go

import (
	"time"
)

type USCalendar struct{}

// Create a new US calendar
func NewUSCalendar() USCalendar {
	return USCalendar{}
}

// Finds the US name for a day of the week.
// Note that time.Weekday starts at 0 with Sunday, not Monday.
func (nc USCalendar) DayName(day time.Weekday) string {
	return day.String()
}

// Finds the US name for a given month
func (nc USCalendar) MonthName(month time.Month) string {
	return month.String()
}

// Checks if a given date is a "red day" (public holiday) in the US calendar.
// Returns true/false, a description and true/false for if it's a flag day.
// The dates will never overlap.
func (nc USCalendar) RedDay(date time.Time) (bool, string, bool) {

	// Source: http://en.wikipedia.org/wiki/Public_holidays_in_the_United_States
	// Source: http://timpanogos.wordpress.com/flag-fly-dates/

	var (
		desc string
		flag bool
	)

	// Sundays
	if date.Weekday() == 0 {
		desc = "Sunday"
	}

	// New Year's Day
	if atMD(date, 1, 1) {
		desc = "New Year's Day"
		flag = true
	}

	// Birthday of Dr. Martin Luther King, Jr.
	if atDrMartinLutherKingJrBirthday(date) {
		desc = "Birthday of Dr. Martin Luther King, Jr."
		flag = true
	}

	// TODO: The rest

	// Red days
	if desc != "" {
		// Then return
		return true, desc, flag
	}

	// Normal days
	return false, desc, false
}

// Some days are not red, but special in one way or another.
// Checks if a given date is notable. Returns true/false if the
// given date is notable, a comma separated description (in case there are more
// than one notable event that day) and true/false depending on if it's a flag
// flying day or not.
func (nc USCalendar) NotableDay(date time.Time) (bool, string, bool) {

	// Source: http://www.timeanddate.no/kalender/merkedag-innhold
	// Source: http://no.wikipedia.org/wiki/Norges_offisielle_flaggdager

	var (
		//descriptions []string
		//flag         bool
	)

	// Since days may overlap, "flaggdager" must come first for the flag
	// flying days to be correct.

	// --- Flag days ---

	// --- Non-flag days ---

	// No notable events
	return false, "", false
}

// Checks if a given date is in a notable time range (summer holidays, for instance)
func (nc USCalendar) NotablePeriod(date time.Time) (bool, string) {
	// TODO:
	// uke 28, 29 og 30, fellesferie
	// jul
	// fastelavn
	// faste
	// sommer/vinter/vår/høst (legg denne først i listen)
	return false, ""
}

// An ordinary day
func (nc USCalendar) NormalDay() string {
	return "Ordinary"
}
<commit_msg>Work in progress<commit_after>package moskus

// locale: en_US

// NOTE: Work in progress!

// Anything that's specific to the US, with the exception of some movable dates which are in movable.go
// This calendar has a corresponding locale code in the NewCalendar function in calendar.go

import (
	"time"
)

type USCalendar struct{}

// Create a new US calendar
func NewUSCalendar() USCalendar {
	return USCalendar{}
}

// Finds the US name for a day of the week.
// Note that time.Weekday starts at 0 with Sunday, not Monday.
func (nc USCalendar) DayName(day time.Weekday) string {
	return day.String()
}

// Finds the US name for a given month
func (nc USCalendar) MonthName(month time.Month) string {
	return month.String()
}

// Checks if a given date is a "red day" (public holiday) in the US calendar.
// Returns true/false, a description and true/false for if it's a flag day.
// The dates will never overlap.
func (nc USCalendar) RedDay(date time.Time) (bool, string, bool) {

	// Source: http://en.wikipedia.org/wiki/Public_holidays_in_the_United_States
	// Source: http://timpanogos.wordpress.com/flag-fly-dates/

	var (
		desc string
		flag bool
	)

	// Sundays
	if date.Weekday() == 0 {
		desc = "Sunday"
	}

	// New Year's Day
	if atMD(date, 1, 1) {
		desc = "New Year's Day"
		flag = true
	}

	// Birthday of Dr. Martin Luther King, Jr.
	if atDrMartinLutherKingJrBirthday(date) {
		desc = "Birthday of Dr. Martin Luther King, Jr."
		flag = true
	}

	// TODO: The rest

	// Elections: 2000, 2004, 2008, 2012 etc

	/*
	First January 20 following a Presidential electionInauguration DayObserved only by federal government employees in Washington, D.C., and the border counties of Maryland and Virginia to relieve congestion that occurs with this major event. Swearing-in of President of the United States and Vice President of the United States. Celebrated every fourth year. Note: Takes place on January 21 if the 20th is a Sunday (although the President is still privately inaugurated on the 20th). If Inauguration Day falls on a Saturday, the preceding Friday is not a federal holiday.

	Third Monday in FebruaryWashington's Birthday/Presidents' DayWashington's Birthday was first declared a federal holiday by an 1879 act of Congress. The Uniform Holidays Act, 1968, shifted the date of the commemoration of Washington's Birthday from February 22 to the third Monday in February (between February 15 and 21, meaning the observed holiday never falls on Washington's actual birthday). Because of this, combined with the fact that President Abraham Lincoln's birthday falls on February 12, many people now refer to this holiday as "Presidents' Day" and consider it a day honoring all American presidents. However, neither the Uniform Holidays Act nor any subsequent law changed the name of the holiday from Washington's Birthday to Presidents' Day.[8]

	Last Monday in MayMemorial DayHonors the nation's war dead from the Civil War onwards; marks the unofficial beginning of the summer season. (traditionally May 30, shifted by the Uniform Holidays Act 1968)

	July 4Independence DayHonorsCelebrates the signing of the Declaration of Independence from British rule, also called the Fourth of July. Firework celebrations are held in many cities throughout the nation.

	First Monday in SeptemberLabor DayHonorsCelebratesCelebrates the achievements of workers and the labor movement; marks the unofficial end of the summer season.

	Second Monday in OctoberColumbus DayHonorsCelebratesCelebratesHonors Christopher Columbus, traditional discoverer of the Americas. In some areas it is also a celebration of Italian culture and heritage. (traditionally October 12)

	November 11Veterans DayHonorsCelebratesCelebratesHonorsHonors all veterans of the United States armed forces. It is observed on November 11 to recall the end of World War I on that date in 1918 (major hostilities of World War I were formally ended at the 11th hour of the 11th day of the 11th month of 1918 with the German signing of the Armistice).

	Fourth Thursday in NovemberThanksgiving DayHonorsCelebratesCelebratesHonorsHonorsTraditionally celebrates the giving of thanks for the autumn harvest. Traditionally includes the sharing of a turkey dinner. Traditional start of the Christmas and holiday season.

	December 25ChristmasThe most widely celebrated holiday of the Christian year, Christmas is observed as a commemoration of the birth of Jesus of Nazareth.
	*/

	// Red days
	if desc != "" {
		// Then return
		return true, desc, flag
	}

	// Normal days
	return false, desc, false
}

// Some days are not red, but special in one way or another.
// Checks if a given date is notable. Returns true/false if the
// given date is notable, a comma separated description (in case there are more
// than one notable event that day) and true/false depending on if it's a flag
// flying day or not.
func (nc USCalendar) NotableDay(date time.Time) (bool, string, bool) {

	// Source: http://www.timeanddate.no/kalender/merkedag-innhold
	// Source: http://no.wikipedia.org/wiki/Norges_offisielle_flaggdager

	var (
		//descriptions []string
		//flag         bool
	)

	// Since days may overlap, "flaggdager" must come first for the flag
	// flying days to be correct.

	// --- Flag days ---

	// --- Non-flag days ---

	// No notable events
	return false, "", false
}

// Checks if a given date is in a notable time range (summer holidays, for instance)
func (nc USCalendar) NotablePeriod(date time.Time) (bool, string) {
	// TODO:
	// uke 28, 29 og 30, fellesferie
	// jul
	// fastelavn
	// faste
	// sommer/vinter/vår/høst (legg denne først i listen)
	return false, ""
}

// An ordinary day
func (nc USCalendar) NormalDay() string {
	return "Ordinary"
}
<|endoftext|>