<commit_before>package main

import (
	"errors"
	"io"
	"os"
	"path/filepath"
	"time"

	"crypto"
	"crypto/rand"
	"crypto/rsa"
	_ "crypto/sha256"
	_ "golang.org/x/crypto/ripemd160"

	"compress/gzip"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
	// Goencrypt app
	app  = kingpin.New("goencrypt", "A command line tool for encrypting files")
	bits = app.Flag("bits", "Bits for keys").Default("4096").Int()
	privateKey = app.Flag("private", "Private key").String()
	publicKey = app.Flag("public", "Public key").String()
	signatureFile = app.Flag("sig", "Signature File").String()

	// Generates new public and private keys
	keyGenCmd       = app.Command("keygen", "Generates a new public/private key pair")
	keyOutputPrefix = keyGenCmd.Arg("prefix", "Prefix of key files").Required().String()
	keyOutputDir    = keyGenCmd.Flag("d", "Output directory of key files").Default(".").String()

	// Base command for entity sub-commands
	entityCmd        = app.Command("entity", "Generates a PGP entity")
	entityOutputFile = entityCmd.Flag("o", "Output file of entity").Default("./entity.asc").String()
	identityName     = entityCmd.Flag("name", "Name of identity").Default("").String()
	identityComment  = entityCmd.Flag("comment", "Command for identity").Default("").String()
	identityEmail    = entityCmd.Flag("email", "Email for identity").Default("").String()

	// Creates a new entity
	createEntityCmd = entityCmd.Command("new", "Create new entity")

	// Adds an identity to an entity
	addIdentityCmd            = entityCmd.Command("add", "Add an identity to the entity")
	addIdentityFile           = addIdentityCmd.Arg("entity file", "Entity to add the identity to").Required().String()
	addIdentityPrivateKeyFile = addIdentityCmd.Flag("key", "Private key of entity").String()

	// Encrypts a file with a public key
	encryptionCmd = app.Command("encrypt", "Encrypt a file")

	// Signs a file with a private key
	signCmd = app.Command("sign", "Sign a file")

	// Verifies a file was signed with the public key
	verifyCmd = app.Command("verify", "Verify a signature")

	// Decrypts a file with a private key
	decryptionCmd = app.Command("decrypt", "Decrypt a file")
)

func main() {
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	// generate keys
	case keyGenCmd.FullCommand():
		generateKeys()
	case createEntityCmd.FullCommand():
		newEntity()
	case encryptionCmd.FullCommand():
		encryptFile()
	case signCmd.FullCommand():
		signFile()
	case verifyCmd.FullCommand():
		verifyFile()
	case decryptionCmd.FullCommand():
		decryptFile()
	case addIdentityCmd.FullCommand():
		privKey := decodePrivateKey(*addIdentityPrivateKeyFile)

		entityList := decodeKeyRing(*addIdentityFile)
		var config packet.Config
		for _, ent := range entityList {
			ent.PrivateKey = privKey
			addIdentity(false, config, ent)

			// Sign identities
			for _, id := range ent.Identities {
				kingpin.Errorf("Identity: %#v", *ent.PrivateKey)
				err := id.SelfSignature.SignUserId(id.UserId.Id, ent.PrimaryKey, ent.PrivateKey, nil)
				if err != nil {
					kingpin.FatalIfError(err, "Error signing identity: %s", err)
					return
				}
			}

			// Sign subkeys
			for _, subkey := range ent.Subkeys {
				err := subkey.Sig.SignKey(subkey.PublicKey, ent.PrivateKey, nil)
				if err != nil {
					kingpin.FatalIfError(err, "Error signing subkey: %s", err)
				}
			}

			entityFile, err := os.Create(*entityOutputFile)
			kingpin.FatalIfError(err, "Error creating entity file: %s", err)
			defer entityFile.Close()

			w, err := armor.Encode(entityFile, openpgp.PublicKeyType, map[string]string{})
			err = ent.Serialize(w)
			kingpin.Errorf("identities: %d", len(ent.Identities))
			kingpin.Errorf("subkeys: %d", len(ent.Subkeys))
			kingpin.FatalIfError(err, "Error serializing entity: %s", err)
			w.Close()

			break
		}

	default:
		kingpin.FatalUsage("Unknown command")
	}
}

func encodePrivateKey(out io.Writer, key *rsa.PrivateKey) {
	w, err := armor.Encode(out, openpgp.PrivateKeyType, make(map[string]string))
	kingpin.FatalIfError(err, "Error creating OpenPGP Armor: %s", err)

	pgpKey := packet.NewRSAPrivateKey(time.Now(), key)
	kingpin.FatalIfError(pgpKey.Serialize(w), "Error serializing private key: %s", err)
	kingpin.FatalIfError(w.Close(), "Error serializing private key: %s", err)
}

func decodePrivateKey(filename string) *packet.PrivateKey {

	// open ascii armored private key
	in, err := os.Open(filename)
	kingpin.FatalIfError(err, "Error opening private key: %s", err)
	defer in.Close()

	block, err := armor.Decode(in)
	kingpin.FatalIfError(err, "Error decoding OpenPGP Armor: %s", err)

	if block.Type != openpgp.PrivateKeyType {
		kingpin.FatalIfError(errors.New("Invalid private key file"), "Error decoding private key")
	}

	reader := packet.NewReader(block.Body)
	pkt, err := reader.Next()
	kingpin.FatalIfError(err, "Error reading private key")

	key, ok := pkt.(*packet.PrivateKey)
	if !ok {
		kingpin.FatalIfError(errors.New("Invalid private key"), "Error parsing private key")
	}
	return key
}

func encodePublicKey(out io.Writer, key *rsa.PrivateKey) {
	w, err := armor.Encode(out, openpgp.PublicKeyType, make(map[string]string))
	kingpin.FatalIfError(err, "Error creating OpenPGP Armor: %s", err)

	pgpKey := packet.NewRSAPublicKey(time.Now(), &key.PublicKey)
	kingpin.FatalIfError(pgpKey.Serialize(w), "Error serializing public key: %s", err)
	kingpin.FatalIfError(w.Close(), "Error serializing public key: %s", err)
}

func decodePublicKey(filename string) *packet.PublicKey {

	// open ascii armored public key
	in, err := os.Open(filename)
	kingpin.FatalIfError(err, "Error opening public key: %s", err)
	defer in.Close()

	block, err := armor.Decode(in)
	kingpin.FatalIfError(err, "Error decoding OpenPGP Armor: %s", err)

	if block.Type != openpgp.PublicKeyType {
		kingpin.FatalIfError(errors.New("Invalid private key file"), "Error decoding private key")
	}

	reader := packet.NewReader(block.Body)
	pkt, err := reader.Next()
	kingpin.FatalIfError(err, "Error reading private key")

	key, ok := pkt.(*packet.PublicKey)
	if !ok {
		kingpin.FatalIfError(errors.New("Invalid public key"), "Error parsing public key")
	}
	return key
}

func decodeSignature(filename string) *packet.Signature {

	// open ascii armored public key
	in, err := os.Open(filename)
	kingpin.FatalIfError(err, "Error opening public key: %s", err)
	defer in.Close()

	block, err := armor.Decode(in)
	kingpin.FatalIfError(err, "Error decoding OpenPGP Armor: %s", err)

	if block.Type != openpgp.SignatureType {
		kingpin.FatalIfError(errors.New("Invalid signature file"), "Error decoding signature")
	}

	reader := packet.NewReader(block.Body)
	pkt, err := reader.Next()
	kingpin.FatalIfError(err, "Error reading signature")

	sig, ok := pkt.(*packet.Signature)
	if !ok {
		kingpin.FatalIfError(errors.New("Invalid signature"), "Error parsing signature")
	}
	return sig
}

func encryptFile() {
	pubKey := decodePublicKey(*publicKey)
	privKey := decodePrivateKey(*privateKey)

	to := createEntityFromKeys(pubKey, privKey)

	w, err := armor.Encode(os.Stdout, "Message", make(map[string]string))
	kingpin.FatalIfError(err, "Error creating OpenPGP Armor: %s", err)
	defer w.Close()

	plain, err := openpgp.Encrypt(w, []*openpgp.Entity{to}, nil, nil, nil)
	kingpin.FatalIfError(err, "Error creating entity for encryption")
	defer plain.Close()

	compressed, err := gzip.NewWriterLevel(plain, gzip.BestCompression)
	kingpin.FatalIfError(err, "Invalid compression level")

	n, err := io.Copy(compressed, os.Stdin)
	kingpin.FatalIfError(err, "Error writing encrypted file")
	kingpin.Errorf("Encrypted %d bytes", n)

	compressed.Close()
}

func decryptFile() {
	pubKey := decodePublicKey(*publicKey)
	privKey := decodePrivateKey(*privateKey)

	entity := createEntityFromKeys(pubKey, privKey)

	block, err := armor.Decode(os.Stdin)
	kingpin.FatalIfError(err, "Error reading OpenPGP Armor: %s", err)

	if block.Type != "Message" {
		kingpin.FatalIfError(err, "Invalid message type")
	}

	var entityList openpgp.EntityList
	entityList = append(entityList, entity)

	md, err := openpgp.ReadMessage(block.Body, entityList, nil, nil)
	kingpin.FatalIfError(err, "Error reading message")

	compressed, err := gzip.NewReader(md.UnverifiedBody)
	kingpin.FatalIfError(err, "Invalid compression level")
	defer compressed.Close()

	n, err := io.Copy(os.Stdout, compressed)
	kingpin.FatalIfError(err, "Error reading encrypted file")
	kingpin.Errorf("Decrypted %d bytes", n)
}

func signFile() {
	pubKey := decodePublicKey(*publicKey)
	privKey := decodePrivateKey(*privateKey)

	signer := createEntityFromKeys(pubKey, privKey)

	err := openpgp.ArmoredDetachSign(os.Stdout, signer, os.Stdin, nil)
	kingpin.FatalIfError(err, "Error signing input")
}

func verifyFile() {
	pubKey := decodePublicKey(*publicKey)
	sig := decodeSignature(*signatureFile)

	hash := sig.Hash.New()
	io.Copy(hash, os.Stdin)

	err := pubKey.VerifySignature(hash, sig)
	kingpin.FatalIfError(err, "Error signing input")
	kingpin.Errorf("Verified signature")
}

func createEntityFromKeys(pubKey *packet.PublicKey, privKey *packet.PrivateKey) *openpgp.Entity {
	config := packet.Config{
		DefaultHash:            crypto.SHA256,
		DefaultCipher:          packet.CipherAES256,
		DefaultCompressionAlgo: packet.CompressionZLIB,
		CompressionConfig: &packet.CompressionConfig{
			Level: 9,
		},
		RSABits: *bits,
	}
	currentTime := config.Now()
	uid := packet.NewUserId("", "", "")

	e := openpgp.Entity{
		PrimaryKey: pubKey,
		PrivateKey: privKey,
		Identities: make(map[string]*openpgp.Identity),
	}
	isPrimaryId := false

	e.Identities[uid.Id] = &openpgp.Identity{
		Name:   uid.Name,
		UserId: uid,
		SelfSignature: &packet.Signature{
			CreationTime: currentTime,
			SigType:      packet.SigTypePositiveCert,
			PubKeyAlgo:   packet.PubKeyAlgoRSA,
			Hash:         config.Hash(),
			IsPrimaryId:  &isPrimaryId,
			FlagsValid:   true,
			FlagSign:     true,
			FlagCertify:  true,
			IssuerKeyId:  &e.PrimaryKey.KeyId,
		},
	}

	keyLifetimeSecs := uint32(86400 * 365)

	e.Subkeys = make([]openpgp.Subkey, 1)
	e.Subkeys[0] = openpgp.Subkey{
		PublicKey: pubKey,
		PrivateKey: privKey,
		Sig: &packet.Signature{
			CreationTime:              currentTime,
			SigType:                   packet.SigTypeSubkeyBinding,
			PubKeyAlgo:                packet.PubKeyAlgoRSA,
			Hash:                      config.Hash(),
			PreferredHash:             []uint8{8}, // SHA-256
			FlagsValid:                true,
			FlagEncryptStorage:        true,
			FlagEncryptCommunications: true,
			IssuerKeyId:               &e.PrimaryKey.KeyId,
			KeyLifetimeSecs:           &keyLifetimeSecs,
		},
	}
	return &e
}

func newEntity() {
	ent, err := openpgp.NewEntity(*identityName, *identityComment, *identityEmail, &packet.Config{
		DefaultCipher:          packet.CipherAES256,
		DefaultCompressionAlgo: packet.CompressionZLIB,
		CompressionConfig: &packet.CompressionConfig{
			Level: 9,
		},
		RSABits: *bits,
	})
	kingpin.FatalIfError(err, "Error creating entity: %s", err)

	// Sign identities
	for _, id := range ent.Identities {
		err := id.SelfSignature.SignUserId(id.UserId.Id, ent.PrimaryKey, ent.PrivateKey, nil)
		if err != nil {
			kingpin.FatalIfError(err, "Error signing identity: %s", err)
			return
		}
	}

	// Sign subkeys
	for _, subkey := range ent.Subkeys {

		priv, err := os.Create(filepath.Join(*entityOutputFile + ".subkey.key"))
		kingpin.FatalIfError(err, "Error writing private key to file: %s", err)
		defer priv.Close()

		encodePrivateKey(priv, subkey.PrivateKey.PrivateKey.(*rsa.PrivateKey))

		err = subkey.Sig.SignKey(subkey.PublicKey, ent.PrivateKey, nil)
		if err != nil {
			kingpin.FatalIfError(err, "Error signing subkey: %s", err)
		}
	}

	entityFile, err := os.Create(*entityOutputFile)
	kingpin.FatalIfError(err, "Error creating entity file: %s", err)
	defer entityFile.Close()

	w, err := armor.Encode(entityFile, openpgp.PublicKeyType, map[string]string{})
	err = ent.Serialize(w)
	kingpin.Errorf("identities: %d", len(ent.Identities))
	kingpin.Errorf("subkeys: %d", len(ent.Subkeys))
	kingpin.FatalIfError(err, "Error serializing entity: %s", err)
	w.Close()

	priv, err := os.Create(filepath.Join(*entityOutputFile + ".key"))
	kingpin.FatalIfError(err, "Error writing private key to file: %s", err)
	defer priv.Close()

	encodePrivateKey(priv, ent.PrivateKey.PrivateKey.(*rsa.PrivateKey))
}

// func createNewEntity() {
// 	key, err := rsa.GenerateKey(rand.Reader, *bits)
// 	kingpin.FatalIfError(err, "Error generating RSA key: %s", err)

// 	currentTime := time.Now()
// 	e := &openpgp.Entity{
// 		PrimaryKey: packet.NewRSAPublicKey(currentTime, &key.PublicKey),
// 		PrivateKey: packet.NewRSAPrivateKey(currentTime, key),
// 		Identities: make(map[string]*openpgp.Identity),
// 	}

// 	var config packet.Config
// 	addIdentity(true, config, e)

// 	entityFile, err := os.Create(*entityOutputFile)
// 	kingpin.FatalIfError(err, "Error creating entity file: %s", err)
// 	defer entityFile.Close()

// 	w, err := armor.Encode(entityFile, openpgp.PublicKeyType, map[string]string{})
// 	e.Serialize(w)
// 	w.Close()

// 	priv, err := os.Create(filepath.Join(*entityOutputFile + ".key"))
// 	kingpin.FatalIfError(err, "Error writing private key to file: %s", err)
// 	defer priv.Close()

// 	encodePrivateKey(priv, key)
// }

func addIdentity(isPrimary bool, config packet.Config, entity *openpgp.Entity) {
	currentTime := config.Now()
	uid := packet.NewUserId(*identityName, *identityComment, *identityEmail)

	encryptingPriv, err := rsa.GenerateKey(rand.Reader, *bits)
	kingpin.FatalIfError(err, "Error generating RSA key: %s", err)

	entity.Identities[uid.Id] = &openpgp.Identity{
		Name:   uid.Name,
		UserId: uid,
		SelfSignature: &packet.Signature{
			CreationTime: currentTime,
			SigType:      packet.SigTypePositiveCert,
			PubKeyAlgo:   packet.PubKeyAlgoRSA,
			Hash:         config.Hash(),
			IsPrimaryId:  &isPrimary,
			FlagsValid:   true,
			FlagSign:     true,
			FlagCertify:  true,
			IssuerKeyId:  &entity.PrimaryKey.KeyId,
		},
	}

	keyLifetimeSecs := uint32(86400 * 365 * 10)
	subkey := openpgp.Subkey{
		PublicKey:  packet.NewRSAPublicKey(currentTime, &encryptingPriv.PublicKey),
		PrivateKey: packet.NewRSAPrivateKey(currentTime, encryptingPriv),

		Sig: &packet.Signature{
			CreationTime:              currentTime,
			SigType:                   packet.SigTypeSubkeyBinding,
			PubKeyAlgo:                packet.PubKeyAlgoRSA,
			Hash:                      config.Hash(),
			FlagsValid:                true,
			FlagEncryptStorage:        true,
			FlagEncryptCommunications: true,
			IssuerKeyId:               &entity.PrimaryKey.KeyId,
			KeyLifetimeSecs:           &keyLifetimeSecs,
		},
	}
	subkey.PublicKey.IsSubkey = true
	subkey.PrivateKey.IsSubkey = true
	entity.Subkeys = append(entity.Subkeys, subkey)
}

func decodeKeyRing(filename string) openpgp.EntityList {

	// open ascii armored private key
	entityFile, err := os.Open(filename)
	kingpin.FatalIfError(err, "Error opening entity: %s", err)
	defer entityFile.Close()

	// read key ring
	entity, err := openpgp.ReadArmoredKeyRing(entityFile)
	kingpin.FatalIfError(err, "Error reading entity: %s", err)

	for _, e := range entity {
		kingpin.Errorf("Entity: %#v", *e)
	}
	return entity
}

func generateKeys() {
	key, err := rsa.GenerateKey(rand.Reader, *bits)
	kingpin.FatalIfError(err, "Error generating RSA key: %s", err)

	priv, err := os.Create(filepath.Join(*keyOutputDir, *keyOutputPrefix+".privkey"))
	kingpin.FatalIfError(err, "Error writing private key to file: %s", err)
	defer priv.Close()

	pub, err := os.Create(filepath.Join(*keyOutputDir, *keyOutputPrefix+".pubkey"))
	kingpin.FatalIfError(err, "Error writing public key to file: %s", err)
	defer pub.Close()

	encodePrivateKey(priv, key)
	encodePublicKey(pub, key)
}
<commit_msg>Removed entity management commands<commit_after>package main

import (
	"errors"
	"io"
	"os"
	"path/filepath"
	"time"

	"crypto"
	"crypto/rand"
	"crypto/rsa"
	_ "crypto/sha256"
	_ "golang.org/x/crypto/ripemd160"

	"compress/gzip"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
	// Goencrypt app
	app  = kingpin.New("goencrypt", "A command line tool for encrypting files")
	bits = app.Flag("bits", "Bits for keys").Default("4096").Int()
	privateKey = app.Flag("private", "Private key").String()
	publicKey = app.Flag("public", "Public key").String()
	signatureFile = app.Flag("sig", "Signature File").String()

	// Generates new public and private keys
	keyGenCmd       = app.Command("keygen", "Generates a new public/private key pair")
	keyOutputPrefix = keyGenCmd.Arg("prefix", "Prefix of key files").Required().String()
	keyOutputDir    = keyGenCmd.Flag("d", "Output directory of key files").Default(".").String()

	// Encrypts a file with a public key
	encryptionCmd = app.Command("encrypt", "Encrypt from stdin")

	// Signs a file with a private key
	signCmd = app.Command("sign", "Sign stdin")

	// Verifies a file was signed with the public key
	verifyCmd = app.Command("verify", "Verify a signature of stdin")

	// Decrypts a file with a private key
	decryptionCmd = app.Command("decrypt", "Decrypt from stdin")
)

func main() {
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	// generate keys
	case keyGenCmd.FullCommand():
		generateKeys()
	// case createEntityCmd.FullCommand():
	// 	newEntity()
	case encryptionCmd.FullCommand():
		encryptFile()
	case signCmd.FullCommand():
		signFile()
	case verifyCmd.FullCommand():
		verifyFile()
	case decryptionCmd.FullCommand():
		decryptFile()
	default:
		kingpin.FatalUsage("Unknown command")
	}
}

func encodePrivateKey(out io.Writer, key *rsa.PrivateKey) {
	w, err := armor.Encode(out, openpgp.PrivateKeyType, make(map[string]string))
	kingpin.FatalIfError(err, "Error creating OpenPGP Armor: %s", err)

	pgpKey := packet.NewRSAPrivateKey(time.Now(), key)
	kingpin.FatalIfError(pgpKey.Serialize(w), "Error serializing private key: %s", err)
	kingpin.FatalIfError(w.Close(), "Error serializing private key: %s", err)
}

func decodePrivateKey(filename string) *packet.PrivateKey {

	// open ascii armored private key
	in, err := os.Open(filename)
	kingpin.FatalIfError(err, "Error opening private key: %s", err)
	defer in.Close()

	block, err := armor.Decode(in)
	kingpin.FatalIfError(err, "Error decoding OpenPGP Armor: %s", err)

	if block.Type != openpgp.PrivateKeyType {
		kingpin.FatalIfError(errors.New("Invalid private key file"), "Error decoding private key")
	}

	reader := packet.NewReader(block.Body)
	pkt, err := reader.Next()
	kingpin.FatalIfError(err, "Error reading private key")

	key, ok := pkt.(*packet.PrivateKey)
	if !ok {
		kingpin.FatalIfError(errors.New("Invalid private key"), "Error parsing private key")
	}
	return key
}

func encodePublicKey(out io.Writer, key *rsa.PrivateKey) {
	w, err := armor.Encode(out, openpgp.PublicKeyType, make(map[string]string))
	kingpin.FatalIfError(err, "Error creating OpenPGP Armor: %s", err)

	pgpKey := packet.NewRSAPublicKey(time.Now(), &key.PublicKey)
	kingpin.FatalIfError(pgpKey.Serialize(w), "Error serializing public key: %s", err)
	kingpin.FatalIfError(w.Close(), "Error serializing public key: %s", err)
}

func decodePublicKey(filename string) *packet.PublicKey {

	// open ascii armored public key
	in, err := os.Open(filename)
	kingpin.FatalIfError(err, "Error opening public key: %s", err)
	defer in.Close()

	block, err := armor.Decode(in)
	kingpin.FatalIfError(err, "Error decoding OpenPGP Armor: %s", err)

	if block.Type != openpgp.PublicKeyType {
		kingpin.FatalIfError(errors.New("Invalid private key file"), "Error decoding private key")
	}

	reader := packet.NewReader(block.Body)
	pkt, err := reader.Next()
	kingpin.FatalIfError(err, "Error reading private key")

	key, ok := pkt.(*packet.PublicKey)
	if !ok {
		kingpin.FatalIfError(errors.New("Invalid public key"), "Error parsing public key")
	}
	return key
}

func decodeSignature(filename string) *packet.Signature {

	// open ascii armored public key
	in, err := os.Open(filename)
	kingpin.FatalIfError(err, "Error opening public key: %s", err)
	defer in.Close()

	block, err := armor.Decode(in)
	kingpin.FatalIfError(err, "Error decoding OpenPGP Armor: %s", err)

	if block.Type != openpgp.SignatureType {
		kingpin.FatalIfError(errors.New("Invalid signature file"), "Error decoding signature")
	}

	reader := packet.NewReader(block.Body)
	pkt, err := reader.Next()
	kingpin.FatalIfError(err, "Error reading signature")

	sig, ok := pkt.(*packet.Signature)
	if !ok {
		kingpin.FatalIfError(errors.New("Invalid signature"), "Error parsing signature")
	}
	return sig
}

func encryptFile() {
	pubKey := decodePublicKey(*publicKey)
	privKey := decodePrivateKey(*privateKey)

	to := createEntityFromKeys(pubKey, privKey)

	w, err := armor.Encode(os.Stdout, "Message", make(map[string]string))
	kingpin.FatalIfError(err, "Error creating OpenPGP Armor: %s", err)
	defer w.Close()

	plain, err := openpgp.Encrypt(w, []*openpgp.Entity{to}, nil, nil, nil)
	kingpin.FatalIfError(err, "Error creating entity for encryption")
	defer plain.Close()

	compressed, err := gzip.NewWriterLevel(plain, gzip.BestCompression)
	kingpin.FatalIfError(err, "Invalid compression level")

	n, err := io.Copy(compressed, os.Stdin)
	kingpin.FatalIfError(err, "Error writing encrypted file")
	kingpin.Errorf("Encrypted %d bytes", n)

	compressed.Close()
}

func decryptFile() {
	pubKey := decodePublicKey(*publicKey)
	privKey := decodePrivateKey(*privateKey)

	entity := createEntityFromKeys(pubKey, privKey)

	block, err := armor.Decode(os.Stdin)
	kingpin.FatalIfError(err, "Error reading OpenPGP Armor: %s", err)

	if block.Type != "Message" {
		kingpin.FatalIfError(err, "Invalid message type")
	}

	var entityList openpgp.EntityList
	entityList = append(entityList, entity)

	md, err := openpgp.ReadMessage(block.Body, entityList, nil, nil)
	kingpin.FatalIfError(err, "Error reading message")

	compressed, err := gzip.NewReader(md.UnverifiedBody)
	kingpin.FatalIfError(err, "Invalid compression level")
	defer compressed.Close()

	n, err := io.Copy(os.Stdout, compressed)
	kingpin.FatalIfError(err, "Error reading encrypted file")
	kingpin.Errorf("Decrypted %d bytes", n)
}

func signFile() {
	pubKey := decodePublicKey(*publicKey)
	privKey := decodePrivateKey(*privateKey)

	signer := createEntityFromKeys(pubKey, privKey)

	err := openpgp.ArmoredDetachSign(os.Stdout, signer, os.Stdin, nil)
	kingpin.FatalIfError(err, "Error signing input")
}

func verifyFile() {
	pubKey := decodePublicKey(*publicKey)
	sig := decodeSignature(*signatureFile)

	hash := sig.Hash.New()
	io.Copy(hash, os.Stdin)

	err := pubKey.VerifySignature(hash, sig)
	kingpin.FatalIfError(err, "Error signing input")
	kingpin.Errorf("Verified signature")
}

func createEntityFromKeys(pubKey *packet.PublicKey, privKey *packet.PrivateKey) *openpgp.Entity {
	config := packet.Config{
		DefaultHash:            crypto.SHA256,
		DefaultCipher:          packet.CipherAES256,
		DefaultCompressionAlgo: packet.CompressionZLIB,
		CompressionConfig: &packet.CompressionConfig{
			Level: 9,
		},
		RSABits: *bits,
	}
	currentTime := config.Now()
	uid := packet.NewUserId("", "", "")

	e := openpgp.Entity{
		PrimaryKey: pubKey,
		PrivateKey: privKey,
		Identities: make(map[string]*openpgp.Identity),
	}
	isPrimaryId := false

	e.Identities[uid.Id] = &openpgp.Identity{
		Name:   uid.Name,
		UserId: uid,
		SelfSignature: &packet.Signature{
			CreationTime: currentTime,
			SigType:      packet.SigTypePositiveCert,
			PubKeyAlgo:   packet.PubKeyAlgoRSA,
			Hash:         config.Hash(),
			IsPrimaryId:  &isPrimaryId,
			FlagsValid:   true,
			FlagSign:     true,
			FlagCertify:  true,
			IssuerKeyId:  &e.PrimaryKey.KeyId,
		},
	}

	keyLifetimeSecs := uint32(86400 * 365)

	e.Subkeys = make([]openpgp.Subkey, 1)
	e.Subkeys[0] = openpgp.Subkey{
		PublicKey: pubKey,
		PrivateKey: privKey,
		Sig: &packet.Signature{
			CreationTime:              currentTime,
			SigType:                   packet.SigTypeSubkeyBinding,
			PubKeyAlgo:                packet.PubKeyAlgoRSA,
			Hash:                      config.Hash(),
			PreferredHash:             []uint8{8}, // SHA-256
			FlagsValid:                true,
			FlagEncryptStorage:        true,
			FlagEncryptCommunications: true,
			IssuerKeyId:               &e.PrimaryKey.KeyId,
			KeyLifetimeSecs:           &keyLifetimeSecs,
		},
	}
	return &e
}

func generateKeys() {
	key, err := rsa.GenerateKey(rand.Reader, *bits)
	kingpin.FatalIfError(err, "Error generating RSA key: %s", err)

	priv, err := os.Create(filepath.Join(*keyOutputDir, *keyOutputPrefix+".privkey"))
	kingpin.FatalIfError(err, "Error writing private key to file: %s", err)
	defer priv.Close()

	pub, err := os.Create(filepath.Join(*keyOutputDir, *keyOutputPrefix+".pubkey"))
	kingpin.FatalIfError(err, "Error writing public key to file: %s", err)
	defer pub.Close()

	encodePrivateKey(priv, key)
	encodePublicKey(pub, key)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/ambientsound/pms/console"
	"github.com/ambientsound/pms/pms"
	"github.com/ambientsound/pms/version"
	"github.com/ambientsound/pms/xdg"

	"github.com/jessevdk/go-flags"
)

var buildVersion = "undefined"

type cliOptions struct {
	Version     bool   `short:"v" long:"version" description:"Print program version"`
	Debug       string `short:"d" long:"debug" description:"Write debugging info to file"`
	MpdHost     string `long:"host" description:"MPD host (MPD_HOST environment variable)" default:"localhost"`
	MpdPort     string `long:"port" description:"MPD port (MPD_PORT environment variable)" default:"6600"`
	MpdPassword string `long:"password" description:"MPD password"`
}

func main() {
	var opts cliOptions

	version.SetVersion(buildVersion)
	fmt.Printf("%s %s\n", version.LongName(), version.Version())

	remainder, err := flags.Parse(&opts)
	if err != nil {
		os.Exit(1)
	}
	if len(remainder) > 0 {
		trailing := strings.Join(remainder, " ")
		fmt.Printf("error: trailing characters: %s\n", trailing)
		os.Exit(1)
	}

	if len(opts.Debug) > 0 {
		err := console.Open(opts.Debug)
		if err != nil {
			fmt.Printf("Error while opening log file: %s", err)
			os.Exit(1)
		}
	}

	if opts.Version {
		os.Exit(0)
	}

	console.Log("Starting Practical Music Search.")

	val, ok := os.LookupEnv("MPD_HOST")
	if ok {
		opts.MpdHost = val
	}
	val, ok = os.LookupEnv("MPD_PORT")
	if ok {
		opts.MpdPort = val
	}

	pms := pms.New()
	defer func() {
		pms.QuitSignal <- 0
	}()

	// Source default configuration.
	pms.Message("Applying default configuration.")
	if err := pms.SourceDefaultConfig(); err != nil {
		panic(fmt.Sprintf("BUG in default config: %s\n", err))
	}

	// Source configuration files from all XDG standard directories.
	configDirs := xdg.ConfigDirectories()
	for _, dir := range configDirs {
		p := path.Join(dir, "pms.conf")
		pms.Message("Reading configuration file '%s'.", p)
		err = pms.SourceConfigFile(p)
		if err != nil {
			pms.Error("Error while reading configuration file '%s': %s", p, err)
		}
	}

	pms.SetConnectionParams(opts.MpdHost, opts.MpdPort, opts.MpdPassword)
	go pms.LoopConnect()

	pms.Main()
	pms.Wait()

	console.Log("Exiting normally.")
}
<commit_msg>Make command-line host and port options override env vars<commit_after>package main

import (
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/ambientsound/pms/console"
	"github.com/ambientsound/pms/pms"
	"github.com/ambientsound/pms/version"
	"github.com/ambientsound/pms/xdg"

	"github.com/jessevdk/go-flags"
)

var buildVersion = "undefined"

type cliOptions struct {
	Version     bool   `short:"v" long:"version" description:"Print program version"`
	Debug       string `short:"d" long:"debug" description:"Write debugging info to file"`
	MpdHost     string `long:"host" description:"MPD host" default-mask:"MPD_HOST environment variable or localhost"`
	MpdPort     string `long:"port" description:"MPD port" default-mask:"MPD_PORT environment variable or 6600"`
	MpdPassword string `long:"password" description:"MPD password"`
}

func main() {
	var opts cliOptions

	version.SetVersion(buildVersion)
	fmt.Printf("%s %s\n", version.LongName(), version.Version())

	remainder, err := flags.Parse(&opts)
	if err != nil {
		os.Exit(1)
	}
	if len(remainder) > 0 {
		trailing := strings.Join(remainder, " ")
		fmt.Printf("error: trailing characters: %s\n", trailing)
		os.Exit(1)
	}

	if len(opts.Debug) > 0 {
		err := console.Open(opts.Debug)
		if err != nil {
			fmt.Printf("Error while opening log file: %s", err)
			os.Exit(1)
		}
	}

	if opts.Version {
		os.Exit(0)
	}

	console.Log("Starting Practical Music Search.")

	if len(opts.MpdHost) == 0 {
		val, ok := os.LookupEnv("MPD_HOST")
		if ok {
			opts.MpdHost = val
		} else {
			opts.MpdHost = "localhost"
		}
	}
	if len(opts.MpdPort) == 0 {
		val, ok := os.LookupEnv("MPD_PORT")
		if ok {
			opts.MpdPort = val
		} else {
			opts.MpdPort = "6600"
		}
	}

	pms := pms.New()
	defer func() {
		pms.QuitSignal <- 0
	}()

	// Source default configuration.
	pms.Message("Applying default configuration.")
	if err := pms.SourceDefaultConfig(); err != nil {
		panic(fmt.Sprintf("BUG in default config: %s\n", err))
	}

	// Source configuration files from all XDG standard directories.
	configDirs := xdg.ConfigDirectories()
	for _, dir := range configDirs {
		p := path.Join(dir, "pms.conf")
		pms.Message("Reading configuration file '%s'.", p)
		err = pms.SourceConfigFile(p)
		if err != nil {
			pms.Error("Error while reading configuration file '%s': %s", p, err)
		}
	}

	pms.SetConnectionParams(opts.MpdHost, opts.MpdPort, opts.MpdPassword)
	go pms.LoopConnect()

	pms.Main()
	pms.Wait()

	console.Log("Exiting normally.")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"github.com/golang/gddo/httputil"
	"github.com/gorilla/mux"
)

const (
	ContentText = 1
	ContentJSON = 2
	ContentYAML = 3

	// The top-level key in the JSON for the default (not client-specific answers)
	DEFAULT_KEY = "default"

	// A key to check for magic traversing of arrays by a string field in them
	// For example, given: { things: [ {name: 'asdf', stuff: 42}, {name: 'zxcv', stuff: 43} ] }
	// Both ../things/0/stuff and ../things/asdf/stuff will return 42 because 'asdf' matched the 'anme' field of one of the 'things'.
	MAGIC_ARRAY_KEY = "name"
)

var (
	debug       = flag.Bool("debug", false, "Debug")
	enableXff   = flag.Bool("xff", false, "X-Forwarded-For header support")
	listen      = flag.String("listen", ":80", "Address to listen to (TCP)")
	answersFile = flag.String("answers", "./answers.yaml", "File containing the answers to respond with")
	logFile     = flag.String("log", "", "Log file")
	pidFile     = flag.String("pid-file", "", "PID to write to")

	router  = mux.NewRouter()
	answers Versions

	wantRevision   = 1
	loadedRevision = 0
	loading        = false
	ticker         = time.NewTicker(1 * time.Second)
)

func main() {
	log.Info("Starting rancher-metadata")
	parseFlags()
	err := loadAnswers()
	if err != nil {
		log.Fatal("Cannot startup without a valid Answers file")
	}
	watchSignals()

	router.HandleFunc("/favicon.ico", http.NotFound)
	router.HandleFunc("/", root).
		Methods("GET", "HEAD").
		Name("Root")

	router.HandleFunc("/{version}", metadata).
		Methods("GET", "HEAD").
		Name("Version")

	router.HandleFunc("/{version}/{key:.*}", metadata).
		Methods("GET", "HEAD").
		Name("Metadata")

	log.Info("Listening on ", *listen)
	log.Fatal(http.ListenAndServe(*listen, router))
}

func parseFlags() {
	flag.Parse()

	if *debug {
		log.SetLevel(log.DebugLevel)
	}

	if *logFile != "" {
		if output, err := os.OpenFile(*logFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666); err != nil {
			log.Fatalf("Failed to log to file %s: %v", *logFile, err)
		} else {
			log.SetOutput(output)
		}
	}

	if *pidFile != "" {
		log.Infof("Writing pid %d to %s", os.Getpid(), *pidFile)
		if err := ioutil.WriteFile(*pidFile, []byte(strconv.Itoa(os.Getpid())), 0644); err != nil {
			log.Fatalf("Failed to write pid file %s: %v", *pidFile, err)
		}
	}
}

func loadAnswers() (err error) {
	log.Info("Loading answers")
	loading = true
	revision := wantRevision
	neu, err := ParseAnswers(*answersFile)
	if err == nil {
		for _, data := range neu {
			defaults, ok := data[DEFAULT_KEY]
			if ok {
				defaultsMap, ok := defaults.(map[string]interface{})
				if ok {
					// Copy the defaults into the per-client info, so there's no
					// complicated merging and lookup logic when retrieving.
					mergeDefaults(&data, defaultsMap)
				}
			}
		}

		answers = neu
		loadedRevision = revision
		loading = false
		log.Infof("Loaded answers for %d versions", len(answers))
	} else {
		log.Errorf("Failed to load answers: %v", err)
	}
	return err
}

func mergeDefaults(clientAnswers *Answers, defaultAnswers map[string]interface{}) {
	for _, client := range *clientAnswers {
		clientMap, ok := client.(map[string]interface{})
		if ok {
			for k, v := range defaultAnswers {
				_, exists := clientMap[k]
				if !exists {
					clientMap[k] = v
				}
			}
		}
	}
}

func watchSignals() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGHUP)

	go func() {
		for _ = range c {
			log.Info("Received HUP signal")
			wantRevision += 1
		}
	}()

	go func() {
		for {
			select {
			case <-ticker.C:
				//log.Info("Ping")
				if !loading && loadedRevision < wantRevision {
					loadAnswers()
				}
			}
		}
	}()

}

func contentType(req *http.Request) int {
	str := httputil.NegotiateContentType(req, []string{
		"text/plain",
		"application/json",
		"application/yaml",
		"application/x-yaml",
		"text/yaml",
		"text/x-yaml",
	}, "text/plain")

	if strings.Contains(str, "json") {
		return ContentJSON
	} else if strings.Contains(str, "yaml") {
		return ContentYAML
	} else {
		return ContentText
	}
}

func root(w http.ResponseWriter, req *http.Request) {
	log.WithFields(log.Fields{"client": requestIp(req), "version": "root"}).Infof("OK: %s", "/")

	m := make(map[string]interface{})
	for _, k := range answers.Versions() {
		url, err := router.Get("Version").URL("version", k)
		if err == nil {
			m[k] = (*url).String()
		} else {
			log.Warn("Error: ", err.Error())
		}
	}

	// If latest isn't in the list, pretend it is
	_, ok := m["latest"]
	if !ok {
		url, err := router.Get("Version").URL("version", "latest")
		if err == nil {
			m["latest"] = (*url).String()
		} else {
			log.Warn("Error: ", err.Error())
		}
	}

	respondSuccess(w, req, m)
}

func metadata(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	clientIp := requestIp(req)

	version := vars["version"]
	_, ok := answers[version]
	if !ok {
		// If a `latest` key is not provided, pick the ASCII-betically highest version and call it that.
		if version == "latest" {
			version = ""
			for _, k := range answers.Versions() {
				if k > version {
					version = k
				}
			}

			log.Debugf("Picked %s for latest version because none provided", version)
		} else {
			respondError(w, req, "Invalid version", http.StatusNotFound)
			return
		}
	}

	key := strings.TrimRight(vars["key"], "/")
	displayKey := "/" + key

	log.WithFields(log.Fields{"version": version, "client": clientIp}).Debugf("Searching for: %s", key)
	val, ok := answers.Matching(version, clientIp, key)

	if ok {
		log.WithFields(log.Fields{"version": version, "client": clientIp}).Infof("OK: %s", displayKey)
		respondSuccess(w, req, val)
	} else {
		log.WithFields(log.Fields{"version": version, "client": clientIp}).Infof("Error: %s", displayKey)
		respondError(w, req, "Not found", http.StatusNotFound)
	}
}

func respondError(w http.ResponseWriter, req *http.Request, msg string, statusCode int) {
	obj := make(map[string]interface{})
	obj["message"] = msg
	obj["type"] = "error"
	obj["code"] = statusCode

	switch contentType(req) {
	case ContentText:
		http.Error(w, msg, statusCode)
	case ContentJSON:
		bytes, err := json.Marshal(obj)
		if err == nil {
			http.Error(w, string(bytes), statusCode)
		} else {
			http.Error(w, "{\"type\": \"error\", \"message\": \"JSON marshal error\"}", http.StatusInternalServerError)
		}
	case ContentYAML:
		bytes, err := yaml.Marshal(obj)
		if err == nil {
			http.Error(w, string(bytes), statusCode)
		} else {
			http.Error(w, "type: \"error\"\nmessage: \"JSON marshal error\"", http.StatusInternalServerError)
		}
	}
}

func respondSuccess(w http.ResponseWriter, req *http.Request, val interface{}) {
	switch contentType(req) {
	case ContentText:
		respondText(w, req, val)
	case ContentJSON:
		respondJSON(w, req, val)
	case ContentYAML:
		respondYAML(w, req, val)
	}
}

func respondText(w http.ResponseWriter, req *http.Request, val interface{}) {
	if val == nil {
		fmt.Fprint(w, "")
		return
	}

	switch v := val.(type) {
	case string:
		fmt.Fprint(w, v)
	case uint, uint8, uint16, uint32, uint64, int, int8, int16, int32, int64:
		fmt.Fprintf(w, "%d", v)
	case float64:
		// The default format has extra trailing zeros
		str := strings.TrimRight(fmt.Sprintf("%f", v), "0")
		str = strings.TrimRight(str, ".")
		fmt.Fprint(w, str)
	case bool:
		if v {
			fmt.Fprint(w, "true")
		} else {
			fmt.Fprint(w, "false")
		}
	case map[string]interface{}:
		out := make([]string, len(v))
		i := 0
		for k, vv := range v {
			_, isMap := vv.(map[string]interface{})
			_, isArray := vv.([]interface{})
			if isMap || isArray {
				out[i] = fmt.Sprintf("%s/\n", k)
			} else {
				out[i] = fmt.Sprintf("%s\n", k)
			}
			i++
		}

		sort.Strings(out)
		for _, vv := range out {
			fmt.Fprint(w, vv)
		}
	case []interface{}:
		for k, vv := range v {
			vvMap, isMap := vv.(map[string]interface{})
			_, isArray := vv.([]interface{})

			if isMap {
				// If the child is a map and has a "name" property, show "key=name"
				name, ok := vvMap[MAGIC_ARRAY_KEY]
				if ok {
					fmt.Fprintf(w, "%d=%s\n", k, name)
					continue
				}
			}

			if isMap || isArray {
				// If the child is a map or array, show "key/"
				fmt.Fprintf(w, "%d/\n", k)
			} else {
				// Otherwise, show "key"
				fmt.Fprintf(w, "%d\n", k)
			}
		}
	default:
		http.Error(w, "Value is of a type I don't know how to handle", http.StatusInternalServerError)
	}
}

func respondJSON(w http.ResponseWriter, req *http.Request, val interface{}) {
	bytes, err := json.Marshal(val)
	if err == nil {
		w.Write(bytes)
	} else {
		respondError(w, req, "Error serializing to JSON: "+err.Error(), http.StatusInternalServerError)
	}
}

func respondYAML(w http.ResponseWriter, req *http.Request, val interface{}) {
	bytes, err := yaml.Marshal(val)
	if err == nil {
		w.Write(bytes)
	} else {
		respondError(w, req, "Error serializing to YAML: "+err.Error(), http.StatusInternalServerError)
	}
}

func requestIp(req *http.Request) string {
	if *enableXff {
		clientIp := req.Header.Get("X-Forwarded-For")
		if len(clientIp) > 0 {
			return clientIp
		}
	}

	clientIp, _, _ := net.SplitHostPort(req.RemoteAddr)
	return clientIp
}
<commit_msg>Listen on 127.0.0.1:8112 for POST /v1/reload, CORS header<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"github.com/golang/gddo/httputil"
	"github.com/gorilla/mux"
)

const (
	ContentText = 1
	ContentJSON = 2
	ContentYAML = 3

	// The top-level key in the JSON for the default (not client-specific answers)
	DEFAULT_KEY = "default"

	// A key to check for magic traversing of arrays by a string field in them
	// For example, given: { things: [ {name: 'asdf', stuff: 42}, {name: 'zxcv', stuff: 43} ] }
	// Both ../things/0/stuff and ../things/asdf/stuff will return 42 because 'asdf' matched the 'anme' field of one of the 'things'.
	MAGIC_ARRAY_KEY = "name"
)

var (
	debug        = flag.Bool("debug", false, "Debug")
	enableXff    = flag.Bool("xff", false, "X-Forwarded-For header support")
	listen       = flag.String("listen", ":80", "Address to listen to (TCP)")
	listenReload = flag.String("listenReload", "127.0.0.1:8112", "Address to listen to for reload requests (TCP)")
	answersFile  = flag.String("answers", "./answers.yaml", "File containing the answers to respond with")
	logFile      = flag.String("log", "", "Log file")
	pidFile      = flag.String("pid-file", "", "PID to write to")

	router  = mux.NewRouter()
	answers Versions

	wantRevision   = 1
	loadedRevision = 0
	loading        = false
	ticker         = time.NewTicker(1 * time.Second)
)

func main() {
	log.Info("Starting rancher-metadata")
	parseFlags()
	err := loadAnswers()
	if err != nil {
		log.Fatal("Cannot startup without a valid Answers file")
	}

	watchSignals()
	watchHttp()

	router.HandleFunc("/favicon.ico", http.NotFound)
	router.HandleFunc("/", root).
		Methods("GET", "HEAD").
		Name("Root")

	router.HandleFunc("/{version}", metadata).
		Methods("GET", "HEAD").
		Name("Version")

	router.HandleFunc("/{version}/{key:.*}", metadata).
		Methods("GET", "HEAD").
		Name("Metadata")

	log.Info("Listening on ", *listen)
	log.Fatal(http.ListenAndServe(*listen, router))
}

func parseFlags() {
	flag.Parse()

	if *debug {
		log.SetLevel(log.DebugLevel)
	}

	if *logFile != "" {
		if output, err := os.OpenFile(*logFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666); err != nil {
			log.Fatalf("Failed to log to file %s: %v", *logFile, err)
		} else {
			log.SetOutput(output)
		}
	}

	if *pidFile != "" {
		log.Infof("Writing pid %d to %s", os.Getpid(), *pidFile)
		if err := ioutil.WriteFile(*pidFile, []byte(strconv.Itoa(os.Getpid())), 0644); err != nil {
			log.Fatalf("Failed to write pid file %s: %v", *pidFile, err)
		}
	}
}

func loadAnswers() (err error) {
	log.Debug("Loading answers")
	loading = true
	revision := wantRevision
	neu, err := ParseAnswers(*answersFile)
	if err == nil {
		for _, data := range neu {
			defaults, ok := data[DEFAULT_KEY]
			if ok {
				defaultsMap, ok := defaults.(map[string]interface{})
				if ok {
					// Copy the defaults into the per-client info, so there's no
					// complicated merging and lookup logic when retrieving.
					mergeDefaults(&data, defaultsMap)
				}
			}
		}

		answers = neu
		loadedRevision = revision
		loading = false
		log.Infof("Loaded answers revision %d for %d versions", revision, len(answers))
	} else {
		log.Errorf("Failed to load answers: %v", err)
	}
	return err
}

func mergeDefaults(clientAnswers *Answers, defaultAnswers map[string]interface{}) {
	for _, client := range *clientAnswers {
		clientMap, ok := client.(map[string]interface{})
		if ok {
			for k, v := range defaultAnswers {
				_, exists := clientMap[k]
				if !exists {
					clientMap[k] = v
				}
			}
		}
	}
}

func watchSignals() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGHUP)

	go func() {
		for _ = range c {
			log.Info("Received HUP signal")
			wantRevision += 1
		}
	}()

	go func() {
		for {
			select {
			case <-ticker.C:
				//log.Info("Ping")
				if !loading && loadedRevision < wantRevision {
					loadAnswers()
				}
			}
		}
	}()

}

func watchHttp() {
	reloadRouter := mux.NewRouter()
	reloadRouter.HandleFunc("/favicon.ico", http.NotFound)
	reloadRouter.HandleFunc("/v1/reload", httpReload).Methods("POST")

	log.Info("Listening for Reload on ", *listenReload)
	go http.ListenAndServe(*listenReload, reloadRouter)
}

func httpReload(w http.ResponseWriter, req *http.Request) {
	wantRevision += 1
	waitFor := wantRevision
	log.Debugf("Received HTTP reload request, wait for %d, ", waitFor)

	for {
		select {
		case <-time.After(100 * time.Millisecond):
			//log.Debugf("Now at %d, waiting for %d", loadedRevision, waitFor)
			if loadedRevision >= waitFor {
				fmt.Fprintf(w, "OK %d\r\n", loadedRevision)
				return
			}
		}
	}
}

func contentType(req *http.Request) int {
	str := httputil.NegotiateContentType(req, []string{
		"text/plain",
		"application/json",
		"application/yaml",
		"application/x-yaml",
		"text/yaml",
		"text/x-yaml",
	}, "text/plain")

	if strings.Contains(str, "json") {
		return ContentJSON
	} else if strings.Contains(str, "yaml") {
		return ContentYAML
	} else {
		return ContentText
	}
}

func root(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")

	log.WithFields(log.Fields{"client": requestIp(req), "version": "root"}).Infof("OK: %s", "/")

	m := make(map[string]interface{})
	for _, k := range answers.Versions() {
		url, err := router.Get("Version").URL("version", k)
		if err == nil {
			m[k] = (*url).String()
		} else {
			log.Warn("Error: ", err.Error())
		}
	}

	// If latest isn't in the list, pretend it is
	_, ok := m["latest"]
	if !ok {
		url, err := router.Get("Version").URL("version", "latest")
		if err == nil {
			m["latest"] = (*url).String()
		} else {
			log.Warn("Error: ", err.Error())
		}
	}

	respondSuccess(w, req, m)
}

func metadata(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")

	vars := mux.Vars(req)
	clientIp := requestIp(req)

	version := vars["version"]
	_, ok := answers[version]
	if !ok {
		// If a `latest` key is not provided, pick the ASCII-betically highest version and call it that.
		if version == "latest" {
			version = ""
			for _, k := range answers.Versions() {
				if k > version {
					version = k
				}
			}

			log.Debugf("Picked %s for latest version because none provided", version)
		} else {
			respondError(w, req, "Invalid version", http.StatusNotFound)
			return
		}
	}

	key := strings.TrimRight(vars["key"], "/")
	displayKey := "/" + key

	log.WithFields(log.Fields{"version": version, "client": clientIp}).Debugf("Searching for: %s", key)
	val, ok := answers.Matching(version, clientIp, key)

	if ok {
		log.WithFields(log.Fields{"version": version, "client": clientIp}).Infof("OK: %s", displayKey)
		respondSuccess(w, req, val)
	} else {
		log.WithFields(log.Fields{"version": version, "client": clientIp}).Infof("Error: %s", displayKey)
		respondError(w, req, "Not found", http.StatusNotFound)
	}
}

func respondError(w http.ResponseWriter, req *http.Request, msg string, statusCode int) {
	obj := make(map[string]interface{})
	obj["message"] = msg
	obj["type"] = "error"
	obj["code"] = statusCode

	switch contentType(req) {
	case ContentText:
		http.Error(w, msg, statusCode)
	case ContentJSON:
		bytes, err := json.Marshal(obj)
		if err == nil {
			http.Error(w, string(bytes), statusCode)
		} else {
			http.Error(w, "{\"type\": \"error\", \"message\": \"JSON marshal error\"}", http.StatusInternalServerError)
		}
	case ContentYAML:
		bytes, err := yaml.Marshal(obj)
		if err == nil {
			http.Error(w, string(bytes), statusCode)
		} else {
			http.Error(w, "type: \"error\"\nmessage: \"JSON marshal error\"", http.StatusInternalServerError)
		}
	}
}

func respondSuccess(w http.ResponseWriter, req *http.Request, val interface{}) {
	switch contentType(req) {
	case ContentText:
		respondText(w, req, val)
	case ContentJSON:
		respondJSON(w, req, val)
	case ContentYAML:
		respondYAML(w, req, val)
	}
}

func respondText(w http.ResponseWriter, req *http.Request, val interface{}) {
	if val == nil {
		fmt.Fprint(w, "")
		return
	}

	switch v := val.(type) {
	case string:
		fmt.Fprint(w, v)
	case uint, uint8, uint16, uint32, uint64, int, int8, int16, int32, int64:
		fmt.Fprintf(w, "%d", v)
	case float64:
		// The default format has extra trailing zeros
		str := strings.TrimRight(fmt.Sprintf("%f", v), "0")
		str = strings.TrimRight(str, ".")
		fmt.Fprint(w, str)
	case bool:
		if v {
			fmt.Fprint(w, "true")
		} else {
			fmt.Fprint(w, "false")
		}
	case map[string]interface{}:
		out := make([]string, len(v))
		i := 0
		for k, vv := range v {
			_, isMap := vv.(map[string]interface{})
			_, isArray := vv.([]interface{})
			if isMap || isArray {
				out[i] = fmt.Sprintf("%s/\n", k)
			} else {
				out[i] = fmt.Sprintf("%s\n", k)
			}
			i++
		}

		sort.Strings(out)
		for _, vv := range out {
			fmt.Fprint(w, vv)
		}
	case []interface{}:
		for k, vv := range v {
			vvMap, isMap := vv.(map[string]interface{})
			_, isArray := vv.([]interface{})

			if isMap {
				// If the child is a map and has a "name" property, show "key=name"
				name, ok := vvMap[MAGIC_ARRAY_KEY]
				if ok {
					fmt.Fprintf(w, "%d=%s\n", k, name)
					continue
				}
			}

			if isMap || isArray {
				// If the child is a map or array, show "key/"
				fmt.Fprintf(w, "%d/\n", k)
			} else {
				// Otherwise, show "key"
				fmt.Fprintf(w, "%d\n", k)
			}
		}
	default:
		http.Error(w, "Value is of a type I don't know how to handle", http.StatusInternalServerError)
	}
}

func respondJSON(w http.ResponseWriter, req *http.Request, val interface{}) {
	bytes, err := json.Marshal(val)
	if err == nil {
		w.Write(bytes)
	} else {
		respondError(w, req, "Error serializing to JSON: "+err.Error(), http.StatusInternalServerError)
	}
}

func respondYAML(w http.ResponseWriter, req *http.Request, val interface{}) {
	bytes, err := yaml.Marshal(val)
	if err == nil {
		w.Write(bytes)
	} else {
		respondError(w, req, "Error serializing to YAML: "+err.Error(), http.StatusInternalServerError)
	}
}

func requestIp(req *http.Request) string {
	if *enableXff {
		clientIp := req.Header.Get("X-Forwarded-For")
		if len(clientIp) > 0 {
			return clientIp
		}
	}

	clientIp, _, _ := net.SplitHostPort(req.RemoteAddr)
	return clientIp
}
<|endoftext|>


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

package main

import (
	"bufio"
	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"code.google.com/p/gopacket/tcpassembly"
	"code.google.com/p/gopacket/tcpassembly/tcpreader"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"
)

/* Command line arguments */
var iface = flag.String("ifce", "en0", "Interface to get packets from")
var inputfile = flag.String("infile", "", "read packets from file")
var logAllPackets = flag.Bool("v", false, "Logs every packet in great detail")
var launchCmd = flag.String("e", "", "Launches the command and logs its traffic")
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

type interval struct {
	req  *http.Request
	resp *http.Response
}

type BidiStream struct {
	key     uint64
	in, out *TCPStream
}

// TCPStream will handle the actual decoding of http requests and responses.
type TCPStream struct {
	netFlow, tcpFlow gopacket.Flow
	readStream       tcpreader.ReaderStream
	storage          *Storage
	bidikey          uint64
	closed           bool
}

/* This reads both HTTP requests and HTTP responses in two separate streams */
func (h *TCPStream) runOut() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64

	for {
		/*		_, err := buf.Peek(1)
				if err == io.EOF {
					return
				}*/
		req, err := http.ReadRequest(buf)
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)

			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			// bodyBytes := tcpreader.DiscardBytesToEOF(req.Body)
			req.Body.Close()
			err = h.storage.SentRequest(h.bidikey, reqID, time.Now(), req)
			if err != nil {
				log.Println("Error storing request", err)
			}

			reqID++
			//			fmt.Print(".")
			// log.Println("Received request from stream", h.netFlow, h.tcpFlow,
			// 	":", req, "with", bodyBytes, "bytes in request body")
		}
	}
}

func (h *TCPStream) runIn() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64
	for {
		/* Don't start reading a response if no data is available */
		_, err := buf.Peek(1)
		if err == io.EOF {
			return
		}
		/* Data available, read response */
		resp, err := http.ReadResponse(buf, nil) // TODO: request
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			//			log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)
			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			_, err := tcpreader.DiscardBytesToFirstError(resp.Body)
			if err != nil && err != io.EOF {
				log.Println("Error discarding bytes ", err)
			}
			resp.Body.Close()
			err = h.storage.ReceivedResponse(h.bidikey, reqID, time.Now(), resp)
			if err != nil {
				log.Println("Error storing response", err)
			}
			reqID++
			// fmt.Print(".")
			//log.Println("Received response from stream", h.netFlow, h.tcpFlow,
			//	":", resp, "with", bodyBytes, "bytes in response body")
		}
		/* Match the response with the next request */
	}

}

/* httpStreamFactory implements tcpassembly.StreamFactory */
type httpStreamFactory struct {
	bidiStreams map[uint64]*BidiStream
	storage     *Storage
	closed      bool
}

func NewStreamFactory(s *Storage) *httpStreamFactory {
	return &httpStreamFactory{bidiStreams: make(map[uint64]*BidiStream),
		storage: s}
}
func (h *httpStreamFactory) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {

	/* Watch out: this function can still get called even after all
	   streams were flushed (via FlushAll) and closed. */
	/*	if h.closed == true {
			return tcpreader.NewReaderStream()
		}
	*/
	/* First the outgoing stream, then the incoming stream */
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	hstream := &TCPStream{
		netFlow:    netFlow,
		tcpFlow:    tcpFlow,
		readStream: tcpreader.NewReaderStream(),
		storage:    h.storage,
		bidikey:    key,
	}

	bds := h.bidiStreams[key]
	if bds == nil {
		//		log.Println("reading stream", netFlow, tcpFlow)
		bds = &BidiStream{out: hstream, key: key}
		h.bidiStreams[key] = bds
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runOut()
	} else {
		//		log.Println("opening TCP conn", netFlow, tcpFlow)
		bds.in = hstream
		err := h.storage.OpenTCPConnection(key, time.Now())
		if err != nil {
			log.Println("Error storing connection", err)
		}
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runIn()
	}

	// ReaderStream implements tcpassembly.Stream, so we can return a pointer to it.
	return &hstream.readStream
}

func (h *httpStreamFactory) LogPacketSize(packet gopacket.Packet) {
	netFlow := packet.NetworkLayer().NetworkFlow()
	tcpFlow := packet.TransportLayer().TransportFlow()
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	ipv4Layer := packet.Layer(layers.LayerTypeIPv4)
	ipv4, _ := ipv4Layer.(*layers.IPv4)

	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	tcp, _ := tcpLayer.(*layers.TCP)

	bds := h.bidiStreams[key]
	if bds == nil || bds.in == nil || bds.out == nil {
		return
	}

	payloadLength := uint32(ipv4.Length - uint16(ipv4.IHL)*4 - uint16(tcp.DataOffset)*4)

	if bds.in.netFlow == netFlow {
		/* incoming */
		err := h.storage.IncomingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	} else {
		/* outgoing */
		err := h.storage.OutgoingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	}
}

func create_process_ended_channel(cmd *exec.Cmd) (cmd_done chan error) {
	cmd_done = make(chan error, 1)
	go func() {
		cmd_done <- cmd.Wait()
	}()
	return cmd_done
}

/* Create a channel and send any CTRL-C signal over it */
func create_ctrl_c_handler() (timeout chan os.Signal) {
	ctrlc := make(chan os.Signal, 1)
	signal.Notify(ctrlc, os.Interrupt)
	return ctrlc
}

/* Create a channel that will contain TRUE after T seconds. */
func create_timeout_channel(t time.Duration) (timeout chan bool) {
	timeout = make(chan bool, 1)
	go func() {
		time.Sleep(t * time.Second)
		timeout <- true
	}()
	return timeout
}

func main() {
	flag.Parse()
	log.SetFlags(log.Ltime | log.Lmicroseconds)
	//	log.SetOutput(ioutil.Discard)

	// Setup profiler
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// Set up storage layer
	storage, err := NewStorage()
	if err != nil {
		panic(err)
	}

	// Set up assembly
	streamFactory := NewStreamFactory(storage)
	streamPool := tcpassembly.NewStreamPool(streamFactory)
	assembler := tcpassembly.NewAssembler(streamPool)

	// Setup CTRL-C handler channel
	ctrlc := create_ctrl_c_handler()

	// Setup channel that reports all socket kernel events (Mac OS X only)
	//	netDescSource := NewOSXNetDescSource()
	//	descriptors := netDescSource.Descriptors()
	var handle *pcap.Handle
	if *inputfile != "" {
		handle, err = pcap.OpenOffline(*inputfile)
	} else {
		log.Printf("starting capture on interface %q", *iface)
		// Setup packet capture
		handle, err = pcap.OpenLive(*iface, 1600, true, pcap.BlockForever)
	}

	if err != nil {
		panic(err)
	} else if err := handle.SetBPFFilter("tcp and port 80"); err != nil {
		panic(err)
	}
	log.Println("reading in packets. Press CTRL-C to end and report.")
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets := packetSource.Packets()

	// Run the external command
	//pid := uint32(0)
	var cmd_done chan error
	var start_time time.Time
	if *launchCmd != "" {
		s := *launchCmd
		//		n := strings.Index(*launchCmd, " ")
		//		fmt.Println("Launching ", s[0:n], " --- ", s[n:])
		//		cmd := exec.Command(s[0:n], strings.Split(s[n:], " "))
		args := strings.Split(s, " ")
		cmd := exec.Command(args[0], args[1:]...)
		start_time = time.Now()
		cmd_stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		err = cmd.Start()
		if err != nil {
			panic(err)
		}
		go io.Copy(os.Stdout, cmd_stdout)
		//		cmd.Stdout = os.Stdout
		//		cmd.Stderr = os.Stderr
		// Create the channel that listens for the end of the command.
		cmd_done = create_process_ended_channel(cmd)
		//pid = uint32(cmd.Process.Pid)
	}

	var timeout chan bool
loop:
	for {
		select {
		/*		case netDesc := <-descriptors:
				if netDesc.Pid == pid {
					fmt.Println("event received ", netDesc)
				}*/
		case packet, ok := <-packets:
			if !ok {
				log.Println("All data read")
				packets = nil
				timeout = create_timeout_channel(0)
				break
			}

			if packet == nil {
				break
			}

			if *logAllPackets {
				log.Println(packet)
			}

			if storage.PacketInScope(packet) {
				streamFactory.LogPacketSize(packet)
				netFlow := packet.NetworkLayer().NetworkFlow()
				tcp := packet.TransportLayer().(*layers.TCP)

				assembler.AssembleWithTimestamp(netFlow, tcp,
					packet.Metadata().Timestamp)
			}
		case err := <-cmd_done:
			if err != nil {
				log.Printf("process done with error = %v\n", err)
			}

			log.Println("Process took: ", time.Now().Sub(start_time))

			/* Wait for a couple of seconds, just enough to get the events
			   handled by the main function. */
			log.Println("Waiting for the remaining responses to arrive.")
			timeout = create_timeout_channel(10)
		case <-ctrlc:
			/* Don't wait. */
			timeout = create_timeout_channel(0)
		case <-timeout:
			break loop
		}
	}

	/* Cleanup the go routines */
	/* Ignore any http request/response parsing errors when closing the streams. */
	streamFactory.closed = true
	for _, v := range streamFactory.bidiStreams {
		if v.in != nil {
			v.in.closed = true
		}
		if v.out != nil {
			v.out.closed = true
		}
	}

	assembler.FlushAll()

	/* Close the storage layer. This will block until all pending inserts in
	   the db are handled. */
	storage.Close()

	reporting, err := NewReporting()
	if err != nil {
		panic(err)
	}

	if err = reporting.Report(); err != nil {
		log.Panic(err)
	}

	os.Exit(0)

}
<commit_msg>Close the control-C signal handler after the main loop, so that long running reports can be interrupted.<commit_after>// Copyright 2014 Lieven Govaerts. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bufio"
	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"code.google.com/p/gopacket/tcpassembly"
	"code.google.com/p/gopacket/tcpassembly/tcpreader"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"
)

/* Command line arguments */
var iface = flag.String("ifce", "en0", "Interface to get packets from")
var inputfile = flag.String("infile", "", "read packets from file")
var logAllPackets = flag.Bool("v", false, "Logs every packet in great detail")
var launchCmd = flag.String("e", "", "Launches the command and logs its traffic")
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

type interval struct {
	req  *http.Request
	resp *http.Response
}

type BidiStream struct {
	key     uint64
	in, out *TCPStream
}

// TCPStream will handle the actual decoding of http requests and responses.
type TCPStream struct {
	netFlow, tcpFlow gopacket.Flow
	readStream       tcpreader.ReaderStream
	storage          *Storage
	bidikey          uint64
	closed           bool
}

/* This reads both HTTP requests and HTTP responses in two separate streams */
func (h *TCPStream) runOut() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64

	for {
		/*		_, err := buf.Peek(1)
				if err == io.EOF {
					return
				}*/
		req, err := http.ReadRequest(buf)
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)

			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			// bodyBytes := tcpreader.DiscardBytesToEOF(req.Body)
			req.Body.Close()
			err = h.storage.SentRequest(h.bidikey, reqID, time.Now(), req)
			if err != nil {
				log.Println("Error storing request", err)
			}

			reqID++
			//			fmt.Print(".")
			// log.Println("Received request from stream", h.netFlow, h.tcpFlow,
			// 	":", req, "with", bodyBytes, "bytes in request body")
		}
	}
}

func (h *TCPStream) runIn() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64
	for {
		/* Don't start reading a response if no data is available */
		_, err := buf.Peek(1)
		if err == io.EOF {
			return
		}
		/* Data available, read response */
		resp, err := http.ReadResponse(buf, nil) // TODO: request
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			//			log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)
			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			_, err := tcpreader.DiscardBytesToFirstError(resp.Body)
			if err != nil && err != io.EOF {
				log.Println("Error discarding bytes ", err)
			}
			resp.Body.Close()
			err = h.storage.ReceivedResponse(h.bidikey, reqID, time.Now(), resp)
			if err != nil {
				log.Println("Error storing response", err)
			}
			reqID++
			// fmt.Print(".")
			//log.Println("Received response from stream", h.netFlow, h.tcpFlow,
			//	":", resp, "with", bodyBytes, "bytes in response body")
		}
		/* Match the response with the next request */
	}

}

/* httpStreamFactory implements tcpassembly.StreamFactory */
type httpStreamFactory struct {
	bidiStreams map[uint64]*BidiStream
	storage     *Storage
	closed      bool
}

func NewStreamFactory(s *Storage) *httpStreamFactory {
	return &httpStreamFactory{bidiStreams: make(map[uint64]*BidiStream),
		storage: s}
}
func (h *httpStreamFactory) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {

	/* Watch out: this function can still get called even after all
	   streams were flushed (via FlushAll) and closed. */
	/*	if h.closed == true {
			return tcpreader.NewReaderStream()
		}
	*/
	/* First the outgoing stream, then the incoming stream */
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	hstream := &TCPStream{
		netFlow:    netFlow,
		tcpFlow:    tcpFlow,
		readStream: tcpreader.NewReaderStream(),
		storage:    h.storage,
		bidikey:    key,
	}

	bds := h.bidiStreams[key]
	if bds == nil {
		//		log.Println("reading stream", netFlow, tcpFlow)
		bds = &BidiStream{out: hstream, key: key}
		h.bidiStreams[key] = bds
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runOut()
	} else {
		//		log.Println("opening TCP conn", netFlow, tcpFlow)
		bds.in = hstream
		err := h.storage.OpenTCPConnection(key, time.Now())
		if err != nil {
			log.Println("Error storing connection", err)
		}
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runIn()
	}

	// ReaderStream implements tcpassembly.Stream, so we can return a pointer to it.
	return &hstream.readStream
}

func (h *httpStreamFactory) LogPacketSize(packet gopacket.Packet) {
	netFlow := packet.NetworkLayer().NetworkFlow()
	tcpFlow := packet.TransportLayer().TransportFlow()
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	ipv4Layer := packet.Layer(layers.LayerTypeIPv4)
	ipv4, _ := ipv4Layer.(*layers.IPv4)

	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	tcp, _ := tcpLayer.(*layers.TCP)

	bds := h.bidiStreams[key]
	if bds == nil || bds.in == nil || bds.out == nil {
		return
	}

	payloadLength := uint32(ipv4.Length - uint16(ipv4.IHL)*4 - uint16(tcp.DataOffset)*4)

	if bds.in.netFlow == netFlow {
		/* incoming */
		err := h.storage.IncomingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	} else {
		/* outgoing */
		err := h.storage.OutgoingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	}
}

func create_process_ended_channel(cmd *exec.Cmd) (cmd_done chan error) {
	cmd_done = make(chan error, 1)
	go func() {
		cmd_done <- cmd.Wait()
	}()
	return cmd_done
}

/* Create a channel and send any CTRL-C signal over it */
func create_ctrl_c_handler() (timeout chan os.Signal) {
	ctrlc := make(chan os.Signal, 1)
	signal.Notify(ctrlc, os.Interrupt)
	return ctrlc
}

/* Create a channel that will contain TRUE after T seconds. */
func create_timeout_channel(t time.Duration) (timeout chan bool) {
	timeout = make(chan bool, 1)
	go func() {
		time.Sleep(t * time.Second)
		timeout <- true
	}()
	return timeout
}

func main() {
	flag.Parse()
	log.SetFlags(log.Ltime | log.Lmicroseconds)
	//	log.SetOutput(ioutil.Discard)

	// Setup profiler
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// Set up storage layer
	storage, err := NewStorage()
	if err != nil {
		panic(err)
	}

	// Set up assembly
	streamFactory := NewStreamFactory(storage)
	streamPool := tcpassembly.NewStreamPool(streamFactory)
	assembler := tcpassembly.NewAssembler(streamPool)

	// Setup CTRL-C handler channel
	ctrlc := create_ctrl_c_handler()

	// Setup channel that reports all socket kernel events (Mac OS X only)
	//	netDescSource := NewOSXNetDescSource()
	//	descriptors := netDescSource.Descriptors()
	var handle *pcap.Handle
	if *inputfile != "" {
		handle, err = pcap.OpenOffline(*inputfile)
	} else {
		log.Printf("starting capture on interface %q", *iface)
		// Setup packet capture
		handle, err = pcap.OpenLive(*iface, 1600, true, pcap.BlockForever)
	}

	if err != nil {
		panic(err)
	} else if err := handle.SetBPFFilter("tcp and port 80"); err != nil {
		panic(err)
	}
	log.Println("reading in packets. Press CTRL-C to end and report.")
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets := packetSource.Packets()

	// Run the external command
	//pid := uint32(0)
	var cmd_done chan error
	var start_time time.Time
	if *launchCmd != "" {
		s := *launchCmd
		//		n := strings.Index(*launchCmd, " ")
		//		fmt.Println("Launching ", s[0:n], " --- ", s[n:])
		//		cmd := exec.Command(s[0:n], strings.Split(s[n:], " "))
		args := strings.Split(s, " ")
		cmd := exec.Command(args[0], args[1:]...)
		start_time = time.Now()
		cmd_stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		err = cmd.Start()
		if err != nil {
			panic(err)
		}
		go io.Copy(os.Stdout, cmd_stdout)
		//		cmd.Stdout = os.Stdout
		//		cmd.Stderr = os.Stderr
		// Create the channel that listens for the end of the command.
		cmd_done = create_process_ended_channel(cmd)
		//pid = uint32(cmd.Process.Pid)
	}

	var timeout chan bool
loop:
	for {
		select {
		/*		case netDesc := <-descriptors:
				if netDesc.Pid == pid {
					fmt.Println("event received ", netDesc)
				}*/
		case packet, ok := <-packets:
			if !ok {
				log.Println("All data read")
				packets = nil
				timeout = create_timeout_channel(0)
				break
			}

			if packet == nil {
				break
			}

			if *logAllPackets {
				log.Println(packet)
			}

			if storage.PacketInScope(packet) {
				streamFactory.LogPacketSize(packet)
				netFlow := packet.NetworkLayer().NetworkFlow()
				tcp := packet.TransportLayer().(*layers.TCP)

				assembler.AssembleWithTimestamp(netFlow, tcp,
					packet.Metadata().Timestamp)
			}
		case err := <-cmd_done:
			if err != nil {
				log.Printf("process done with error = %v\n", err)
			}

			log.Println("Process took: ", time.Now().Sub(start_time))

			/* Wait for a couple of seconds, just enough to get the events
			   handled by the main function. */
			log.Println("Waiting for the remaining responses to arrive.")
			timeout = create_timeout_channel(10)
		case <-ctrlc:
			/* Don't wait. */
			timeout = create_timeout_channel(0)
		case <-timeout:
			break loop
		}
	}

	signal.Stop(ctrlc)

	/* Cleanup the go routines */
	/* Ignore any http request/response parsing errors when closing the streams. */
	streamFactory.closed = true
	for _, v := range streamFactory.bidiStreams {
		if v.in != nil {
			v.in.closed = true
		}
		if v.out != nil {
			v.out.closed = true
		}
	}

	assembler.FlushAll()

	/* Close the storage layer. This will block until all pending inserts in
	   the db are handled. */
	storage.Close()

	reporting, err := NewReporting()
	if err != nil {
		panic(err)
	}

	if err = reporting.Report(); err != nil {
		log.Panic(err)
	}

	os.Exit(0)

}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/hadalhw17/AlexBot/commands"
)

func init() {
	flag.StringVar(&token, "t", "", "Bot Token")
	flag.Parse()
	dontDie()
}

const ANNOUNCEID = "238048572413575168"

var token string

func main() {
	if token == "" {
		fmt.Println("No token provided. Please run: AlexBot -t <bot token>")
		return
	}

	// Create a new Discord session using the provided bot token.
	dg, err := discordgo.New("Bot " + token)
	if err != nil {
		fmt.Println("Error creating Discord session: ", err)
		return
	}
	//dg.GuildMemberRoleRemove("238048572413575168","299602264039882752","295907725022593024")

	// Register ready as a callback for the ready events.
	dg.AddHandler(ready)

	// Register messageCreate as a callback for the messageCreate events.
	dg.AddHandler(messageCreate)

	// Register guildCreate as a callback for the guildCreate events.
	dg.AddHandler(guildCreate)

	//Register newMember as a callback for the GuildMemberAdd
	dg.AddHandler(newMember)
	//
	// Open the websocket and begin listening.
	err = dg.Open()
	if err != nil {
		fmt.Println("Error opening Discord session: ", err)
	}
	//Every 7 days, the bot will send a reminder about our pending coursework.
	ticker := time.NewTicker(time.Day * 7)
	var msg string
	var num int
	msg, num = commands.ReadCal("courseworks.ics")
	go func() {
		for range ticker.C {
			if num != 0 {
				dg.ChannelMessageSend(ANNOUNCEID, "you have a lot of courseworks to do!"+"\n"+msg)
			} else {
				dg.ChannelMessageSend(ANNOUNCEID, " no courseworks!! You are free!")
			}
		}

	}()
	ticker1 := time.NewTicker(time.Minute * 10)
	go func() {
		for range ticker1.C {
			fmt.Println("...")
		}
	}()
	fmt.Print("Bot is running.... Thanks god!!")
	// Simple way to keep program running until CTRL-C is pressed.
	<-make(chan struct{})
	return
}

func ready(s *discordgo.Session, event *discordgo.Ready) {
	// Set the playing status.
	_ = s.UpdateStatus(0, "@AlexBot !info|HadalHW17")

}

// This function will be called (due to AddHandler above) every time a new
// message is created on any channel that the autenticated bot has access to.
func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	if m.Author.ID == s.State.User.ID {
		// Ignore self
		return
	}

	reply, isExist := commands.ParceForCommands(s, m)
	if isExist {
		s.ChannelMessageSend(m.ChannelID, reply)
	}
	return
}

// This function will be called (due to AddHandler above) every time a new
// guild is joined.+
func guildCreate(s *discordgo.Session, event *discordgo.GuildCreate) {
	if event.Guild.Unavailable {
		return
	}

	for _, channel := range event.Guild.Channels {
		if channel.ID == event.Guild.ID {
			_, _ = s.ChannelMessageSend(ANNOUNCEID, "AlexBot is ready! Type !help to get some usefull shit about me.")
			return
		}
	}
}

// This function will be called (due to AddHandler above) every time a new
// member is joined.+
func newMember(s *discordgo.Session, event *discordgo.GuildMemberAdd) {
	s.ChannelMessageSend(ANNOUNCEID, "New member joined. Welcome @"+event.Member.User.Username)

}

func dontDie() {
	ticker := time.NewTicker(time.Minute * 20)
	go func() {
		for range ticker.C {
			fmt.Println("...")
		}
	}()
}
<commit_msg>I whoopsed. time.Day isn't a thing.<commit_after>package main

import (
	"flag"
	"fmt"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/hadalhw17/AlexBot/commands"
)

func init() {
	flag.StringVar(&token, "t", "", "Bot Token")
	flag.Parse()
	dontDie()
}

const ANNOUNCEID = "238048572413575168"

var token string

func main() {
	if token == "" {
		fmt.Println("No token provided. Please run: AlexBot -t <bot token>")
		return
	}

	// Create a new Discord session using the provided bot token.
	dg, err := discordgo.New("Bot " + token)
	if err != nil {
		fmt.Println("Error creating Discord session: ", err)
		return
	}
	//dg.GuildMemberRoleRemove("238048572413575168","299602264039882752","295907725022593024")

	// Register ready as a callback for the ready events.
	dg.AddHandler(ready)

	// Register messageCreate as a callback for the messageCreate events.
	dg.AddHandler(messageCreate)

	// Register guildCreate as a callback for the guildCreate events.
	dg.AddHandler(guildCreate)

	//Register newMember as a callback for the GuildMemberAdd
	dg.AddHandler(newMember)
	//
	// Open the websocket and begin listening.
	err = dg.Open()
	if err != nil {
		fmt.Println("Error opening Discord session: ", err)
	}
	//Every 7 days, the bot will send a reminder about our pending coursework.
	ticker := time.NewTicker(time.Hour * 168)
	var msg string
	var num int
	msg, num = commands.ReadCal("courseworks.ics")
	go func() {
		for range ticker.C {
			if num != 0 {
				dg.ChannelMessageSend(ANNOUNCEID, "you have a lot of courseworks to do!"+"\n"+msg)
			} else {
				dg.ChannelMessageSend(ANNOUNCEID, " no courseworks!! You are free!")
			}
		}

	}()
	ticker1 := time.NewTicker(time.Minute * 10)
	go func() {
		for range ticker1.C {
			fmt.Println("...")
		}
	}()
	fmt.Print("Bot is running.... Thanks god!!")
	// Simple way to keep program running until CTRL-C is pressed.
	<-make(chan struct{})
	return
}

func ready(s *discordgo.Session, event *discordgo.Ready) {
	// Set the playing status.
	_ = s.UpdateStatus(0, "@AlexBot !info|HadalHW17")

}

// This function will be called (due to AddHandler above) every time a new
// message is created on any channel that the autenticated bot has access to.
func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	if m.Author.ID == s.State.User.ID {
		// Ignore self
		return
	}

	reply, isExist := commands.ParceForCommands(s, m)
	if isExist {
		s.ChannelMessageSend(m.ChannelID, reply)
	}
	return
}

// This function will be called (due to AddHandler above) every time a new
// guild is joined.+
func guildCreate(s *discordgo.Session, event *discordgo.GuildCreate) {
	if event.Guild.Unavailable {
		return
	}

	for _, channel := range event.Guild.Channels {
		if channel.ID == event.Guild.ID {
			_, _ = s.ChannelMessageSend(ANNOUNCEID, "AlexBot is ready! Type !help to get some usefull shit about me.")
			return
		}
	}
}

// This function will be called (due to AddHandler above) every time a new
// member is joined.+
func newMember(s *discordgo.Session, event *discordgo.GuildMemberAdd) {
	s.ChannelMessageSend(ANNOUNCEID, "New member joined. Welcome @"+event.Member.User.Username)

}

func dontDie() {
	ticker := time.NewTicker(time.Minute * 20)
	go func() {
		for range ticker.C {
			fmt.Println("...")
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"github.com/badoux/checkmail"
	"github.com/gorilla/mux"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
)

// From config.json
type Config struct {
	Protocol string `json:"protocol"`
	ApiKey   string `json:"api-key"`
	Url      string `json:"url"`
}

// Send message
type Message struct {
	Status string `json:"status"`
	Body   string `json:"message"`
}

func main() {
	r := mux.NewRouter()

	// ALL /
	// Sends a Hello, Crema
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		sendJson(w, r, Message{
			Status: "ok",
			Body:   "Hello, Crema!",
		})
	})

	// GET /{email}
	// Both test validates the email and reminds user to POST to /{email}
	r.HandleFunc("/{email}", func(w http.ResponseWriter, r *http.Request) {
		email := mux.Vars(r)["email"]

		// Validates email and sends an error message if it is not validated
		if err := checkmail.ValidateFormat(email); err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Email in incorrect format.",
			})
			return
		}

		// If the email is validated, a reminder for the user to POST to the account
		sendJson(w, r, Message{
			Status: "ok",
			Body:   "Thanks for using our service, " + email + ". Please make sure your form has the method=POST attribute",
		})
	}).Methods("GET")

	// POST /{email}
	// Main function of crema forms - validates the email, executes the template, sends template through email
	r.HandleFunc("/{email}", func(w http.ResponseWriter, r *http.Request) {
		email := mux.Vars(r)["email"]

		// Validates email and sens an error message if it is not validated
		if err := checkmail.ValidateFormat(email); err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Email in incorrect format.",
			})
			return
		}

		// Parse the form and execute it into the email template
		r.ParseMultipartForm((1 << 10) * 24)
		var doc bytes.Buffer
		t, err := template.ParseFiles("email.tmpl")
		if err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Could not parse template..",
			})
		}
		t.Execute(&doc, r.Form)
		html := doc.String()

		config := Config{
			Protocol: os.Getenv("PROTOCOL"),
			ApiKey:   os.Getenv("APIKEY"),
			Url:      os.Getenv("URL"),
		}
		// Post to Mailchimp
		resp, err := http.PostForm(config.Protocol+config.ApiKey+"@"+config.Url, url.Values{
			"from":    {"crema@koostudios.com"},
			"to":      {email},
			"subject": {"New submission"},
			"html":    {html},
		})
		if err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Could not send email.",
			})
		}
		defer resp.Body.Close()

		// Read response from Mailchimp
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "There was an error sending the email.",
			})
			return
		}

		// Sends response from Mailchimp
		sendJson(w, r, Message{
			Status: "ok",
			Body:   string(body),
		})
	})

	// Starts the Application
	port := os.Getenv("PORT")
	if port == "" {
		log.Fatal("$PORT must be set")
	}
	log.Fatal(http.ListenAndServe(":"+port, r))
}

// Function sendJson sets the content type and relevant headers before sending a Message
func sendJson(w http.ResponseWriter, r *http.Request, msg Message) {
	w.Header().Set("Content-Type", "application/json;charset=UTF-8")
	w.Header().Set("Access-Control-Allow-Credentials", "true")

	// Set Access-Control-Allow-Origin
	if origin := r.Header.Get("Origin"); origin != "" {
		w.Header().Set("Access-Control-Allow-Origin", origin)
	}

	// Exposes the required headers
	if expose := r.URL.Query()["access-control-expose-headers"]; len(expose) != 0 {
		w.Header()["Access-Control-Expose-Headers"] = expose
		for key := range expose {
			w.Header()[expose[key]] = r.URL.Query()["__amp_source_origin"]
		}
	}

	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(msg)
}
<commit_msg>Added some log prints<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/badoux/checkmail"
	"github.com/gorilla/mux"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
)

// From config.json
type Config struct {
	Protocol string `json:"protocol"`
	ApiKey   string `json:"api-key"`
	Url      string `json:"url"`
}

// Send message
type Message struct {
	Status string `json:"status"`
	Body   string `json:"message"`
}

func main() {
	r := mux.NewRouter()

	// ALL /
	// Sends a Hello, Crema
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		sendJson(w, r, Message{
			Status: "ok",
			Body:   "Hello, Crema!",
		})
	})

	// GET /{email}
	// Both test validates the email and reminds user to POST to /{email}
	r.HandleFunc("/{email}", func(w http.ResponseWriter, r *http.Request) {
		email := mux.Vars(r)["email"]

		// Validates email and sends an error message if it is not validated
		if err := checkmail.ValidateFormat(email); err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Email in incorrect format.",
			})
			return
		}

		// If the email is validated, a reminder for the user to POST to the account
		sendJson(w, r, Message{
			Status: "ok",
			Body:   "Thanks for using our service, " + email + ". Please make sure your form has the method=POST attribute",
		})
	}).Methods("GET")

	// POST /{email}
	// Main function of crema forms - validates the email, executes the template, sends template through email
	r.HandleFunc("/{email}", func(w http.ResponseWriter, r *http.Request) {
		email := mux.Vars(r)["email"]

		// Validates email and sens an error message if it is not validated
		if err := checkmail.ValidateFormat(email); err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Email in incorrect format.",
			})
			return
		}

		// Parse the form and execute it into the email template
		r.ParseMultipartForm((1 << 10) * 24)
		var doc bytes.Buffer
		t, err := template.ParseFiles("email.tmpl")
		if err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Could not parse template..",
			})
		}
		t.Execute(&doc, r.Form)
		html := doc.String()

		config := Config{
			Protocol: os.Getenv("PROTOCOL"),
			ApiKey:   os.Getenv("APIKEY"),
			Url:      os.Getenv("URL"),
		}

		fmt.Println(config)
		// Post to Mailchimp
		resp, err := http.PostForm(config.Protocol+config.ApiKey+"@"+config.Url, url.Values{
			"from":    {"crema@koostudios.com"},
			"to":      {email},
			"subject": {"New submission"},
			"html":    {html},
		})
		if err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "Could not send email.",
			})
		}
		defer resp.Body.Close()

		// Read response from Mailchimp
		body, err := ioutil.ReadAll(resp.Body)
		fmt.Println(body)
		if err != nil {
			sendJson(w, r, Message{
				Status: "error",
				Body:   "There was an error sending the email.",
			})
			return
		}

		// Sends response from Mailchimp
		sendJson(w, r, Message{
			Status: "ok",
			Body:   string(body),
		})
	})

	// Starts the Application
	port := os.Getenv("PORT")
	if port == "" {
		log.Fatal("$PORT must be set")
	}
	log.Fatal(http.ListenAndServe(":"+port, r))
}

// Function sendJson sets the content type and relevant headers before sending a Message
func sendJson(w http.ResponseWriter, r *http.Request, msg Message) {
	w.Header().Set("Content-Type", "application/json;charset=UTF-8")
	w.Header().Set("Access-Control-Allow-Credentials", "true")

	// Set Access-Control-Allow-Origin
	if origin := r.Header.Get("Origin"); origin != "" {
		w.Header().Set("Access-Control-Allow-Origin", origin)
	}

	// Exposes the required headers
	if expose := r.URL.Query()["access-control-expose-headers"]; len(expose) != 0 {
		w.Header()["Access-Control-Expose-Headers"] = expose
		for key := range expose {
			w.Header()[expose[key]] = r.URL.Query()["__amp_source_origin"]
		}
	}

	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(msg)
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"github.com/albertyw/reaction-pics/server"
	"github.com/joho/godotenv"
	newrelic "github.com/newrelic/go-agent"
	"github.com/rollbar/rollbar-go"
)

func setupEnv() {
	err := godotenv.Load()
	if err != nil {
		panic(err)
	}
}

func getNewRelicApp() newrelic.Application {
	newrelicKey := os.Getenv("NEWRELIC_KEY")
	config := newrelic.NewConfig("Reaction.pics", newrelicKey)
	app, err := newrelic.NewApplication(config)
	if err != nil {
		panic(err)
	}
	return app
}

func setupRollbar() {
	rollbar.SetToken(os.Getenv("ROLLBAR_SERVER_TOKEN"))
	rollbar.SetEnvironment(os.Getenv("ENVIRONMENT"))
}

func main() {
	setupEnv()
	newrelicApp := getNewRelicApp()
	server.Run(newrelicApp)
}
<commit_msg>Make sure rollbar is setup correctly<commit_after>package main

import (
	"os"

	"github.com/albertyw/reaction-pics/server"
	"github.com/joho/godotenv"
	newrelic "github.com/newrelic/go-agent"
	"github.com/rollbar/rollbar-go"
)

func setupEnv() {
	err := godotenv.Load()
	if err != nil {
		panic(err)
	}
}

func getNewRelicApp() newrelic.Application {
	newrelicKey := os.Getenv("NEWRELIC_KEY")
	config := newrelic.NewConfig("Reaction.pics", newrelicKey)
	app, err := newrelic.NewApplication(config)
	if err != nil {
		panic(err)
	}
	return app
}

func setupRollbar() {
	rollbar.SetToken(os.Getenv("ROLLBAR_SERVER_TOKEN"))
	rollbar.SetEnvironment(os.Getenv("ENVIRONMENT"))
}

func main() {
	setupEnv()
	setupRollbar()
	newrelicApp := getNewRelicApp()
	server.Run(newrelicApp)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/carbocation/forum.git/forum"
	"github.com/carbocation/util.git/datatypes/binarytree"
	"github.com/carbocation/util.git/datatypes/closuretable"
	"net/http"
	"strconv"
	"time"
	//"html/template"
	"database/sql"
	_ "github.com/lib/pq"
)

var db *sql.DB

func main() {
	// Initialize the DB and keep here so that it will presumably kept open in a pool
	db = initdb()
	defer db.Close()

	http.HandleFunc("/thread/", threadHandler)
	http.HandleFunc("/hello/", commentHandler)
	http.HandleFunc("/css/", cssHandler)
	http.HandleFunc("/", defaultHandler)
	http.ListenAndServe("localhost:9999", nil)
}

func initdb() *sql.DB {
	db, err := sql.Open("postgres", "dbname=forumtest sslmode=disable")
	if err != nil {
		panic(err)
	}

	return db
}

func defaultHandler(w http.ResponseWriter, r *http.Request) {
	remPartOfURL := r.URL.Path[len("/"):]
	fmt.Fprintf(w, "<link rel=\"stylesheet\" href=\"/css/main.css\">
Welcome, %s
Say hello", remPartOfURL)

	fmt.Fprint(w, "")
}

func cssHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/css")

	docname := r.URL.Path[len("/css/"):]

	switch {
	case docname == "main.css":
		fmt.Fprintf(w, "%s", mainCss())
	}

}

func mainCss() string {
	return `
div .comment {
	padding-left: 100px;
}
`
}

func commentHandler(w http.ResponseWriter, r *http.Request) {
	remPartOfURL := r.URL.Path[len("/hello/"):] //get everything after the /hello/ part of the URL
	//w.Header().Set("Content-Type", "text/html")

	fmt.Fprint(w, "<link rel=\"stylesheet\" href=\"/css/main.css\">")
	fmt.Fprintf(w, "Hello %s!", remPartOfURL)

	PrintNestedComments(w, ClosureTree())

	fmt.Fprint(w, "")
}

func threadHandler(w http.ResponseWriter, r *http.Request) {
	unsafeId := r.URL.Path[len("/thread/"):]

	//If the thread ID is not parseable as an integer, stop immediately
	id, err := strconv.ParseInt(unsafeId, 10, 64)
	if err != nil {
		return
	}

	// Generate a closuretable from the root requested id
	ct := closuretable.New(id)
	// Pull down the remaining elements in the closure table that are descendants of this node
	q := `select *
from entry_closures
where descendant in (
select descendant
from entry_closures
where ancestor=$1
)
and ancestor in (
select descendant
from entry_closures
where ancestor=$1
)
and depth = 1`
	stmt, err := db.Prepare(q)
	if err != nil {
		//fmt.Printf("Statement Preparation Error: %s", err)
		return
	}

	rows, err := stmt.Query(unsafeId)
	if err != nil {
		//fmt.Printf("Query Error: %v", err)
		return
	}

	//Populate the closuretable
	for rows.Next() {
		var ancestor, descendant int64
		var depth int
		err = rows.Scan(&ancestor, &descendant, &depth)
		if err != nil {
			//fmt.Printf("Rowscan error: %s", err)
			return
		}

		err = ct.AddChild(closuretable.Child{Parent: ancestor, Child: descendant})

		//err = ct.AddRelationship(closuretable.Relationship{Ancestor: ancestor, Descendant: descendant, Depth: depth})
		if err != nil {
			//fmt.Fprintf(w, "Error: %s", err)
			return
		}
	}

	fmt.Printf("Closure table is: %#v", ct)

	id, entries, err := forum.RetrieveDescendantEntries(unsafeId, db)
	if err != nil {
		//fmt.Fprintf(w, "Error: %s", err)
		return
	}

	//fmt.Printf("Entries: %#v, %s", entries, err)

	//Obligatory boxing step
	interfaceEntries := map[int64]interface{}{}
	for k, v := range entries {
		interfaceEntries[k] = v
	}

	tree, err := ct.TableToTree(interfaceEntries)
	fmt.Printf("ClosureTree is : %#v", ct)
	if err != nil {
		fmt.Printf("TableToTree error: %s", err)
	}

	fmt.Fprint(w, "<link rel=\"stylesheet\" href=\"/css/main.css\">")
	fmt.Printf("Entry: %#v\n", tree)
	PrintNestedComments(w, tree)
	fmt.Fprint(w, "")
}

func PrintNestedComments(w http.ResponseWriter, el *binarytree.Tree) {
	if el == nil {
		return
	}

	fmt.Fprint(w, "
")
	//Self
	fmt.Printf("Entry: %#v\n", el)
	e := el.Value.(forum.Entry)
	fmt.Fprintf(w, "Title: %s", e.Title)

	//Children are nested
	PrintNestedComments(w, el.Left())
	fmt.Fprint(w, "
")

	//Siblings are parallel
	PrintNestedComments(w, el.Right())
}

func ClosureTree() *binarytree.Tree {
	//Make some entries
	entries := map[int64]forum.Entry{
		0: forum.Entry{Id: 100, Title: "Title 100", Body: "Body 100", Created: time.Now(), AuthorId: 0},
		1: forum.Entry{Id: 101, Title: "Title 101", Body: "Body 101", Created: time.Now(), AuthorId: 1},
		2: forum.Entry{Id: 102, Title: "Title 102", Body: "Body 102", Created: time.Now(), AuthorId: 2},
		3: forum.Entry{Id: 103, Title: "Title 103", Body: "Body 103", Created: time.Now(), AuthorId: 3},
	}

	ct := closuretable.New(0)
	ct.AddChild(closuretable.Child{Parent: 0, Child: 1})
	ct.AddChild(closuretable.Child{Parent: 0, Child: 2})
	ct.AddChild(closuretable.Child{Parent: 1, Child: 3})

	// Obligatory boxing step
	// Convert to interface type so the generic TableToTree method can be called on these entries
	boxedEntries := map[int64]interface{}{}
	for k, v := range entries {
		boxedEntries[k] = v
	}

	tree, _ := ct.TableToTree(boxedEntries)

	return tree
}
<commit_msg>Removed debugging code.<commit_after>package main

import (
	"fmt"
	"github.com/carbocation/forum.git/forum"
	"github.com/carbocation/util.git/datatypes/binarytree"
	"github.com/carbocation/util.git/datatypes/closuretable"
	"net/http"
	"strconv"
	"time"
	//"html/template"
	"database/sql"
	_ "github.com/lib/pq"
)

var db *sql.DB

func main() {
	// Initialize the DB and keep here so that it will presumably kept open in a pool
	db = initdb()
	defer db.Close()

	http.HandleFunc("/thread/", threadHandler)
	http.HandleFunc("/hello/", commentHandler)
	http.HandleFunc("/css/", cssHandler)
	http.HandleFunc("/", defaultHandler)
	http.ListenAndServe("localhost:9999", nil)
}

func initdb() *sql.DB {
	db, err := sql.Open("postgres", "dbname=forumtest sslmode=disable")
	if err != nil {
		panic(err)
	}

	return db
}

func defaultHandler(w http.ResponseWriter, r *http.Request) {
	remPartOfURL := r.URL.Path[len("/"):]
	fmt.Fprintf(w, "<link rel=\"stylesheet\" href=\"/css/main.css\">
Welcome, %s
Say hello", remPartOfURL)

	fmt.Fprint(w, "")
}

func cssHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/css")

	docname := r.URL.Path[len("/css/"):]

	switch {
	case docname == "main.css":
		fmt.Fprintf(w, "%s", mainCss())
	}

}

func mainCss() string {
	return `
div .comment {
	padding-left: 100px;
}
`
}

func commentHandler(w http.ResponseWriter, r *http.Request) {
	remPartOfURL := r.URL.Path[len("/hello/"):] //get everything after the /hello/ part of the URL
	//w.Header().Set("Content-Type", "text/html")

	fmt.Fprint(w, "<link rel=\"stylesheet\" href=\"/css/main.css\">")
	fmt.Fprintf(w, "Hello %s!", remPartOfURL)

	PrintNestedComments(w, ClosureTree())

	fmt.Fprint(w, "")
}

func threadHandler(w http.ResponseWriter, r *http.Request) {
	unsafeId := r.URL.Path[len("/thread/"):]

	//If the thread ID is not parseable as an integer, stop immediately
	id, err := strconv.ParseInt(unsafeId, 10, 64)
	if err != nil {
		return
	}

	// Generate a closuretable from the root requested id
	ct := closuretable.New(id)
	// Pull down the remaining elements in the closure table that are descendants of this node
	q := `select *
from entry_closures
where descendant in (
select descendant
from entry_closures
where ancestor=$1
)
and ancestor in (
select descendant
from entry_closures
where ancestor=$1
)
and depth = 1`
	stmt, err := db.Prepare(q)
	if err != nil {
		//fmt.Printf("Statement Preparation Error: %s", err)
		return
	}

	rows, err := stmt.Query(unsafeId)
	if err != nil {
		//fmt.Printf("Query Error: %v", err)
		return
	}

	//Populate the closuretable
	for rows.Next() {
		var ancestor, descendant int64
		var depth int
		err = rows.Scan(&ancestor, &descendant, &depth)
		if err != nil {
			//fmt.Printf("Rowscan error: %s", err)
			return
		}

		err = ct.AddChild(closuretable.Child{Parent: ancestor, Child: descendant})

		//err = ct.AddRelationship(closuretable.Relationship{Ancestor: ancestor, Descendant: descendant, Depth: depth})
		if err != nil {
			//fmt.Fprintf(w, "Error: %s", err)
			return
		}
	}

	id, entries, err := forum.RetrieveDescendantEntries(unsafeId, db)
	if err != nil {
		//fmt.Fprintf(w, "Error: %s", err)
		return
	}

	//fmt.Printf("Entries: %#v, %s", entries, err)

	//Obligatory boxing step
	interfaceEntries := map[int64]interface{}{}
	for k, v := range entries {
		interfaceEntries[k] = v
	}

	tree, err := ct.TableToTree(interfaceEntries)
	if err != nil {
		//fmt.Printf("TableToTree error: %s", err)
		return
	}

	fmt.Fprint(w, "<link rel=\"stylesheet\" href=\"/css/main.css\">")
	PrintNestedComments(w, tree)
	fmt.Fprint(w, "")
}

func PrintNestedComments(w http.ResponseWriter, el *binarytree.Tree) {
	if el == nil {
		return
	}

	fmt.Fprint(w, "
")
	//Self
	e := el.Value.(forum.Entry)
	fmt.Fprintf(w, "Title: %s", e.Title)

	//Children are nested
	PrintNestedComments(w, el.Left())
	fmt.Fprint(w, "
")

	//Siblings are parallel
	PrintNestedComments(w, el.Right())
}

func ClosureTree() *binarytree.Tree {
	//Make some entries
	entries := map[int64]forum.Entry{
		0: forum.Entry{Id: 100, Title: "Title 100", Body: "Body 100", Created: time.Now(), AuthorId: 0},
		1: forum.Entry{Id: 101, Title: "Title 101", Body: "Body 101", Created: time.Now(), AuthorId: 1},
		2: forum.Entry{Id: 102, Title: "Title 102", Body: "Body 102", Created: time.Now(), AuthorId: 2},
		3: forum.Entry{Id: 103, Title: "Title 103", Body: "Body 103", Created: time.Now(), AuthorId: 3},
	}

	ct := closuretable.New(0)
	ct.AddChild(closuretable.Child{Parent: 0, Child: 1})
	ct.AddChild(closuretable.Child{Parent: 0, Child: 2})
	ct.AddChild(closuretable.Child{Parent: 1, Child: 3})

	// Obligatory boxing step
	// Convert to interface type so the generic TableToTree method can be called on these entries
	boxedEntries := map[int64]interface{}{}
	for k, v := range entries {
		boxedEntries[k] = v
	}

	tree, _ := ct.TableToTree(boxedEntries)

	return tree
}
<|endoftext|>