<commit_before>// +build !plan9

package sftp

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/subtle"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/cmd/serve/proxy"
	"github.com/rclone/rclone/cmd/serve/proxy/proxyflags"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/config"
	"github.com/rclone/rclone/lib/env"
	"github.com/rclone/rclone/vfs"
	"github.com/rclone/rclone/vfs/vfsflags"
	"golang.org/x/crypto/ssh"
)

// server contains everything to run the server
type server struct {
	f        fs.Fs
	opt      Options
	vfs      *vfs.VFS
	ctx      context.Context // for global config
	config   *ssh.ServerConfig
	listener net.Listener
	waitChan chan struct{} // for waiting on the listener to close
	proxy    *proxy.Proxy
}

func newServer(ctx context.Context, f fs.Fs, opt *Options) *server {
	s := &server{
		f:        f,
		ctx:      ctx,
		opt:      *opt,
		waitChan: make(chan struct{}),
	}
	if proxyflags.Opt.AuthProxy != "" {
		s.proxy = proxy.New(ctx, &proxyflags.Opt)
	} else {
		s.vfs = vfs.New(f, &vfsflags.Opt)
	}
	return s
}

// getVFS gets the vfs from s or the proxy
func (s *server) getVFS(what string, sshConn *ssh.ServerConn) (VFS *vfs.VFS) {
	if s.proxy == nil {
		return s.vfs
	}
	if sshConn.Permissions == nil && sshConn.Permissions.Extensions == nil {
		fs.Infof(what, "SSH Permissions Extensions not found")
		return nil
	}
	key := sshConn.Permissions.Extensions["_vfsKey"]
	if key == "" {
		fs.Infof(what, "VFS key not found")
		return nil
	}
	VFS = s.proxy.Get(key)
	if VFS == nil {
		fs.Infof(what, "failed to read VFS from cache")
		return nil
	}
	return VFS
}

func (s *server) acceptConnections() {
	for {
		nConn, err := s.listener.Accept()
		if err != nil {
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			fs.Errorf(nil, "Failed to accept incoming connection: %v", err)
			continue
		}
		what := describeConn(nConn)

		// Before use, a handshake must be performed on the incoming net.Conn.
		sshConn, chans, reqs, err := ssh.NewServerConn(nConn, s.config)
		if err != nil {
			fs.Errorf(what, "SSH login failed: %v", err)
			continue
		}

		fs.Infof(what, "SSH login from %s using %s", sshConn.User(), sshConn.ClientVersion())

		// Discard all global out-of-band Requests
		go ssh.DiscardRequests(reqs)

		c := &conn{
			what: what,
			vfs:  s.getVFS(what, sshConn),
		}
		if c.vfs == nil {
			fs.Infof(what, "Closing unauthenticated connection (couldn't find VFS)")
			_ = nConn.Close()
			continue
		}
		c.handlers = newVFSHandler(c.vfs)

		// Accept all channels
		go c.handleChannels(chans)
	}
}

// Based on example server code from golang.org/x/crypto/ssh and server_standalone
func (s *server) serve() (err error) {
	var authorizedKeysMap map[string]struct{}

	// ensure the user isn't trying to use conflicting flags
	if proxyflags.Opt.AuthProxy != "" && s.opt.AuthorizedKeys != "" && s.opt.AuthorizedKeys != DefaultOpt.AuthorizedKeys {
		return errors.New("--auth-proxy and --authorized-keys cannot be used at the same time")
	}

	// Load the authorized keys
	if s.opt.AuthorizedKeys != "" && proxyflags.Opt.AuthProxy == "" {
		authKeysFile := env.ShellExpand(s.opt.AuthorizedKeys)
		authorizedKeysMap, err = loadAuthorizedKeys(authKeysFile)
		// If user set the flag away from the default then report an error
		if err != nil && s.opt.AuthorizedKeys != DefaultOpt.AuthorizedKeys {
			return err
		}
		fs.Logf(nil, "Loaded %d authorized keys from %q", len(authorizedKeysMap), authKeysFile)
	}

	if !s.opt.NoAuth && len(authorizedKeysMap) == 0 && s.opt.User == "" && s.opt.Pass == "" && s.proxy == nil {
		return errors.New("no authorization found, use --user/--pass or --authorized-keys or --no-auth or --auth-proxy")
	}

	// An SSH server is represented by a ServerConfig, which holds
	// certificate details and handles authentication of ServerConns.
	s.config = &ssh.ServerConfig{
		ServerVersion: "SSH-2.0-" + fs.GetConfig(s.ctx).UserAgent,
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			fs.Debugf(describeConn(c), "Password login attempt for %s", c.User())
			if s.proxy != nil {
				// query the proxy for the config
				_, vfsKey, err := s.proxy.Call(c.User(), string(pass), false)
				if err != nil {
					return nil, err
				}
				// just return the Key so we can get it back from the cache
				return &ssh.Permissions{
					Extensions: map[string]string{
						"_vfsKey": vfsKey,
					},
				}, nil
			} else if s.opt.User != "" && s.opt.Pass != "" {
				userOK := subtle.ConstantTimeCompare([]byte(c.User()), []byte(s.opt.User))
				passOK := subtle.ConstantTimeCompare(pass, []byte(s.opt.Pass))
				if (userOK & passOK) == 1 {
					return nil, nil
				}
			}
			return nil, fmt.Errorf("password rejected for %q", c.User())
		},
		PublicKeyCallback: func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			fs.Debugf(describeConn(c), "Public key login attempt for %s", c.User())
			if s.proxy != nil {
				//query the proxy for the config
				_, vfsKey, err := s.proxy.Call(
					c.User(),
					base64.StdEncoding.EncodeToString(pubKey.Marshal()),
					true,
				)
				if err != nil {
					return nil, err
				}
				// just return the Key so we can get it back from the cache
				return &ssh.Permissions{
					Extensions: map[string]string{
						"_vfsKey": vfsKey,
					},
				}, nil
			}
			if _, ok := authorizedKeysMap[string(pubKey.Marshal())]; ok {
				return &ssh.Permissions{
					// Record the public key used for authentication.
					Extensions: map[string]string{
						"pubkey-fp": ssh.FingerprintSHA256(pubKey),
					},
				}, nil
			}
			return nil, fmt.Errorf("unknown public key for %q", c.User())
		},
		AuthLogCallback: func(conn ssh.ConnMetadata, method string, err error) {
			status := "OK"
			if err != nil {
				status = err.Error()
			}
			fs.Debugf(describeConn(conn), "ssh auth %q from %q: %s", method, conn.ClientVersion(), status)
		},
		NoClientAuth: s.opt.NoAuth,
	}

	// Load the private key, from the cache if not explicitly configured
	keyPaths := s.opt.HostKeys
	cachePath := filepath.Join(config.CacheDir, "serve-sftp")
	if len(keyPaths) == 0 {
		keyPaths = []string{filepath.Join(cachePath, "id_rsa")}
	}
	for _, keyPath := range keyPaths {
		private, err := loadPrivateKey(keyPath)
		if err != nil && len(s.opt.HostKeys) == 0 {
			fs.Debugf(nil, "Failed to load %q: %v", keyPath, err)
			// If loading a cached key failed, make the keys and retry
			err = os.MkdirAll(cachePath, 0700)
			if err != nil {
				return errors.Wrap(err, "failed to create cache path")
			}
			const bits = 2048
			fs.Logf(nil, "Generating %d bit key pair at %q", bits, keyPath)
			err = makeSSHKeyPair(bits, keyPath+".pub", keyPath)
			if err != nil {
				return errors.Wrap(err, "failed to create SSH key pair")
			}
			// reload the new keys
			private, err = loadPrivateKey(keyPath)
		}
		if err != nil {
			return err
		}
		fs.Debugf(nil, "Loaded private key from %q", keyPath)

		s.config.AddHostKey(private)
	}

	// Once a ServerConfig has been configured, connections can be
	// accepted.
	s.listener, err = net.Listen("tcp", s.opt.ListenAddr)
	if err != nil {
		return errors.Wrap(err, "failed to listen for connection")
	}
	fs.Logf(nil, "SFTP server listening on %v\n", s.listener.Addr())

	go s.acceptConnections()

	return nil
}

// Addr returns the address the server is listening on
func (s *server) Addr() string {
	return s.listener.Addr().String()
}

// Serve runs the sftp server in the background.
//
// Use s.Close() and s.Wait() to shutdown server
func (s *server) Serve() error {
	err := s.serve()
	if err != nil {
		return err
	}
	return nil
}

// Wait blocks while the listener is open.
func (s *server) Wait() {
	<-s.waitChan
}

// Close shuts the running server down
func (s *server) Close() {
	err := s.listener.Close()
	if err != nil {
		fs.Errorf(nil, "Error on closing SFTP server: %v", err)
		return
	}
	close(s.waitChan)
}

func loadPrivateKey(keyPath string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load private key")
	}
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse private key")
	}
	return private, nil
}

// Public key authentication is done by comparing
// the public key of a received connection
// with the entries in the authorized_keys file.
func loadAuthorizedKeys(authorizedKeysPath string) (authorizedKeysMap map[string]struct{}, err error) {
	authorizedKeysBytes, err := ioutil.ReadFile(authorizedKeysPath)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load authorized keys")
	}
	authorizedKeysMap = make(map[string]struct{})
	for len(authorizedKeysBytes) > 0 {
		pubKey, _, _, rest, err := ssh.ParseAuthorizedKey(authorizedKeysBytes)
		if err != nil {
			return nil, errors.Wrap(err, "failed to parse authorized keys")
		}
		authorizedKeysMap[string(pubKey.Marshal())] = struct{}{}
		authorizedKeysBytes = bytes.TrimSpace(rest)
	}
	return authorizedKeysMap, nil
}

// makeSSHKeyPair make a pair of public and private keys for SSH access.
// Public key is encoded in the format for inclusion in an OpenSSH authorized_keys file.
// Private Key generated is PEM encoded
//
// Originally from: https://stackoverflow.com/a/34347463/164234
func makeSSHKeyPair(bits int, pubKeyPath, privateKeyPath string) (err error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}

	// generate and write private key as PEM
	privateKeyFile, err := os.OpenFile(privateKeyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer fs.CheckClose(privateKeyFile, &err)
	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	if err := pem.Encode(privateKeyFile, privateKeyPEM); err != nil {
		return err
	}

	// generate and write public key
	pub, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(pubKeyPath, ssh.MarshalAuthorizedKey(pub), 0644)
}
<commit_msg>serve sftp: fix authentication on one connection blocking others - fixes #4882<commit_after>// +build !plan9

package sftp

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/subtle"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/cmd/serve/proxy"
	"github.com/rclone/rclone/cmd/serve/proxy/proxyflags"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/config"
	"github.com/rclone/rclone/lib/env"
	"github.com/rclone/rclone/vfs"
	"github.com/rclone/rclone/vfs/vfsflags"
	"golang.org/x/crypto/ssh"
)

// server contains everything to run the server
type server struct {
	f        fs.Fs
	opt      Options
	vfs      *vfs.VFS
	ctx      context.Context // for global config
	config   *ssh.ServerConfig
	listener net.Listener
	waitChan chan struct{} // for waiting on the listener to close
	proxy    *proxy.Proxy
}

func newServer(ctx context.Context, f fs.Fs, opt *Options) *server {
	s := &server{
		f:        f,
		ctx:      ctx,
		opt:      *opt,
		waitChan: make(chan struct{}),
	}
	if proxyflags.Opt.AuthProxy != "" {
		s.proxy = proxy.New(ctx, &proxyflags.Opt)
	} else {
		s.vfs = vfs.New(f, &vfsflags.Opt)
	}
	return s
}

// getVFS gets the vfs from s or the proxy
func (s *server) getVFS(what string, sshConn *ssh.ServerConn) (VFS *vfs.VFS) {
	if s.proxy == nil {
		return s.vfs
	}
	if sshConn.Permissions == nil && sshConn.Permissions.Extensions == nil {
		fs.Infof(what, "SSH Permissions Extensions not found")
		return nil
	}
	key := sshConn.Permissions.Extensions["_vfsKey"]
	if key == "" {
		fs.Infof(what, "VFS key not found")
		return nil
	}
	VFS = s.proxy.Get(key)
	if VFS == nil {
		fs.Infof(what, "failed to read VFS from cache")
		return nil
	}
	return VFS
}

// Accept a single connection - run in a go routine as the ssh
// authentication can block
func (s *server) acceptConnection(nConn net.Conn) {
	what := describeConn(nConn)

	// Before use, a handshake must be performed on the incoming net.Conn.
	sshConn, chans, reqs, err := ssh.NewServerConn(nConn, s.config)
	if err != nil {
		fs.Errorf(what, "SSH login failed: %v", err)
		return
	}

	fs.Infof(what, "SSH login from %s using %s", sshConn.User(), sshConn.ClientVersion())

	// Discard all global out-of-band Requests
	go ssh.DiscardRequests(reqs)

	c := &conn{
		what: what,
		vfs:  s.getVFS(what, sshConn),
	}
	if c.vfs == nil {
		fs.Infof(what, "Closing unauthenticated connection (couldn't find VFS)")
		_ = nConn.Close()
		return
	}
	c.handlers = newVFSHandler(c.vfs)

	// Accept all channels
	go c.handleChannels(chans)
}

// Accept connections and call them in a go routine
func (s *server) acceptConnections() {
	for {
		nConn, err := s.listener.Accept()
		if err != nil {
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			fs.Errorf(nil, "Failed to accept incoming connection: %v", err)
			continue
		}
		go s.acceptConnection(nConn)
	}
}

// Based on example server code from golang.org/x/crypto/ssh and server_standalone
func (s *server) serve() (err error) {
	var authorizedKeysMap map[string]struct{}

	// ensure the user isn't trying to use conflicting flags
	if proxyflags.Opt.AuthProxy != "" && s.opt.AuthorizedKeys != "" && s.opt.AuthorizedKeys != DefaultOpt.AuthorizedKeys {
		return errors.New("--auth-proxy and --authorized-keys cannot be used at the same time")
	}

	// Load the authorized keys
	if s.opt.AuthorizedKeys != "" && proxyflags.Opt.AuthProxy == "" {
		authKeysFile := env.ShellExpand(s.opt.AuthorizedKeys)
		authorizedKeysMap, err = loadAuthorizedKeys(authKeysFile)
		// If user set the flag away from the default then report an error
		if err != nil && s.opt.AuthorizedKeys != DefaultOpt.AuthorizedKeys {
			return err
		}
		fs.Logf(nil, "Loaded %d authorized keys from %q", len(authorizedKeysMap), authKeysFile)
	}

	if !s.opt.NoAuth && len(authorizedKeysMap) == 0 && s.opt.User == "" && s.opt.Pass == "" && s.proxy == nil {
		return errors.New("no authorization found, use --user/--pass or --authorized-keys or --no-auth or --auth-proxy")
	}

	// An SSH server is represented by a ServerConfig, which holds
	// certificate details and handles authentication of ServerConns.
	s.config = &ssh.ServerConfig{
		ServerVersion: "SSH-2.0-" + fs.GetConfig(s.ctx).UserAgent,
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			fs.Debugf(describeConn(c), "Password login attempt for %s", c.User())
			if s.proxy != nil {
				// query the proxy for the config
				_, vfsKey, err := s.proxy.Call(c.User(), string(pass), false)
				if err != nil {
					return nil, err
				}
				// just return the Key so we can get it back from the cache
				return &ssh.Permissions{
					Extensions: map[string]string{
						"_vfsKey": vfsKey,
					},
				}, nil
			} else if s.opt.User != "" && s.opt.Pass != "" {
				userOK := subtle.ConstantTimeCompare([]byte(c.User()), []byte(s.opt.User))
				passOK := subtle.ConstantTimeCompare(pass, []byte(s.opt.Pass))
				if (userOK & passOK) == 1 {
					return nil, nil
				}
			}
			return nil, fmt.Errorf("password rejected for %q", c.User())
		},
		PublicKeyCallback: func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			fs.Debugf(describeConn(c), "Public key login attempt for %s", c.User())
			if s.proxy != nil {
				//query the proxy for the config
				_, vfsKey, err := s.proxy.Call(
					c.User(),
					base64.StdEncoding.EncodeToString(pubKey.Marshal()),
					true,
				)
				if err != nil {
					return nil, err
				}
				// just return the Key so we can get it back from the cache
				return &ssh.Permissions{
					Extensions: map[string]string{
						"_vfsKey": vfsKey,
					},
				}, nil
			}
			if _, ok := authorizedKeysMap[string(pubKey.Marshal())]; ok {
				return &ssh.Permissions{
					// Record the public key used for authentication.
					Extensions: map[string]string{
						"pubkey-fp": ssh.FingerprintSHA256(pubKey),
					},
				}, nil
			}
			return nil, fmt.Errorf("unknown public key for %q", c.User())
		},
		AuthLogCallback: func(conn ssh.ConnMetadata, method string, err error) {
			status := "OK"
			if err != nil {
				status = err.Error()
			}
			fs.Debugf(describeConn(conn), "ssh auth %q from %q: %s", method, conn.ClientVersion(), status)
		},
		NoClientAuth: s.opt.NoAuth,
	}

	// Load the private key, from the cache if not explicitly configured
	keyPaths := s.opt.HostKeys
	cachePath := filepath.Join(config.CacheDir, "serve-sftp")
	if len(keyPaths) == 0 {
		keyPaths = []string{filepath.Join(cachePath, "id_rsa")}
	}
	for _, keyPath := range keyPaths {
		private, err := loadPrivateKey(keyPath)
		if err != nil && len(s.opt.HostKeys) == 0 {
			fs.Debugf(nil, "Failed to load %q: %v", keyPath, err)
			// If loading a cached key failed, make the keys and retry
			err = os.MkdirAll(cachePath, 0700)
			if err != nil {
				return errors.Wrap(err, "failed to create cache path")
			}
			const bits = 2048
			fs.Logf(nil, "Generating %d bit key pair at %q", bits, keyPath)
			err = makeSSHKeyPair(bits, keyPath+".pub", keyPath)
			if err != nil {
				return errors.Wrap(err, "failed to create SSH key pair")
			}
			// reload the new keys
			private, err = loadPrivateKey(keyPath)
		}
		if err != nil {
			return err
		}
		fs.Debugf(nil, "Loaded private key from %q", keyPath)

		s.config.AddHostKey(private)
	}

	// Once a ServerConfig has been configured, connections can be
	// accepted.
	s.listener, err = net.Listen("tcp", s.opt.ListenAddr)
	if err != nil {
		return errors.Wrap(err, "failed to listen for connection")
	}
	fs.Logf(nil, "SFTP server listening on %v\n", s.listener.Addr())

	go s.acceptConnections()

	return nil
}

// Addr returns the address the server is listening on
func (s *server) Addr() string {
	return s.listener.Addr().String()
}

// Serve runs the sftp server in the background.
//
// Use s.Close() and s.Wait() to shutdown server
func (s *server) Serve() error {
	err := s.serve()
	if err != nil {
		return err
	}
	return nil
}

// Wait blocks while the listener is open.
func (s *server) Wait() {
	<-s.waitChan
}

// Close shuts the running server down
func (s *server) Close() {
	err := s.listener.Close()
	if err != nil {
		fs.Errorf(nil, "Error on closing SFTP server: %v", err)
		return
	}
	close(s.waitChan)
}

func loadPrivateKey(keyPath string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load private key")
	}
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse private key")
	}
	return private, nil
}

// Public key authentication is done by comparing
// the public key of a received connection
// with the entries in the authorized_keys file.
func loadAuthorizedKeys(authorizedKeysPath string) (authorizedKeysMap map[string]struct{}, err error) {
	authorizedKeysBytes, err := ioutil.ReadFile(authorizedKeysPath)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load authorized keys")
	}
	authorizedKeysMap = make(map[string]struct{})
	for len(authorizedKeysBytes) > 0 {
		pubKey, _, _, rest, err := ssh.ParseAuthorizedKey(authorizedKeysBytes)
		if err != nil {
			return nil, errors.Wrap(err, "failed to parse authorized keys")
		}
		authorizedKeysMap[string(pubKey.Marshal())] = struct{}{}
		authorizedKeysBytes = bytes.TrimSpace(rest)
	}
	return authorizedKeysMap, nil
}

// makeSSHKeyPair make a pair of public and private keys for SSH access.
// Public key is encoded in the format for inclusion in an OpenSSH authorized_keys file.
// Private Key generated is PEM encoded
//
// Originally from: https://stackoverflow.com/a/34347463/164234
func makeSSHKeyPair(bits int, pubKeyPath, privateKeyPath string) (err error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}

	// generate and write private key as PEM
	privateKeyFile, err := os.OpenFile(privateKeyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer fs.CheckClose(privateKeyFile, &err)
	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	if err := pem.Encode(privateKeyFile, privateKeyPEM); err != nil {
		return err
	}

	// generate and write public key
	pub, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(pubKeyPath, ssh.MarshalAuthorizedKey(pub), 0644)
}
<|endoftext|>


<commit_before>package protocols

import (
	"context"
	"encoding/binary"
	"net"

	"go.uber.org/zap"
)

type bittorrentMsg struct {
	Length             uint8
	ProtocolIdentifier [19]uint8
	Reserved           [8]uint8
	InfoHash           [20]uint8
	PeerID             [20]uint8
}

// HandleBittorrent handles a Bittorrent connection
func HandleBittorrent(ctx context.Context, conn net.Conn, logger Logger, h Honeypot) error {
	var err error
	defer func() {
		if err = conn.Close(); err != nil {
			logger.Error("failed to close connection", zap.Error(err), zap.String("handler", "bittorrent"))
			return
		}
	}()
	for {
		msg := bittorrentMsg{}
		if err := binary.Read(conn, binary.BigEndian, &msg); err != nil {
			logger.Error("failed to read message", zap.Error(err), zap.String("handler", "bittorrent"))
			break
		}
		md, err := h.MetadataByConnection(conn)
		if err != nil {
			return err
		}
		if err = h.Produce(conn, md, []byte{}); err != nil {
			logger.Error("failed to produce message", zap.Error(err), zap.String("handler", "bittorrent"))
		}

		logger.Info(
			"telnet send",
			zap.String("handler", "bittorrent"),
			zap.Uint8s("peer_id", msg.PeerID[:]),
			zap.Uint8s("inf_hash", msg.InfoHash[:]),
		)

		if err = binary.Write(conn, binary.BigEndian, msg); err != nil {
			logger.Error("failed to write message", zap.Error(err), zap.String("handler", "bittorrent"))
			break
		}
	}
	return nil
}
<commit_msg>produce bittorrent buffer<commit_after>package protocols

import (
	"bytes"
	"context"
	"encoding/binary"
	"net"

	"go.uber.org/zap"
)

type bittorrentMsg struct {
	Length             uint8
	ProtocolIdentifier [19]uint8
	Reserved           [8]uint8
	InfoHash           [20]uint8
	PeerID             [20]uint8
}

// HandleBittorrent handles a Bittorrent connection
func HandleBittorrent(ctx context.Context, conn net.Conn, logger Logger, h Honeypot) error {
	var err error
	defer func() {
		if err = conn.Close(); err != nil {
			logger.Error("failed to close connection", zap.Error(err), zap.String("handler", "bittorrent"))
			return
		}
	}()
	buffer := make([]byte, 1024)
	var r *bytes.Reader
	msg := bittorrentMsg{}
	md, err := h.MetadataByConnection(conn)
	if err != nil {
		return err
	}
	for {
		n, err := conn.Read(buffer)
		if err == nil || n > 0 {
			r = bytes.NewReader(buffer)
			if err := binary.Read(r, binary.BigEndian, &msg); err != nil {
				logger.Error("failed to read message", zap.Error(err), zap.String("handler", "bittorrent"))
				break
			}
			if err = h.Produce(conn, md, buffer); err != nil {
				logger.Error("failed to produce message", zap.Error(err), zap.String("handler", "bittorrent"))
			}

			logger.Info(
				"telnet send",
				zap.String("handler", "bittorrent"),
				zap.Uint8s("peer_id", msg.PeerID[:]),
				zap.Uint8s("inf_hash", msg.InfoHash[:]),
			)

			if err = binary.Write(conn, binary.BigEndian, msg); err != nil {
				logger.Error("failed to write message", zap.Error(err), zap.String("handler", "bittorrent"))
				break
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package ui

import (
	"fmt"

	"github.com/fabiofalci/sconsify/sconsify"
	"github.com/jroimartin/gocui"
)

func keybindings() error {
	views := []string{VIEW_TRACKS, VIEW_PLAYLISTS, VIEW_QUEUE}
	allViews := ""
	for _, view := range views {
		if err := gui.g.SetKeybinding(view, 'p', 0, pauseCurrentSelectedTrack); err != nil {
			return err
		}
		if err := gui.g.SetKeybinding(view, 'r', 0, setRandomMode); err != nil {
			return err
		}
		if err := gui.g.SetKeybinding(view, 'R', 0, setAllRandomMode); err != nil {
			return err
		}
		if err := gui.g.SetKeybinding(view, '>', 0, nextCommand); err != nil {
			return err
		}
		if err := gui.g.SetKeybinding(view, '/', 0, enableSearchInputCommand); err != nil {
			return err
		}
		if err := gui.g.SetKeybinding(view, 'j', 0, cursorDown); err != nil {
			return err
		}
		if err := gui.g.SetKeybinding(view, 'k', 0, cursorUp); err != nil {
			return err
		}
		if err := gui.g.SetKeybinding(view, 'q', 0, quit); err != nil {
			return err
		}
	}

	if err := gui.g.SetKeybinding(VIEW_TRACKS, gocui.KeySpace, 0, playCurrentSelectedTrack); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_TRACKS, gocui.KeyEnter, 0, playCurrentSelectedTrack); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(VIEW_TRACKS, 'u', 0, queueCommand); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_PLAYLISTS, 'd', 0, removeFromPlaylistsCommand); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_QUEUE, 'd', 0, removeFromQueueCommand); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_QUEUE, 'D', 0, removeAllFromQueueCommand); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_STATUS, gocui.KeyEnter, 0, searchCommand); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(allViews, gocui.KeyHome, 0, cursorHome); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(allViews, gocui.KeyEnd, 0, cursorEnd); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(allViews, gocui.KeyPgup, 0, cursorPgup); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(allViews, gocui.KeyPgdn, 0, cursorPgdn); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(allViews, gocui.KeyArrowDown, 0, cursorDown); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(allViews, gocui.KeyArrowUp, 0, cursorUp); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(VIEW_TRACKS, gocui.KeyArrowLeft, 0, mainNextViewLeft); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_QUEUE, gocui.KeyArrowLeft, 0, nextView); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(VIEW_PLAYLISTS, gocui.KeyArrowRight, 0, nextView); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_TRACKS, gocui.KeyArrowRight, 0, mainNextViewRight); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(VIEW_TRACKS, 'h', 0, mainNextViewLeft); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_QUEUE, 'h', 0, nextView); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_PLAYLISTS, 'l', 0, nextView); err != nil {
		return err
	}
	if err := gui.g.SetKeybinding(VIEW_TRACKS, 'l', 0, mainNextViewRight); err != nil {
		return err
	}

	if err := gui.g.SetKeybinding(allViews, gocui.KeyCtrlC, 0, quit); err != nil {
		return err
	}

	return nil
}

func playCurrentSelectedTrack(g *gocui.Gui, v *gocui.View) error {
	if track := gui.getCurrentSelectedTrack(); track != nil {
		events.Play(track)
	}
	return nil
}

func pauseCurrentSelectedTrack(g *gocui.Gui, v *gocui.View) error {
	events.Pause()
	return nil
}

func setRandomMode(g *gocui.Gui, v *gocui.View) error {
	playlists.InvertMode(sconsify.RandomMode)
	gui.updateCurrentStatus()
	return nil
}

func setAllRandomMode(g *gocui.Gui, v *gocui.View) error {
	playlists.InvertMode(sconsify.AllRandomMode)
	gui.updateCurrentStatus()
	return nil
}

func nextCommand(g *gocui.Gui, v *gocui.View) error {
	gui.playNext()
	return nil
}

func queueCommand(g *gocui.Gui, v *gocui.View) error {
	if track := gui.getCurrentSelectedTrack(); track != nil {
		fmt.Fprintf(gui.queueView, "%v", track.GetTitle())
		queue.Add(track)
	}
	return nil
}

func removeAllFromQueueCommand(g *gocui.Gui, v *gocui.View) error {
	queue.RemoveAll()
	gui.updateQueueView()
	return gui.enableTracksView()
}

func removeFromQueueCommand(g *gocui.Gui, v *gocui.View) error {
	if index := gui.getQueueSelectedTrackIndex(); index > -1 {
		queue.Remove(index)
		gui.updateQueueView()
	}
	return nil
}

func removeFromPlaylistsCommand(g *gocui.Gui, v *gocui.View) error {
	if playlist := gui.getSelectedPlaylist(); playlist != nil && playlist.IsSearch() {
		playlists.Remove(playlist.Name())
		gui.updatePlaylistsView()
		gui.updateTracksView()
	}
	return nil
}

func enableSearchInputCommand(g *gocui.Gui, v *gocui.View) error {
	gui.clearStatusView()
	gui.statusView.Editable = true
	gui.g.SetCurrentView(VIEW_STATUS)
	return nil
}

func searchCommand(g *gocui.Gui, v *gocui.View) error {
	// after user hit Enter, the typed command is at position -1
	if query, _ := gui.statusView.Line(-1); query != "" {
		events.Search(query)
	}
	gui.enableSideView()
	gui.clearStatusView()
	gui.statusView.Editable = false
	gui.updateCurrentStatus()
	return nil
}

func quit(g *gocui.Gui, v *gocui.View) error {
	consoleUserInterface.Shutdown()
	// TODO wait for shutdown
	// <-events.ShutdownUpdates()
	return gocui.ErrorQuit
}
<commit_msg>Improve keybinding<commit_after>package ui

import (
	"fmt"

	"github.com/fabiofalci/sconsify/sconsify"
	"github.com/jroimartin/gocui"
)

type KeyMapping struct {
	key  interface{}
	h    gocui.KeybindingHandler
	view string
}

func keybindings() error {
	keys := make([]*KeyMapping, 0)

	for _, view := range []string{VIEW_TRACKS, VIEW_PLAYLISTS, VIEW_QUEUE} {
		addKeyBinding(&keys, newKeyMapping('p', pauseCurrentSelectedTrack, view))
		addKeyBinding(&keys, newKeyMapping('r', setRandomMode, view))
		addKeyBinding(&keys, newKeyMapping('R', setAllRandomMode, view))
		addKeyBinding(&keys, newKeyMapping('>', nextCommand, view))
		addKeyBinding(&keys, newKeyMapping('/', enableSearchInputCommand, view))
		addKeyBinding(&keys, newKeyMapping('j', cursorDown, view))
		addKeyBinding(&keys, newKeyMapping('k', cursorUp, view))
		addKeyBinding(&keys, newKeyMapping('q', quit, view))
	}

	allViews := ""
	addKeyBinding(&keys, newKeyMapping(gocui.KeySpace, playCurrentSelectedTrack, VIEW_TRACKS))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyEnter, playCurrentSelectedTrack, VIEW_TRACKS))
	addKeyBinding(&keys, newKeyMapping('u', queueCommand, VIEW_TRACKS))
	addKeyBinding(&keys, newKeyMapping('d', removeFromPlaylistsCommand, VIEW_PLAYLISTS))
	addKeyBinding(&keys, newKeyMapping('d', removeFromQueueCommand, VIEW_QUEUE))
	addKeyBinding(&keys, newKeyMapping('D', removeAllFromQueueCommand, VIEW_QUEUE))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyEnter, searchCommand, VIEW_STATUS))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyHome, cursorHome, allViews))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyEnd, cursorEnd, allViews))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyPgup, cursorPgup, allViews))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyPgdn, cursorPgdn, allViews))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyArrowDown, cursorDown, allViews))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyArrowUp, cursorUp, allViews))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyArrowLeft, mainNextViewLeft, VIEW_TRACKS))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyArrowLeft, nextView, VIEW_QUEUE))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyArrowRight, nextView, VIEW_PLAYLISTS))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyArrowRight, mainNextViewRight, VIEW_TRACKS))
	addKeyBinding(&keys, newKeyMapping('h', mainNextViewLeft, VIEW_TRACKS))
	addKeyBinding(&keys, newKeyMapping('h', nextView, VIEW_QUEUE))
	addKeyBinding(&keys, newKeyMapping('l', nextView, VIEW_PLAYLISTS))
	addKeyBinding(&keys, newKeyMapping('l', mainNextViewRight, VIEW_TRACKS))
	addKeyBinding(&keys, newKeyMapping(gocui.KeyCtrlC, quit, allViews))

	for _, key := range keys {
		// it needs to copy the key because closures copy var references and we don't
		// want to execute always the last action
		keyCopy := key
		if err := gui.g.SetKeybinding(key.view, key.key, 0,
			func(g *gocui.Gui, v *gocui.View) error {
				return keyCopy.h(g, v)
			}); err != nil {
			return err
		}
	}

	return nil
}

func addKeyBinding(keys *[]*KeyMapping, key *KeyMapping) {
	*keys = append(*keys, key)
}

func newKeyMapping(key interface{}, h gocui.KeybindingHandler, view string) *KeyMapping {
	return &KeyMapping{key: key, h: h, view: view}
}

func playCurrentSelectedTrack(g *gocui.Gui, v *gocui.View) error {
	if track := gui.getCurrentSelectedTrack(); track != nil {
		events.Play(track)
	}
	return nil
}

func pauseCurrentSelectedTrack(g *gocui.Gui, v *gocui.View) error {
	events.Pause()
	return nil
}

func setRandomMode(g *gocui.Gui, v *gocui.View) error {
	playlists.InvertMode(sconsify.RandomMode)
	gui.updateCurrentStatus()
	return nil
}

func setAllRandomMode(g *gocui.Gui, v *gocui.View) error {
	playlists.InvertMode(sconsify.AllRandomMode)
	gui.updateCurrentStatus()
	return nil
}

func nextCommand(g *gocui.Gui, v *gocui.View) error {
	gui.playNext()
	return nil
}

func queueCommand(g *gocui.Gui, v *gocui.View) error {
	if track := gui.getCurrentSelectedTrack(); track != nil {
		fmt.Fprintf(gui.queueView, "%v", track.GetTitle())
		queue.Add(track)
	}
	return nil
}

func removeAllFromQueueCommand(g *gocui.Gui, v *gocui.View) error {
	queue.RemoveAll()
	gui.updateQueueView()
	return gui.enableTracksView()
}

func removeFromQueueCommand(g *gocui.Gui, v *gocui.View) error {
	if index := gui.getQueueSelectedTrackIndex(); index > -1 {
		queue.Remove(index)
		gui.updateQueueView()
	}
	return nil
}

func removeFromPlaylistsCommand(g *gocui.Gui, v *gocui.View) error {
	if playlist := gui.getSelectedPlaylist(); playlist != nil && playlist.IsSearch() {
		playlists.Remove(playlist.Name())
		gui.updatePlaylistsView()
		gui.updateTracksView()
	}
	return nil
}

func enableSearchInputCommand(g *gocui.Gui, v *gocui.View) error {
	gui.clearStatusView()
	gui.statusView.Editable = true
	gui.g.SetCurrentView(VIEW_STATUS)
	return nil
}

func searchCommand(g *gocui.Gui, v *gocui.View) error {
	// after user hit Enter, the typed command is at position -1
	if query, _ := gui.statusView.Line(-1); query != "" {
		events.Search(query)
	}
	gui.enableSideView()
	gui.clearStatusView()
	gui.statusView.Editable = false
	gui.updateCurrentStatus()
	return nil
}

func quit(g *gocui.Gui, v *gocui.View) error {
	consoleUserInterface.Shutdown()
	// TODO wait for shutdown
	// <-events.ShutdownUpdates()
	return gocui.ErrorQuit
}
<|endoftext|>


<commit_before>package invoke
import (
	"github.com/pkg/errors"
	"github.com/hyperledger/fabric/core/chaincode/shim"
	"encoding/json"
	"build-chaincode/data"
)

var logger = shim.NewLogger("invoke")

func main() {
	logger.SetLevel(shim.LogInfo)
}

var Functions = map[string]func(shim.ChaincodeStubInterface,[]string)([]byte, error) {
    "add_user": add_user,
    "add_thing": add_thing,
    "add_test_data": add_test_data,
}

// Invoke function.
func Invoke(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	logger.Infof("-- Invoking function %v with args %v", function, args)

	if function == "init" {
		return Init(stub, "init", args)
    } else {
        return Functions[function](stub,args)
    }

	return nil, errors.New("Received unknown invoke function name")
}

//==============================================================================================================================
//	Init Function - Called when the user deploys the chaincode
//==============================================================================================================================

func Init(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	logger.Infof("Deployed chaincode.")
	return nil, data.ResetIndexes(stub)
}

//==============================================================================================================================
//		Invoke Functions
//==============================================================================================================================

func add_test_data(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
    var usersIndex = args[0]
    var thingsIndex = args[1]
		logger.Debugf(usersIndex)
    var users []data.User
    err := json.Unmarshal([]byte(usersIndex), &users)
	if err != nil { return nil, err }
    for _,user := range users {
        data.Save(stub, user)
    }

    var things []data.Thing
    err = json.Unmarshal([]byte(thingsIndex), &things)
	if err != nil { return nil, err }
    for _,thing := range things {
        data.Save(stub, thing)
    }
    return nil,err
}

// args 0 is the caller id (not anymore needed in fabric v. 0.6)
func add_user(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var item data.User
	err := json.Unmarshal([]byte(args[1]), &item)
	if err != nil { return nil, err }
	return nil, data.Save(stub, item)
}
func add_thing(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var item data.Thing
	err := json.Unmarshal([]byte(args[0]), &item)
	if err != nil { return nil, err }
	return nil, data.Save(stub, item)
}
<commit_msg>updating array index<commit_after>package invoke
import (
	"github.com/pkg/errors"
	"github.com/hyperledger/fabric/core/chaincode/shim"
	"encoding/json"
	"build-chaincode/data"
)

var logger = shim.NewLogger("invoke")

func main() {
	logger.SetLevel(shim.LogInfo)
}

var Functions = map[string]func(shim.ChaincodeStubInterface,[]string)([]byte, error) {
    "add_user": add_user,
    "add_thing": add_thing,
    "add_test_data": add_test_data,
}

// Invoke function.
func Invoke(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	logger.Infof("-- Invoking function %v with args %v", function, args)

	if function == "init" {
		return Init(stub, "init", args)
    } else {
        return Functions[function](stub,args)
    }

	return nil, errors.New("Received unknown invoke function name")
}

//==============================================================================================================================
//	Init Function - Called when the user deploys the chaincode
//==============================================================================================================================

func Init(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	logger.Infof("Deployed chaincode.")
	return nil, data.ResetIndexes(stub)
}

//==============================================================================================================================
//		Invoke Functions
//==============================================================================================================================

func add_test_data(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
    var usersIndex = args[0]
    var thingsIndex = args[1]
		logger.Debugf(usersIndex)
    var users []data.User
    err := json.Unmarshal([]byte(usersIndex), &users)
	if err != nil { return nil, err }
    for _,user := range users {
        data.Save(stub, user)
    }

    var things []data.Thing
    err = json.Unmarshal([]byte(thingsIndex), &things)
	if err != nil { return nil, err }
    for _,thing := range things {
        data.Save(stub, thing)
    }
    return nil,err
}

// args 0 is the caller id (not anymore needed in fabric v. 0.6)
func add_user(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var item data.User
	err := json.Unmarshal([]byte(args[1]), &item)
	if err != nil { return nil, err }
	return nil, data.Save(stub, item)
}
func add_thing(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var item data.Thing
	err := json.Unmarshal([]byte(args[1]), &item)
	if err != nil { return nil, err }
	return nil, data.Save(stub, item)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen_test

import (
	"io/ioutil"
	"os"
	"testing"

	"path/filepath"

	"github.com/stretchr/testify/assert"
	"github.com/uber/zanzibar/codegen"
	testlib "github.com/uber/zanzibar/test/lib"
)

const (
	clientThrift   = "../examples/example-gateway/idl/github.com/uber/zanzibar/clients/bar/bar.thrift"
	endpointThrift = "../examples/example-gateway/idl/github.com/uber/zanzibar/endpoints/bar/bar.thrift"
	tmpDir         = "../../.tmp_gen"
)

func cmpGoldenFile(t *testing.T, actualFile string, goldenFileDir string) {
	b, err := ioutil.ReadFile(actualFile)
	if !assert.NoError(t, err, "failed to read genereated file %s", err) {
		return
	}
	goldenFile := filepath.Join(goldenFileDir, filepath.Base(actualFile)+"gen")
	testlib.CompareGoldenFile(t, goldenFile, b)
}

func TestGenerateBar(t *testing.T) {
	assert.NoError(t, os.RemoveAll(tmpDir), "failed to clean temporary directory")
	if err := os.MkdirAll(tmpDir, os.ModePerm); !assert.NoError(t, err, "failed to create temporary directory", err) {
		return
	}

	relativeGatewayPath := "../examples/example-gateway"
	absGatewayPath, err := filepath.Abs(relativeGatewayPath)
	if !assert.NoError(t, err, "failed to get abs path %s", err) {
		return
	}

	packageHelper, err := codegen.NewPackageHelper(
		"github.com/uber/zanzibar/examples/example-gateway",
		absGatewayPath,
		"middlewares/middleware-config.json",
		"./idl",
		"github.com/uber/zanzibar/examples/example-gateway/build/gen-code",
		tmpDir,
		testCopyrightHeader,
	)
	if !assert.NoError(t, err, "failed to create package helper", err) {
		return
	}

	moduleSystem, err := codegen.NewDefaultModuleSystem(packageHelper)
	if !assert.NoError(t, err, "failed to create module system", err) {
		return
	}

	_, buildErr := moduleSystem.GenerateBuild(
		"github.com/uber/zanzibar/examples/example-gateway",
		absGatewayPath,
		packageHelper.CodeGenTargetPath(),
	)
	if !assert.NoError(t, buildErr, "failed to create clients init %s", buildErr) {
		return
	}

	/* TODO: delete this
	gateway, err := codegen.NewGatewaySpec(
		moduleInstances,
		packageHelper,
		absGatewayPath,
		"./endpoints",
		"./middlewares/middleware-config.json",
		"example-gateway",
	)
	*/

	if !assert.NoError(t, err, "failed to create gateway spec %s", err) {
		return
	}

	cmpGoldenFile(
		t,
		filepath.Join(absGatewayPath, tmpDir, "clients", "bar", "bar.go"),
		"./test_data/clients",
	)

	/*
		err = gateway.GenerateEndpointRegisterFile()
		if !assert.NoError(t, err, "failed to create endpoint index %s", err) {
			return
		}
	*/

	endpoints, err := ioutil.ReadDir(
		filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar"),
	)
	if !assert.NoError(t, err, "cannot read dir %s", err) {
		return
	}

	for _, file := range endpoints {
		if file.IsDir() || isTestFile(file.Name()) {
			continue
		}

		cmpGoldenFile(
			t,
			filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar", file.Name()),
			"./test_data/endpoints",
		)
	}
	cmpGoldenFile(
		t,
		filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar", "bar_structs.go"),
		"./test_data/endpoints",
	)

	for _, file := range endpoints {
		if file.IsDir() || !isTestFile(file.Name()) {
			continue
		}

		cmpGoldenFile(
			t,
			filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar", file.Name()),
			"./test_data/endpoint_tests",
		)
	}
}

func isTestFile(filename string) bool {
	if len(filename) < 8 {
		return false
	}
	footer := filename[len(filename)-8 : len(filename)]
	return footer == "_test.go"
}
<commit_msg>stylefix test<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen_test

import (
	"io/ioutil"
	"os"
	"testing"

	"path/filepath"

	"github.com/stretchr/testify/assert"
	"github.com/uber/zanzibar/codegen"
	testlib "github.com/uber/zanzibar/test/lib"
)

const (
	clientThrift   = "../examples/example-gateway/idl/github.com/uber/zanzibar/clients/bar/bar.thrift"
	endpointThrift = "../examples/example-gateway/idl/github.com/uber/zanzibar/endpoints/bar/bar.thrift"
	tmpDir         = "../../.tmp_gen"
)

func cmpGoldenFile(t *testing.T, actualFile string, goldenFileDir string) {
	b, err := ioutil.ReadFile(actualFile)
	if !assert.NoError(t, err, "failed to read genereated file %s", err) {
		return
	}
	goldenFile := filepath.Join(goldenFileDir, filepath.Base(actualFile)+"gen")
	testlib.CompareGoldenFile(t, goldenFile, b)
}

func TestGenerateBar(t *testing.T) {
	assert.NoError(t, os.RemoveAll(tmpDir), "failed to clean temporary directory")
	if err := os.MkdirAll(tmpDir, os.ModePerm); !assert.NoError(t, err, "failed to create temporary directory", err) {
		return
	}

	relativeGatewayPath := "../examples/example-gateway"
	absGatewayPath, err := filepath.Abs(relativeGatewayPath)
	if !assert.NoError(t, err, "failed to get abs path %s", err) {
		return
	}

	packageHelper, err := codegen.NewPackageHelper(
		"github.com/uber/zanzibar/examples/example-gateway",
		absGatewayPath,
		"middlewares/middleware-config.json",
		"./idl",
		"github.com/uber/zanzibar/examples/example-gateway/build/gen-code",
		tmpDir,
		testCopyrightHeader,
	)
	if !assert.NoError(t, err, "failed to create package helper", err) {
		return
	}

	moduleSystem, err := codegen.NewDefaultModuleSystem(packageHelper)
	if !assert.NoError(t, err, "failed to create module system", err) {
		return
	}

	_, buildErr := moduleSystem.GenerateBuild(
		"github.com/uber/zanzibar/examples/example-gateway",
		absGatewayPath,
		packageHelper.CodeGenTargetPath(),
	)
	if !assert.NoError(t, buildErr, "failed to create clients init %s", buildErr) {
		return
	}

	/* TODO: delete this
	gateway, err := codegen.NewGatewaySpec(
		moduleInstances,
		packageHelper,
		absGatewayPath,
		"./endpoints",
		"./middlewares/middleware-config.json",
		"example-gateway",
	)
	*/

	if !assert.NoError(t, err, "failed to create gateway spec %s", err) {
		return
	}

	cmpGoldenFile(
		t,
		filepath.Join(absGatewayPath, tmpDir, "clients", "bar", "bar.go"),
		"./test_data/clients",
	)

	/*
		err = gateway.GenerateEndpointRegisterFile()
		if !assert.NoError(t, err, "failed to create endpoint index %s", err) {
			return
		}
	*/

	endpoints, err := ioutil.ReadDir(
		filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar"),
	)
	if !assert.NoError(t, err, "cannot read dir %s", err) {
		return
	}

	for _, file := range endpoints {
		if file.IsDir() || isTestFile(file.Name()) {
			continue
		}

		cmpGoldenFile(
			t,
			filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar", file.Name()),
			"./test_data/endpoints",
		)
	}
	cmpGoldenFile(
		t,
		filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar", "bar_structs.go"),
		"./test_data/endpoints",
	)

	for _, file := range endpoints {
		if file.IsDir() || !isTestFile(file.Name()) {
			continue
		}

		cmpGoldenFile(
			t,
			filepath.Join(absGatewayPath, tmpDir, "endpoints", "bar", file.Name()),
			"./test_data/endpoint_tests",
		)
	}
}

func isTestFile(filename string) bool {
	if len(filename) < 8 {
		return false
	}
	footer := filename[len(filename)-8:]
	return footer == "_test.go"
}
<|endoftext|>


<commit_before>package application

import (
	"cf"
	"cf/api"
	"cf/configuration"
	"cf/errors"
	"cf/models"
	"cf/requirements"
	"cf/terminal"
	"fmt"
	"github.com/cloudfoundry/loggregatorlib/logmessage"
	"github.com/codegangsta/cli"
	"os"
	"strconv"
	"strings"
	"time"
)

const (
	DefaultStagingTimeout = 15 * time.Minute
	DefaultStartupTimeout = 5 * time.Minute
	DefaultPingerThrottle = 5 * time.Second
)

const LogMessageTypeStaging = "STG"

type Start struct {
	ui               terminal.UI
	config           configuration.Reader
	appDisplayer     ApplicationDisplayer
	appReq           requirements.ApplicationRequirement
	appRepo          api.ApplicationRepository
	appInstancesRepo api.AppInstancesRepository
	logRepo          api.LogsRepository

	StartupTimeout time.Duration
	StagingTimeout time.Duration
	PingerThrottle time.Duration
}

type ApplicationStarter interface {
	SetStartTimeoutSeconds(timeout int)
	ApplicationStart(app models.Application) (updatedApp models.Application, err error)
}

func NewStart(ui terminal.UI, config configuration.Reader, appDisplayer ApplicationDisplayer, appRepo api.ApplicationRepository, appInstancesRepo api.AppInstancesRepository, logRepo api.LogsRepository) (cmd *Start) {
	cmd = new(Start)
	cmd.ui = ui
	cmd.config = config
	cmd.appDisplayer = appDisplayer
	cmd.appRepo = appRepo
	cmd.appInstancesRepo = appInstancesRepo
	cmd.logRepo = logRepo

	cmd.PingerThrottle = DefaultPingerThrottle

	if os.Getenv("CF_STAGING_TIMEOUT") != "" {
		duration, err := strconv.ParseInt(os.Getenv("CF_STAGING_TIMEOUT"), 10, 64)
		if err != nil {
			cmd.ui.Failed("invalid value for env var CF_STAGING_TIMEOUT\n%s", err)
		}
		cmd.StagingTimeout = time.Duration(duration) * time.Minute
	} else {
		cmd.StagingTimeout = DefaultStagingTimeout
	}

	if os.Getenv("CF_STARTUP_TIMEOUT") != "" {
		duration, err := strconv.ParseInt(os.Getenv("CF_STARTUP_TIMEOUT"), 10, 64)
		if err != nil {
			cmd.ui.Failed("invalid value for env var CF_STARTUP_TIMEOUT\n%s", err)
		}
		cmd.StartupTimeout = time.Duration(duration) * time.Minute
	} else {
		cmd.StartupTimeout = DefaultStartupTimeout
	}

	return
}

func (cmd *Start) GetRequirements(reqFactory requirements.Factory, c *cli.Context) (reqs []requirements.Requirement, err error) {
	if len(c.Args()) == 0 {
		err = errors.New("Incorrect Usage")
		cmd.ui.FailWithUsage(c, "start")
		return
	}

	cmd.appReq = reqFactory.NewApplicationRequirement(c.Args()[0])

	reqs = []requirements.Requirement{cmd.appReq}
	return
}

func (cmd *Start) Run(c *cli.Context) {
	cmd.ApplicationStart(cmd.appReq.GetApplication())
}

func (cmd *Start) ApplicationStart(app models.Application) (updatedApp models.Application, err error) {
	if app.State == "started" {
		cmd.ui.Say(terminal.WarningColor("App " + app.Name + " is already started"))
		return
	}

	stopLoggingChan := make(chan bool, 1)
	defer close(stopLoggingChan)
	loggingStartedChan := make(chan bool)
	defer close(loggingStartedChan)

	go cmd.tailStagingLogs(app, loggingStartedChan, stopLoggingChan)

	<-loggingStartedChan

	cmd.ui.Say("Starting app %s in org %s / space %s as %s...",
		terminal.EntityNameColor(app.Name),
		terminal.EntityNameColor(cmd.config.OrganizationFields().Name),
		terminal.EntityNameColor(cmd.config.SpaceFields().Name),
		terminal.EntityNameColor(cmd.config.Username()),
	)

	state := "STARTED"
	updatedApp, apiErr := cmd.appRepo.Update(app.Guid, models.AppParams{State: &state})
	if apiErr != nil {
		cmd.ui.Failed(apiErr.Error())
		return
	}

	cmd.ui.Ok()

	cmd.waitForInstancesToStage(updatedApp)
	stopLoggingChan <- true

	cmd.ui.Say("")

	cmd.waitForOneRunningInstance(updatedApp)
	cmd.ui.Say(terminal.HeaderColor("\nApp started\n"))

	cmd.appDisplayer.ShowApp(updatedApp)
	return
}

func (cmd *Start) SetStartTimeoutSeconds(timeout int) {
	cmd.StartupTimeout = time.Duration(timeout) * time.Second
}

func (cmd Start) tailStagingLogs(app models.Application, startChan chan bool, stopChan chan bool) {
	logChan := make(chan *logmessage.Message, 1000)
	go func() {
		defer close(logChan)

		onConnect := func() {
			startChan <- true
		}

		err := cmd.logRepo.TailLogsFor(app.Guid, onConnect, logChan, stopChan, 1)
		if err != nil {
			cmd.ui.Warn("Warning: error tailing logs")
			cmd.ui.Say("%s", err)
			startChan <- true
		}
	}()

	cmd.displayLogMessages(logChan)
}

func (cmd Start) displayLogMessages(logChan chan *logmessage.Message) {
	for msg := range logChan {
		if msg.GetLogMessage().GetSourceName() != "STG" {
			continue
		}
		cmd.ui.Say(simpleLogMessageOutput(msg))
	}
}

func (cmd Start) waitForInstancesToStage(app models.Application) {
	stagingStartTime := time.Now()
	_, err := cmd.appInstancesRepo.GetInstances(app.Guid)

	for err != nil && time.Since(stagingStartTime) < cmd.StagingTimeout {
		if err, ok := err.(errors.HttpError); ok && err.ErrorCode() != errors.APP_NOT_STAGED {
			cmd.ui.Say("")
			cmd.ui.Failed(fmt.Sprintf("%s\n\nTIP: use '%s' for more information",
				err.Error(),
				terminal.CommandColor(fmt.Sprintf("%s logs %s --recent", cf.Name(), app.Name))))
			return
		}
		cmd.ui.Wait(cmd.PingerThrottle)
		_, err = cmd.appInstancesRepo.GetInstances(app.Guid)
	}
	return
}

func (cmd Start) waitForOneRunningInstance(app models.Application) {
	var runningCount, startingCount, flappingCount, downCount int
	startupStartTime := time.Now()

	for runningCount == 0 {
		if time.Since(startupStartTime) > cmd.StartupTimeout {
			cmd.ui.Failed(fmt.Sprintf("Start app timeout\n\nTIP: use '%s' for more information", terminal.CommandColor(fmt.Sprintf("%s logs %s --recent", cf.Name(), app.Name))))
			return
		}

		instances, apiErr := cmd.appInstancesRepo.GetInstances(app.Guid)
		if apiErr != nil {
			cmd.ui.Wait(cmd.PingerThrottle)
			continue
		}

		totalCount := len(instances)
		runningCount, startingCount, flappingCount, downCount = 0, 0, 0, 0

		for _, inst := range instances {
			switch inst.State {
			case models.InstanceRunning:
				runningCount++
			case models.InstanceStarting:
				startingCount++
			case models.InstanceFlapping:
				flappingCount++
			case models.InstanceDown:
				downCount++
			}
		}

		cmd.ui.Say(instancesDetails(startingCount, downCount, runningCount, flappingCount, totalCount))

		if flappingCount > 0 {
			cmd.ui.Failed(fmt.Sprintf("Start unsuccessful\n\nTIP: use '%s' for more information", terminal.CommandColor(fmt.Sprintf("%s logs %s --recent", cf.Name(), app.Name))))
			return
		}
	}
}

func instancesDetails(startingCount, downCount, runningCount, flappingCount, totalCount int) string {
	details := []string{fmt.Sprintf("%d of %d instances running", runningCount, totalCount)}

	if startingCount > 0 {
		details = append(details, fmt.Sprintf("%d starting", startingCount))
	}

	if downCount > 0 {
		details = append(details, fmt.Sprintf("%d down", downCount))
	}

	if flappingCount > 0 {
		details = append(details, fmt.Sprintf("%d failing", flappingCount))
	}

	return strings.Join(details, ", ")
}
<commit_msg>Use constants instead of ignoring them<commit_after>package application

import (
	"cf"
	"cf/api"
	"cf/configuration"
	"cf/errors"
	"cf/models"
	"cf/requirements"
	"cf/terminal"
	"fmt"
	"github.com/cloudfoundry/loggregatorlib/logmessage"
	"github.com/codegangsta/cli"
	"os"
	"strconv"
	"strings"
	"time"
)

const (
	DefaultStagingTimeout = 15 * time.Minute
	DefaultStartupTimeout = 5 * time.Minute
	DefaultPingerThrottle = 5 * time.Second
)

const LogMessageTypeStaging = "STG"

type Start struct {
	ui               terminal.UI
	config           configuration.Reader
	appDisplayer     ApplicationDisplayer
	appReq           requirements.ApplicationRequirement
	appRepo          api.ApplicationRepository
	appInstancesRepo api.AppInstancesRepository
	logRepo          api.LogsRepository

	StartupTimeout time.Duration
	StagingTimeout time.Duration
	PingerThrottle time.Duration
}

type ApplicationStarter interface {
	SetStartTimeoutSeconds(timeout int)
	ApplicationStart(app models.Application) (updatedApp models.Application, err error)
}

func NewStart(ui terminal.UI, config configuration.Reader, appDisplayer ApplicationDisplayer, appRepo api.ApplicationRepository, appInstancesRepo api.AppInstancesRepository, logRepo api.LogsRepository) (cmd *Start) {
	cmd = new(Start)
	cmd.ui = ui
	cmd.config = config
	cmd.appDisplayer = appDisplayer
	cmd.appRepo = appRepo
	cmd.appInstancesRepo = appInstancesRepo
	cmd.logRepo = logRepo

	cmd.PingerThrottle = DefaultPingerThrottle

	if os.Getenv("CF_STAGING_TIMEOUT") != "" {
		duration, err := strconv.ParseInt(os.Getenv("CF_STAGING_TIMEOUT"), 10, 64)
		if err != nil {
			cmd.ui.Failed("invalid value for env var CF_STAGING_TIMEOUT\n%s", err)
		}
		cmd.StagingTimeout = time.Duration(duration) * time.Minute
	} else {
		cmd.StagingTimeout = DefaultStagingTimeout
	}

	if os.Getenv("CF_STARTUP_TIMEOUT") != "" {
		duration, err := strconv.ParseInt(os.Getenv("CF_STARTUP_TIMEOUT"), 10, 64)
		if err != nil {
			cmd.ui.Failed("invalid value for env var CF_STARTUP_TIMEOUT\n%s", err)
		}
		cmd.StartupTimeout = time.Duration(duration) * time.Minute
	} else {
		cmd.StartupTimeout = DefaultStartupTimeout
	}

	return
}

func (cmd *Start) GetRequirements(reqFactory requirements.Factory, c *cli.Context) (reqs []requirements.Requirement, err error) {
	if len(c.Args()) == 0 {
		err = errors.New("Incorrect Usage")
		cmd.ui.FailWithUsage(c, "start")
		return
	}

	cmd.appReq = reqFactory.NewApplicationRequirement(c.Args()[0])

	reqs = []requirements.Requirement{cmd.appReq}
	return
}

func (cmd *Start) Run(c *cli.Context) {
	cmd.ApplicationStart(cmd.appReq.GetApplication())
}

func (cmd *Start) ApplicationStart(app models.Application) (updatedApp models.Application, err error) {
	if app.State == "started" {
		cmd.ui.Say(terminal.WarningColor("App " + app.Name + " is already started"))
		return
	}

	stopLoggingChan := make(chan bool, 1)
	defer close(stopLoggingChan)
	loggingStartedChan := make(chan bool)
	defer close(loggingStartedChan)

	go cmd.tailStagingLogs(app, loggingStartedChan, stopLoggingChan)

	<-loggingStartedChan

	cmd.ui.Say("Starting app %s in org %s / space %s as %s...",
		terminal.EntityNameColor(app.Name),
		terminal.EntityNameColor(cmd.config.OrganizationFields().Name),
		terminal.EntityNameColor(cmd.config.SpaceFields().Name),
		terminal.EntityNameColor(cmd.config.Username()),
	)

	state := "STARTED"
	updatedApp, apiErr := cmd.appRepo.Update(app.Guid, models.AppParams{State: &state})
	if apiErr != nil {
		cmd.ui.Failed(apiErr.Error())
		return
	}

	cmd.ui.Ok()

	cmd.waitForInstancesToStage(updatedApp)
	stopLoggingChan <- true

	cmd.ui.Say("")

	cmd.waitForOneRunningInstance(updatedApp)
	cmd.ui.Say(terminal.HeaderColor("\nApp started\n"))

	cmd.appDisplayer.ShowApp(updatedApp)
	return
}

func (cmd *Start) SetStartTimeoutSeconds(timeout int) {
	cmd.StartupTimeout = time.Duration(timeout) * time.Second
}

func (cmd Start) tailStagingLogs(app models.Application, startChan chan bool, stopChan chan bool) {
	logChan := make(chan *logmessage.Message, 1000)
	go func() {
		defer close(logChan)

		onConnect := func() {
			startChan <- true
		}

		err := cmd.logRepo.TailLogsFor(app.Guid, onConnect, logChan, stopChan, 1)
		if err != nil {
			cmd.ui.Warn("Warning: error tailing logs")
			cmd.ui.Say("%s", err)
			startChan <- true
		}
	}()

	cmd.displayLogMessages(logChan)
}

func (cmd Start) displayLogMessages(logChan chan *logmessage.Message) {
	for msg := range logChan {
		if msg.GetLogMessage().GetSourceName() != LogMessageTypeStaging {
			continue
		}
		cmd.ui.Say(simpleLogMessageOutput(msg))
	}
}

func (cmd Start) waitForInstancesToStage(app models.Application) {
	stagingStartTime := time.Now()
	_, err := cmd.appInstancesRepo.GetInstances(app.Guid)

	for err != nil && time.Since(stagingStartTime) < cmd.StagingTimeout {
		if err, ok := err.(errors.HttpError); ok && err.ErrorCode() != errors.APP_NOT_STAGED {
			cmd.ui.Say("")
			cmd.ui.Failed(fmt.Sprintf("%s\n\nTIP: use '%s' for more information",
				err.Error(),
				terminal.CommandColor(fmt.Sprintf("%s logs %s --recent", cf.Name(), app.Name))))
			return
		}
		cmd.ui.Wait(cmd.PingerThrottle)
		_, err = cmd.appInstancesRepo.GetInstances(app.Guid)
	}
	return
}

func (cmd Start) waitForOneRunningInstance(app models.Application) {
	var runningCount, startingCount, flappingCount, downCount int
	startupStartTime := time.Now()

	for runningCount == 0 {
		if time.Since(startupStartTime) > cmd.StartupTimeout {
			cmd.ui.Failed(fmt.Sprintf("Start app timeout\n\nTIP: use '%s' for more information", terminal.CommandColor(fmt.Sprintf("%s logs %s --recent", cf.Name(), app.Name))))
			return
		}

		instances, apiErr := cmd.appInstancesRepo.GetInstances(app.Guid)
		if apiErr != nil {
			cmd.ui.Wait(cmd.PingerThrottle)
			continue
		}

		totalCount := len(instances)
		runningCount, startingCount, flappingCount, downCount = 0, 0, 0, 0

		for _, inst := range instances {
			switch inst.State {
			case models.InstanceRunning:
				runningCount++
			case models.InstanceStarting:
				startingCount++
			case models.InstanceFlapping:
				flappingCount++
			case models.InstanceDown:
				downCount++
			}
		}

		cmd.ui.Say(instancesDetails(startingCount, downCount, runningCount, flappingCount, totalCount))

		if flappingCount > 0 {
			cmd.ui.Failed(fmt.Sprintf("Start unsuccessful\n\nTIP: use '%s' for more information", terminal.CommandColor(fmt.Sprintf("%s logs %s --recent", cf.Name(), app.Name))))
			return
		}
	}
}

func instancesDetails(startingCount, downCount, runningCount, flappingCount, totalCount int) string {
	details := []string{fmt.Sprintf("%d of %d instances running", runningCount, totalCount)}

	if startingCount > 0 {
		details = append(details, fmt.Sprintf("%d starting", startingCount))
	}

	if downCount > 0 {
		details = append(details, fmt.Sprintf("%d down", downCount))
	}

	if flappingCount > 0 {
		details = append(details, fmt.Sprintf("%d failing", flappingCount))
	}

	return strings.Join(details, ", ")
}
<|endoftext|>


<commit_before>//Underscore is a package that transforms a given string to underscore, escaping special characteres.
//
// Example of usage:
//
// package main

// import(
//     "fmt"
//     "github.com/hgsigner/stringfy/underscore"
// )

// func main() {
//     un := underscore.PerformOn("São Paulo")
//     fmt.Prinln(un) // sao_paulo
// }
package underscore
<commit_msg>fix doc for underscore<commit_after>//Underscore is a package that transforms a given string to underscore, escaping special characteres.
//
// Example of usage:
//
// import(
//     "fmt"
//     "github.com/hgsigner/stringfy/underscore"
// )
//
// func main() {
//     un := underscore.PerformOn("São Paulo")
//     fmt.Prinln(un) // sao_paulo
// }
package underscore
<|endoftext|>


<commit_before>package build

import (
	"bytes"
	"flag"
	"fmt"
	cmdcommon "github.com/mitchellh/packer/common/command"
	"github.com/mitchellh/packer/packer"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
)

type Command byte

func (Command) Help() string {
	return strings.TrimSpace(helpText)
}

func (c Command) Run(env packer.Environment, args []string) int {
	var cfgDebug bool
	var cfgForce bool
	buildOptions := new(cmdcommon.BuildOptions)

	cmdFlags := flag.NewFlagSet("build", flag.ContinueOnError)
	cmdFlags.Usage = func() { env.Ui().Say(c.Help()) }
	cmdFlags.BoolVar(&cfgDebug, "debug", false, "debug mode for builds")
	cmdFlags.BoolVar(&cfgForce, "force", false, "force a build if artifacts exist")
	cmdcommon.BuildOptionFlags(cmdFlags, buildOptions)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	args = cmdFlags.Args()
	if len(args) != 1 {
		cmdFlags.Usage()
		return 1
	}

	if err := buildOptions.Validate(); err != nil {
		env.Ui().Error(err.Error())
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	userVars, err := buildOptions.AllUserVars()
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Error compiling user variables: %s", err))
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	// Read the file into a byte array so that we can parse the template
	log.Printf("Reading template: %s", args[0])
	tpl, err := packer.ParseTemplateFile(args[0])
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to parse template: %s", err))
		return 1
	}

	// The component finder for our builds
	components := &packer.ComponentFinder{
		Builder:       env.Builder,
		Hook:          env.Hook,
		PostProcessor: env.PostProcessor,
		Provisioner:   env.Provisioner,
	}

	// Go through each builder and compile the builds that we care about
	builds, err := buildOptions.Builds(tpl, components)
	if err != nil {
		env.Ui().Error(err.Error())
		return 1
	}

	if cfgDebug {
		env.Ui().Say("Debug mode enabled. Builds will not be parallelized.")
	}

	// Compile all the UIs for the builds
	colors := [5]packer.UiColor{
		packer.UiColorGreen,
		packer.UiColorCyan,
		packer.UiColorMagenta,
		packer.UiColorYellow,
		packer.UiColorBlue,
	}

	buildUis := make(map[string]packer.Ui)
	for i, b := range builds {
		ui := &packer.ColoredUi{
			Color: colors[i%len(colors)],
			Ui:    env.Ui(),
		}

		buildUis[b.Name()] = ui
		ui.Say(fmt.Sprintf("%s output will be in this color.", b.Name()))
	}

	// Add a newline between the color output and the actual output
	env.Ui().Say("")

	log.Printf("Build debug mode: %v", cfgDebug)
	log.Printf("Force build: %v", cfgForce)

	// Set the debug and force mode and prepare all the builds
	for _, b := range builds {
		log.Printf("Preparing build: %s", b.Name())
		b.SetDebug(cfgDebug)
		b.SetForce(cfgForce)
		// TODO(mitchellh): Handle warnings
		_, err := b.Prepare(userVars)
		if err != nil {
			env.Ui().Error(err.Error())
			return 1
		}
	}

	// Run all the builds in parallel and wait for them to complete
	var interruptWg, wg sync.WaitGroup
	interrupted := false
	artifacts := make(map[string][]packer.Artifact)
	errors := make(map[string]error)
	for _, b := range builds {
		// Increment the waitgroup so we wait for this item to finish properly
		wg.Add(1)

		// Handle interrupts for this build
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, os.Interrupt)
		defer signal.Stop(sigCh)
		go func(b packer.Build) {
			<-sigCh
			interruptWg.Add(1)
			defer interruptWg.Done()
			interrupted = true

			log.Printf("Stopping build: %s", b.Name())
			b.Cancel()
			log.Printf("Build cancelled: %s", b.Name())
		}(b)

		// Run the build in a goroutine
		go func(b packer.Build) {
			defer wg.Done()

			name := b.Name()
			log.Printf("Starting build run: %s", name)
			ui := buildUis[name]
			runArtifacts, err := b.Run(ui, env.Cache())

			if err != nil {
				ui.Error(fmt.Sprintf("Build '%s' errored: %s", name, err))
				errors[name] = err
			} else {
				ui.Say(fmt.Sprintf("Build '%s' finished.", name))
				artifacts[name] = runArtifacts
			}
		}(b)

		if cfgDebug {
			log.Printf("Debug enabled, so waiting for build to finish: %s", b.Name())
			wg.Wait()
		}

		if interrupted {
			log.Println("Interrupted, not going to start any more builds.")
			break
		}
	}

	// Wait for both the builds to complete and the interrupt handler,
	// if it is interrupted.
	log.Printf("Waiting on builds to complete...")
	wg.Wait()

	log.Printf("Builds completed. Waiting on interrupt barrier...")
	interruptWg.Wait()

	if interrupted {
		env.Ui().Say("Cleanly cancelled builds after being interrupted.")
		return 1
	}

	if len(errors) > 0 {
		env.Ui().Machine("error-count", strconv.FormatInt(int64(len(errors)), 10))

		env.Ui().Error("\n==> Some builds didn't complete successfully and had errors:")
		for name, err := range errors {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			ui.Machine("error", err.Error())

			env.Ui().Error(fmt.Sprintf("--> %s: %s", name, err))
		}
	}

	if len(artifacts) > 0 {
		env.Ui().Say("\n==> Builds finished. The artifacts of successful builds are:")
		for name, buildArtifacts := range artifacts {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			// Machine-readable helpful
			ui.Machine("artifact-count", strconv.FormatInt(int64(len(buildArtifacts)), 10))

			for i, artifact := range buildArtifacts {
				var message bytes.Buffer
				fmt.Fprintf(&message, "--> %s: ", name)

				if artifact != nil {
					fmt.Fprintf(&message, artifact.String())
				} else {
					fmt.Fprint(&message, "<nothing>")
				}

				iStr := strconv.FormatInt(int64(i), 10)
				if artifact != nil {
					ui.Machine("artifact", iStr, "builder-id", artifact.BuilderId())
					ui.Machine("artifact", iStr, "id", artifact.Id())
					ui.Machine("artifact", iStr, "string", artifact.String())

					files := artifact.Files()
					ui.Machine("artifact",
						iStr,
						"files-count", strconv.FormatInt(int64(len(files)), 10))
					for fi, file := range files {
						fiStr := strconv.FormatInt(int64(fi), 10)
						ui.Machine("artifact", iStr, "file", fiStr, file)
					}
				} else {
					ui.Machine("artifact", iStr, "nil")
				}

				ui.Machine("artifact", iStr, "end")
				env.Ui().Say(message.String())
			}
		}
	} else {
		env.Ui().Say("\n==> Builds finished but no artifacts were created.")
	}

	if len(errors) > 0 {
		// If any errors occurred, exit with a non-zero exit status
		return 1
	}

	return 0
}

func (Command) Synopsis() string {
	return "build image(s) from template"
}
<commit_msg>command/build: output warnings<commit_after>package build

import (
	"bytes"
	"flag"
	"fmt"
	cmdcommon "github.com/mitchellh/packer/common/command"
	"github.com/mitchellh/packer/packer"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
)

type Command byte

func (Command) Help() string {
	return strings.TrimSpace(helpText)
}

func (c Command) Run(env packer.Environment, args []string) int {
	var cfgDebug bool
	var cfgForce bool
	buildOptions := new(cmdcommon.BuildOptions)

	cmdFlags := flag.NewFlagSet("build", flag.ContinueOnError)
	cmdFlags.Usage = func() { env.Ui().Say(c.Help()) }
	cmdFlags.BoolVar(&cfgDebug, "debug", false, "debug mode for builds")
	cmdFlags.BoolVar(&cfgForce, "force", false, "force a build if artifacts exist")
	cmdcommon.BuildOptionFlags(cmdFlags, buildOptions)
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	args = cmdFlags.Args()
	if len(args) != 1 {
		cmdFlags.Usage()
		return 1
	}

	if err := buildOptions.Validate(); err != nil {
		env.Ui().Error(err.Error())
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	userVars, err := buildOptions.AllUserVars()
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Error compiling user variables: %s", err))
		env.Ui().Error("")
		env.Ui().Error(c.Help())
		return 1
	}

	// Read the file into a byte array so that we can parse the template
	log.Printf("Reading template: %s", args[0])
	tpl, err := packer.ParseTemplateFile(args[0])
	if err != nil {
		env.Ui().Error(fmt.Sprintf("Failed to parse template: %s", err))
		return 1
	}

	// The component finder for our builds
	components := &packer.ComponentFinder{
		Builder:       env.Builder,
		Hook:          env.Hook,
		PostProcessor: env.PostProcessor,
		Provisioner:   env.Provisioner,
	}

	// Go through each builder and compile the builds that we care about
	builds, err := buildOptions.Builds(tpl, components)
	if err != nil {
		env.Ui().Error(err.Error())
		return 1
	}

	if cfgDebug {
		env.Ui().Say("Debug mode enabled. Builds will not be parallelized.")
	}

	// Compile all the UIs for the builds
	colors := [5]packer.UiColor{
		packer.UiColorGreen,
		packer.UiColorCyan,
		packer.UiColorMagenta,
		packer.UiColorYellow,
		packer.UiColorBlue,
	}

	buildUis := make(map[string]packer.Ui)
	for i, b := range builds {
		ui := &packer.ColoredUi{
			Color: colors[i%len(colors)],
			Ui:    env.Ui(),
		}

		buildUis[b.Name()] = ui
		ui.Say(fmt.Sprintf("%s output will be in this color.", b.Name()))
	}

	// Add a newline between the color output and the actual output
	env.Ui().Say("")

	log.Printf("Build debug mode: %v", cfgDebug)
	log.Printf("Force build: %v", cfgForce)

	// Set the debug and force mode and prepare all the builds
	for _, b := range builds {
		log.Printf("Preparing build: %s", b.Name())
		b.SetDebug(cfgDebug)
		b.SetForce(cfgForce)

		warnings, err := b.Prepare(userVars)
		if err != nil {
			env.Ui().Error(err.Error())
			return 1
		}
		if len(warnings) > 0 {
			ui := buildUis[b.Name()]
			ui.Say(fmt.Sprintf("Warnings for build '%s':\n", b.Name()))
			for _, warning := range warnings {
				ui.Say(fmt.Sprintf("* %s", warning))
			}
			ui.Say("")
		}
	}

	// Run all the builds in parallel and wait for them to complete
	var interruptWg, wg sync.WaitGroup
	interrupted := false
	artifacts := make(map[string][]packer.Artifact)
	errors := make(map[string]error)
	for _, b := range builds {
		// Increment the waitgroup so we wait for this item to finish properly
		wg.Add(1)

		// Handle interrupts for this build
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, os.Interrupt)
		defer signal.Stop(sigCh)
		go func(b packer.Build) {
			<-sigCh
			interruptWg.Add(1)
			defer interruptWg.Done()
			interrupted = true

			log.Printf("Stopping build: %s", b.Name())
			b.Cancel()
			log.Printf("Build cancelled: %s", b.Name())
		}(b)

		// Run the build in a goroutine
		go func(b packer.Build) {
			defer wg.Done()

			name := b.Name()
			log.Printf("Starting build run: %s", name)
			ui := buildUis[name]
			runArtifacts, err := b.Run(ui, env.Cache())

			if err != nil {
				ui.Error(fmt.Sprintf("Build '%s' errored: %s", name, err))
				errors[name] = err
			} else {
				ui.Say(fmt.Sprintf("Build '%s' finished.", name))
				artifacts[name] = runArtifacts
			}
		}(b)

		if cfgDebug {
			log.Printf("Debug enabled, so waiting for build to finish: %s", b.Name())
			wg.Wait()
		}

		if interrupted {
			log.Println("Interrupted, not going to start any more builds.")
			break
		}
	}

	// Wait for both the builds to complete and the interrupt handler,
	// if it is interrupted.
	log.Printf("Waiting on builds to complete...")
	wg.Wait()

	log.Printf("Builds completed. Waiting on interrupt barrier...")
	interruptWg.Wait()

	if interrupted {
		env.Ui().Say("Cleanly cancelled builds after being interrupted.")
		return 1
	}

	if len(errors) > 0 {
		env.Ui().Machine("error-count", strconv.FormatInt(int64(len(errors)), 10))

		env.Ui().Error("\n==> Some builds didn't complete successfully and had errors:")
		for name, err := range errors {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			ui.Machine("error", err.Error())

			env.Ui().Error(fmt.Sprintf("--> %s: %s", name, err))
		}
	}

	if len(artifacts) > 0 {
		env.Ui().Say("\n==> Builds finished. The artifacts of successful builds are:")
		for name, buildArtifacts := range artifacts {
			// Create a UI for the machine readable stuff to be targetted
			ui := &packer.TargettedUi{
				Target: name,
				Ui:     env.Ui(),
			}

			// Machine-readable helpful
			ui.Machine("artifact-count", strconv.FormatInt(int64(len(buildArtifacts)), 10))

			for i, artifact := range buildArtifacts {
				var message bytes.Buffer
				fmt.Fprintf(&message, "--> %s: ", name)

				if artifact != nil {
					fmt.Fprintf(&message, artifact.String())
				} else {
					fmt.Fprint(&message, "<nothing>")
				}

				iStr := strconv.FormatInt(int64(i), 10)
				if artifact != nil {
					ui.Machine("artifact", iStr, "builder-id", artifact.BuilderId())
					ui.Machine("artifact", iStr, "id", artifact.Id())
					ui.Machine("artifact", iStr, "string", artifact.String())

					files := artifact.Files()
					ui.Machine("artifact",
						iStr,
						"files-count", strconv.FormatInt(int64(len(files)), 10))
					for fi, file := range files {
						fiStr := strconv.FormatInt(int64(fi), 10)
						ui.Machine("artifact", iStr, "file", fiStr, file)
					}
				} else {
					ui.Machine("artifact", iStr, "nil")
				}

				ui.Machine("artifact", iStr, "end")
				env.Ui().Say(message.String())
			}
		}
	} else {
		env.Ui().Say("\n==> Builds finished but no artifacts were created.")
	}

	if len(errors) > 0 {
		// If any errors occurred, exit with a non-zero exit status
		return 1
	}

	return 0
}

func (Command) Synopsis() string {
	return "build image(s) from template"
}
<|endoftext|>


<commit_before>package commandHandlers

// RouterStruct is struct for handling command to handler
type RouterStruct struct {
	routes map[string]CommandHandler
}

// Go returns from router to work with
func (router RouterStruct) Go(command string) CommandHandler {
	handler, found := router.routes[command]
	if found {
		return handler
	}
	return Custom
}

// Router explicitly maps input chat command to a handler
var Router = RouterStruct{routes: map[string]CommandHandler{}}

//"new":   New,
//"alias": Alias}}
<commit_msg>full migration to twitch api v5<commit_after>package commandHandlers

// RouterStruct is struct for handling command to handler
type RouterStruct struct {
	routes map[string]CommandHandler
}

// Go returns from router to work with
func (router RouterStruct) Go(command string) CommandHandler {
	handler, found := router.routes[command]
	if found {
		return handler
	}
	return Custom
}

// Router explicitly maps input chat command to a handler
var Router = RouterStruct{routes: map[string]CommandHandler{
	"new":   New,
	"alias": Alias}}
<|endoftext|>


<commit_before>package definitions

import (
	"bytes"
	"errors"
	"fmt"
	"strings"
	"text/template"

	"github.com/cihangir/gene/generators/common"
	"github.com/cihangir/gene/writers"
	"github.com/cihangir/schema"
	"github.com/cihangir/stringext"
)

const generatorName = "sql-definition"

type Generator struct {
	DatabaseName  string
	SchemaName    string
	TableName     string
	RoleName      string
	FieldNameCase string `default:"snake"`
}

func New() *Generator {
	return &Generator{}
}

func (g *Generator) Name() string {
	return generatorName
}

// Generate generates the basic CRUD statements for the models
func (g *Generator) Generate(context *common.Context, s *schema.Schema) ([]common.Output, error) {
	outputs := make([]common.Output, 0)

	if s.Title == "" {
		return outputs, errors.New("Title should be set")
	}

	moduleName := context.FieldNameFunc(s.Title)

	settings := g.generateSettings(moduleName, s)

	for _, name := range schema.SortedKeys(s.Definitions) {
		def := s.Definitions[name]

		// schema should have our generator
		if !def.Generators.Has(g.Name()) {
			continue
		}

		settingsDef := g.setDefaultSettings(settings, def)
		settingsDef.Set("tableName", context.FieldNameFunc(def.Title))

		//
		// generate roles
		//
		role, err := DefineRole(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content:     role,
			Path:        fmt.Sprintf("%s/001-%s_roles.sql", context.Config.Target, settingsDef.Get("databaseName").(string)),
			DoNotFormat: true,
		})

		//
		// generate database
		//
		db, err := DefineDatabase(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content:     db,
			Path:        fmt.Sprintf("%s/002-%s_database.sql", context.Config.Target, settingsDef.Get("databaseName").(string)),
			DoNotFormat: true,
		})

		//
		// generate extenstions
		//
		extenstions, err := DefineExtensions(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: extenstions,
			Path: fmt.Sprintf(
				"%s/003-%s_extensions.sql",
				context.Config.Target,
				settingsDef.Get("databaseName").(string)),
			DoNotFormat: true,
		})

		//
		// generate schema
		//
		sc, err := DefineSchema(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: sc,
			Path: fmt.Sprintf(
				"%s/%s/004-schema.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate sequences
		//
		sequence, err := DefineSequence(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: sequence,
			Path: fmt.Sprintf(
				"%s/%s/005-%s-sequence.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate types
		//
		types, err := DefineTypes(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: types,
			Path: fmt.Sprintf(
				"%s/%s/006-%s-types.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate tables
		//
		table, err := DefineTable(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: table,
			Path: fmt.Sprintf(
				"%s/%s/007-%s-table.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate constraints
		//
		constraints, err := DefineConstraints(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: constraints,
			Path: fmt.Sprintf(
				"%s/%s/007-%s-constraints.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})
	}

	return outputs, nil
}

func GenerateDefinitions(context *common.Context, settings schema.Generator, s *schema.Schema) ([]byte, error) {
	context.TemplateFuncs["DefineSQLTable"] = DefineTable
	context.TemplateFuncs["DefineSQLSchema"] = DefineSchema
	context.TemplateFuncs["DefineSQLExtensions"] = DefineExtensions
	context.TemplateFuncs["DefineSQLTypes"] = DefineTypes
	context.TemplateFuncs["DefineSQLSequnce"] = DefineSequence

	temp := template.New("create_statement.tmpl").Funcs(context.TemplateFuncs)
	if _, err := temp.Parse(CreateStatementTemplate); err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	data := struct {
		Context  *common.Context
		Schema   *schema.Schema
		Settings schema.Generator
	}{
		Context:  context,
		Schema:   s,
		Settings: settings,
	}
	if err := temp.ExecuteTemplate(&buf, "create_statement.tmpl", data); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func clean(b []byte) []byte {
	b = writers.NewLinesRegex.ReplaceAll(b, []byte(""))

	// convert tabs to 4 spaces
	b = bytes.Replace(b, []byte("\t"), []byte("    "), -1)

	// clean extra spaces
	b = bytes.Replace(b, []byte("  ,"), []byte(","), -1)
	b = bytes.Replace(b, []byte(" ,"), []byte(","), -1)

	// replace last trailing comma
	b = bytes.Replace(b, []byte(",\n)"), []byte("\n)"), -1)

	return b
}

// CreateStatementTemplate holds the template for the create sql statement generator
var CreateStatementTemplate = `{{DefineSQLSchema .Context .Settings .Schema}}

{{DefineSQLSequnce .Context .Settings .Schema}}

{{DefineSQLExtensions .Context .Settings .Schema}}

{{DefineSQLTypes .Context .Settings .Schema}}

{{DefineSQLTable .Context .Settings .Schema}}
`

func GetFieldNameFunc(name string) func(string) string {
	switch name {
	case "lower":
		return strings.ToLower
	default:
		return stringext.ToFieldName
	}
}

func (g *Generator) generateSettings(moduleName string, s *schema.Schema) schema.Generator {
	settings, ok := s.Generators.Get(g.Name())
	if !ok {
		settings = schema.Generator{}
	}
	settings.SetNX("databaseName", stringext.ToFieldName(moduleName))
	settings.SetNX("schemaName", stringext.ToFieldName(moduleName))
	settings.SetNX("tableName", stringext.ToFieldName(s.Title))
	settings.SetNX("roleName", stringext.ToFieldName(moduleName))

	// convert []interface to []string
	grants := settings.GetWithDefault("grants", []string{"ALL"})
	grantsI, ok := grants.([]interface{})
	grantsS := make([]string, 0)

	if ok {
		for _, t := range grantsI {
			grantsS = append(grantsS, t.(string))
		}
	} else {
		grantsS = grants.([]string)
	}

	settings.Set("grants", grantsS)

	return settings
}

func (g *Generator) setDefaultSettings(defaultSettings schema.Generator, s *schema.Schema) schema.Generator {
	settings, _ := s.Generators.Get(g.Name())

	settings.SetNX("databaseName", defaultSettings.Get("databaseName").(string))
	settings.SetNX("schemaName", defaultSettings.Get("schemaName").(string))
	settings.SetNX("tableName", defaultSettings.Get("tableName").(string))
	settings.SetNX("roleName", defaultSettings.Get("roleName").(string))

	// convert []interface to []string
	grants := settings.GetWithDefault("grants", defaultSettings.Get("grants").([]string))
	grantsI, ok := grants.([]interface{})
	grantsS := make([]string, 0)

	if ok {
		for _, t := range grantsI {
			grantsS = append(grantsS, t.(string))
		}
	} else {
		grantsS = grants.([]string)
	}

	settings.Set("grants", grantsS)

	return settings
}
<commit_msg>Gene: move helper function to bottom<commit_after>package definitions

import (
	"bytes"
	"errors"
	"fmt"
	"strings"
	"text/template"

	"github.com/cihangir/gene/generators/common"
	"github.com/cihangir/gene/writers"
	"github.com/cihangir/schema"
	"github.com/cihangir/stringext"
)

const generatorName = "sql-definition"

type Generator struct {
	DatabaseName  string
	SchemaName    string
	TableName     string
	RoleName      string
	FieldNameCase string `default:"snake"`
}

func New() *Generator {
	return &Generator{}
}

func (g *Generator) Name() string {
	return generatorName
}

// Generate generates the basic CRUD statements for the models
func (g *Generator) Generate(context *common.Context, s *schema.Schema) ([]common.Output, error) {
	outputs := make([]common.Output, 0)

	if s.Title == "" {
		return outputs, errors.New("Title should be set")
	}

	moduleName := context.FieldNameFunc(s.Title)

	settings := g.generateSettings(moduleName, s)

	for _, name := range schema.SortedKeys(s.Definitions) {
		def := s.Definitions[name]

		// schema should have our generator
		if !def.Generators.Has(g.Name()) {
			continue
		}

		settingsDef := g.setDefaultSettings(settings, def)
		settingsDef.Set("tableName", context.FieldNameFunc(def.Title))

		//
		// generate roles
		//
		role, err := DefineRole(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content:     role,
			Path:        fmt.Sprintf("%s/001-%s_roles.sql", context.Config.Target, settingsDef.Get("databaseName").(string)),
			DoNotFormat: true,
		})

		//
		// generate database
		//
		db, err := DefineDatabase(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content:     db,
			Path:        fmt.Sprintf("%s/002-%s_database.sql", context.Config.Target, settingsDef.Get("databaseName").(string)),
			DoNotFormat: true,
		})

		//
		// generate extenstions
		//
		extenstions, err := DefineExtensions(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: extenstions,
			Path: fmt.Sprintf(
				"%s/003-%s_extensions.sql",
				context.Config.Target,
				settingsDef.Get("databaseName").(string)),
			DoNotFormat: true,
		})

		//
		// generate schema
		//
		sc, err := DefineSchema(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: sc,
			Path: fmt.Sprintf(
				"%s/%s/004-schema.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate sequences
		//
		sequence, err := DefineSequence(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: sequence,
			Path: fmt.Sprintf(
				"%s/%s/005-%s-sequence.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate types
		//
		types, err := DefineTypes(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: types,
			Path: fmt.Sprintf(
				"%s/%s/006-%s-types.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate tables
		//
		table, err := DefineTable(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: table,
			Path: fmt.Sprintf(
				"%s/%s/007-%s-table.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})

		//
		// generate constraints
		//
		constraints, err := DefineConstraints(context, settingsDef, def)
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: constraints,
			Path: fmt.Sprintf(
				"%s/%s/007-%s-constraints.sql",
				context.Config.Target,
				settingsDef.Get("schemaName").(string),
				settingsDef.Get("tableName").(string),
			),
			DoNotFormat: true,
		})
	}

	return outputs, nil
}

func GenerateDefinitions(context *common.Context, settings schema.Generator, s *schema.Schema) ([]byte, error) {
	context.TemplateFuncs["DefineSQLTable"] = DefineTable
	context.TemplateFuncs["DefineSQLSchema"] = DefineSchema
	context.TemplateFuncs["DefineSQLExtensions"] = DefineExtensions
	context.TemplateFuncs["DefineSQLTypes"] = DefineTypes
	context.TemplateFuncs["DefineSQLSequnce"] = DefineSequence

	temp := template.New("create_statement.tmpl").Funcs(context.TemplateFuncs)
	if _, err := temp.Parse(CreateStatementTemplate); err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	data := struct {
		Context  *common.Context
		Schema   *schema.Schema
		Settings schema.Generator
	}{
		Context:  context,
		Schema:   s,
		Settings: settings,
	}
	if err := temp.ExecuteTemplate(&buf, "create_statement.tmpl", data); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// CreateStatementTemplate holds the template for the create sql statement generator
var CreateStatementTemplate = `{{DefineSQLSchema .Context .Settings .Schema}}

{{DefineSQLSequnce .Context .Settings .Schema}}

{{DefineSQLExtensions .Context .Settings .Schema}}

{{DefineSQLTypes .Context .Settings .Schema}}

{{DefineSQLTable .Context .Settings .Schema}}
`

func GetFieldNameFunc(name string) func(string) string {
	switch name {
	case "lower":
		return strings.ToLower
	default:
		return stringext.ToFieldName
	}
}

func (g *Generator) generateSettings(moduleName string, s *schema.Schema) schema.Generator {
	settings, ok := s.Generators.Get(g.Name())
	if !ok {
		settings = schema.Generator{}
	}
	settings.SetNX("databaseName", stringext.ToFieldName(moduleName))
	settings.SetNX("schemaName", stringext.ToFieldName(moduleName))
	settings.SetNX("tableName", stringext.ToFieldName(s.Title))
	settings.SetNX("roleName", stringext.ToFieldName(moduleName))

	// convert []interface to []string
	grants := settings.GetWithDefault("grants", []string{"ALL"})
	grantsI, ok := grants.([]interface{})
	grantsS := make([]string, 0)

	if ok {
		for _, t := range grantsI {
			grantsS = append(grantsS, t.(string))
		}
	} else {
		grantsS = grants.([]string)
	}

	settings.Set("grants", grantsS)

	return settings
}

func (g *Generator) setDefaultSettings(defaultSettings schema.Generator, s *schema.Schema) schema.Generator {
	settings, _ := s.Generators.Get(g.Name())

	settings.SetNX("databaseName", defaultSettings.Get("databaseName").(string))
	settings.SetNX("schemaName", defaultSettings.Get("schemaName").(string))
	settings.SetNX("tableName", defaultSettings.Get("tableName").(string))
	settings.SetNX("roleName", defaultSettings.Get("roleName").(string))

	// convert []interface to []string
	grants := settings.GetWithDefault("grants", defaultSettings.Get("grants").([]string))
	grantsI, ok := grants.([]interface{})
	grantsS := make([]string, 0)

	if ok {
		for _, t := range grantsI {
			grantsS = append(grantsS, t.(string))
		}
	} else {
		grantsS = grants.([]string)
	}

	settings.Set("grants", grantsS)

	return settings
}

func clean(b []byte) []byte {
	b = writers.NewLinesRegex.ReplaceAll(b, []byte(""))

	// convert tabs to 4 spaces
	b = bytes.Replace(b, []byte("\t"), []byte("    "), -1)

	// clean extra spaces
	b = bytes.Replace(b, []byte("  ,"), []byte(","), -1)
	b = bytes.Replace(b, []byte(" ,"), []byte(","), -1)

	// replace last trailing comma
	b = bytes.Replace(b, []byte(",\n)"), []byte("\n)"), -1)

	return b
}
<|endoftext|>