<commit_before>package comm

import (
    "bufio"
    "errors"
    "fmt"
    ev "github.com/hhkbp2/rafted/event"
    logging "github.com/hhkbp2/rafted/logging"
    ps "github.com/hhkbp2/rafted/persist"
    "github.com/ugorji/go/codec"
    "io"
    "net"
    "sync"
)

const (
    DefaultTransportBufferSize = 16
)

type TransportChunk struct {
    Data     []byte
    SourceCh chan []byte
}

type MemoryServerTransport struct {
    addr       net.Addr
    ConsumeCh  chan *TransportChunk
    ResponseCh chan []byte
    register   *MemoryTransportRegister
}

func NewMemoryServerTransport(
    addr net.Addr, register *MemoryTransportRegister) *MemoryServerTransport {

    return &MemoryServerTransport{
        addr:       addr,
        ConsumeCh:  make(chan *TransportChunk, DefaultTransportBufferSize),
        ResponseCh: make(chan []byte),
        register:   register,
    }
}

func (self *MemoryServerTransport) Open() error {
    self.register.Register(self.addr.String(), self)
    return nil
}

func (self *MemoryServerTransport) ReadChunk() (*TransportChunk, error) {
    chunk := <-self.ConsumeCh
    return chunk, nil
}

func (self *MemoryServerTransport) WriteChunk(chunk *TransportChunk) error {
    self.ConsumeCh <- chunk
    return nil
}

func (self *MemoryServerTransport) Read(p []byte) (int, error) {
    chunk, _ := self.ReadChunk()
    p = chunk.Data
    return len(p), nil
}

func (self *MemoryServerTransport) Write(p []byte) (int, error) {
    self.ResponseCh <- p
    return len(p), nil
}

func (self *MemoryServerTransport) Close() error {
    return self.register.Unregister(self.addr.String())
}

func (self *MemoryServerTransport) Addr() net.Addr {
    return self.addr
}

type MemoryTransportRegister struct {
    sync.RWMutex
    transports map[string]*MemoryServerTransport
}

func NewMemoryTransportRegister() *MemoryTransportRegister {
    return &MemoryTransportRegister{
        transports: make(map[string]*MemoryServerTransport),
    }
}

func (self *MemoryTransportRegister) Register(
    id string, transport *MemoryServerTransport) {

    self.Lock()
    defer self.Unlock()
    self.transports[id] = transport
}

func (self *MemoryTransportRegister) Unregister(id string) error {
    self.Lock()
    defer self.Unlock()
    if _, ok := self.transports[id]; ok {
        delete(self.transports, id)
        return nil
    }
    return errors.New(fmt.Sprintf("no transport for id: %v", id))
}

func (self *MemoryTransportRegister) Get(
    id string) (transport *MemoryServerTransport, ok bool) {

    self.RLock()
    defer self.RUnlock()
    transport, ok = self.transports[id]
    return transport, ok
}

type MemoryTransport struct {
    addr      net.Addr
    consumeCh chan []byte
    register  *MemoryTransportRegister
    peer      *MemoryServerTransport
}

func NewMemoryTransport(addr net.Addr, register *MemoryTransportRegister) *MemoryTransport {
    return &MemoryTransport{
        addr:      addr,
        consumeCh: make(chan []byte, DefaultTransportBufferSize),
        register:  register,
    }
}

func (self *MemoryTransport) PeerAddr() net.Addr {
    return self.addr
}

func (self *MemoryTransport) Open() error {
    if transport, ok := self.register.Get(self.addr.String()); ok {
        self.peer = transport
        return nil
    }
    return errors.New(
        fmt.Sprintf("no server transport for id: %v", self.addr.String()))
}

func (self *MemoryTransport) Close() error {
    // empty body
    return nil
}

func (self *MemoryTransport) Read(b []byte) (int, error) {
    b = <-self.consumeCh
    return len(b), nil
}

func (self *MemoryTransport) Write(b []byte) (int, error) {
    chunk := &TransportChunk{b, self.consumeCh}
    if err := self.peer.WriteChunk(chunk); err != nil {
        return len(b), err
    }
    return len(b), nil
}

type MemoryConnection struct {
    *MemoryTransport
    reader  *bufio.Reader
    writer  *bufio.Writer
    encoder Encoder
    decoder Decoder
}

func NewMemoryConnection(
    addr net.Addr, register *MemoryTransportRegister) *MemoryConnection {

    conn := &MemoryConnection{
        MemoryTransport: NewMemoryTransport(addr, register),
    }
    conn.reader = bufio.NewReader(conn.MemoryTransport)
    conn.writer = bufio.NewWriter(conn.MemoryTransport)
    conn.decoder = codec.NewDecoder(conn.reader, &codec.MsgpackHandle{})
    conn.encoder = codec.NewEncoder(conn.writer, &codec.MsgpackHandle{})
    return conn
}

func (self *MemoryConnection) CallRPC(
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    if err := WriteEvent(self.writer, self.encoder, request); err != nil {
        self.Close()
        return nil, err
    }

    event, err := ReadResponse(self.reader, self.decoder)
    if err != nil {
        self.Close()
        return nil, err
    }
    return event, nil
}

type MemoryClient struct {
    connectionPool     map[string][]*MemoryConnection
    connectionPoolLock sync.Mutex

    poolSize int
    register *MemoryTransportRegister
}

func NewMemoryClient(
    poolSize int, register *MemoryTransportRegister) *MemoryClient {

    return &MemoryClient{
        connectionPool: make(map[string][]*MemoryConnection),
        poolSize:       poolSize,
        register:       register,
    }
}

func (self *MemoryClient) CallRPCTo(
    target net.Addr,
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    connection, err := self.getConnection(target)
    if err != nil {
        return nil, err
    }

    response, err = connection.CallRPC(request)
    if err == nil {
        self.returnConnectionToPool(connection)
    }
    return response, err
}

func (self *MemoryClient) CloseAll(target net.Addr) error {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := target.String()
    connections, ok := self.connectionPool[key]
    if ok {
        var err error
        for _, connection := range connections {
            err = connection.Close()
        }
        delete(self.connectionPool, key)
        return err
    }
    return nil
}

func (self *MemoryClient) getConnectionFromPool(
    target net.Addr) (*MemoryConnection, error) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := target.String()
    connections, ok := self.connectionPool[key]

    if !ok || len(connections) == 0 {
        return nil, errors.New("no connection for this target")
    }

    connection := connections[len(connections)-1]
    self.connectionPool[key] = connections[:len(connections)-1]
    return connection, nil
}

func (self *MemoryClient) returnConnectionToPool(
    connection *MemoryConnection) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := connection.PeerAddr().String()
    connections, ok := self.connectionPool[key]
    if !ok {
        connections = make([]*MemoryConnection, 0)
        self.connectionPool[key] = connections
    }

    if len(connections) < self.poolSize {
        self.connectionPool[key] = append(connections, connection)
    } else {
        connection.Close()
    }
}

func (self *MemoryClient) getConnection(
    target net.Addr) (*MemoryConnection, error) {

    connection, err := self.getConnectionFromPool(target)
    if connection != nil && err == nil {
        return connection, nil
    }

    connection = NewMemoryConnection(target, self.register)
    if err := connection.Open(); err != nil {
        return nil, err
    }

    return connection, nil
}

type MemoryServer struct {
    transport           *MemoryServerTransport
    acceptedConnections map[chan []byte]*MemoryServerTransport

    eventHandler func(ev.RaftRequestEvent)
    register     *MemoryTransportRegister
    logger       logging.Logger
}

func NewMemoryServer(
    bindAddr net.Addr,
    eventHandler func(ev.RaftRequestEvent),
    register *MemoryTransportRegister,
    logger logging.Logger) *MemoryServer {

    transport := NewMemoryServerTransport(bindAddr, register)
    return &MemoryServer{
        transport:    transport,
        eventHandler: eventHandler,
        register:     register,
        logger:       logger,
    }
}

func (self *MemoryServer) Serve() {
    for {
        chunk, err := self.transport.ReadChunk()
        if err != nil {
            // TODO add log
            continue
        }
        transport, ok := self.acceptedConnections[chunk.SourceCh]
        if ok {
            // connection already accepted

            continue
        } else {
            addr := ps.ServerAddr{
                Protocol: self.transport.Addr().Network(),
                IP: fmt.Sprintf(
                    "%s.%#v", self.transport.Addr().String(), chunk.SourceCh),
            }
            transport := NewMemoryServerTransport(&addr, self.register)
            transport.ResponseCh = chunk.SourceCh
            self.acceptedConnections[chunk.SourceCh] = transport
            go self.handleConn(transport)
        }
        transport.WriteChunk(chunk)
    }
}

func (self *MemoryServer) handleConn(
    transport *MemoryServerTransport) {

    defer func() {
        transport.Close()
        delete(self.acceptedConnections, transport.ResponseCh)
    }()
    reader := bufio.NewReader(transport)
    writer := bufio.NewWriter(transport)
    decoder := codec.NewDecoder(reader, &codec.MsgpackHandle{})
    encoder := codec.NewEncoder(writer, &codec.MsgpackHandle{})

    for {
        if err := self.handleCommand(
            reader, writer, decoder, encoder); err != nil {

            if err != io.EOF {
                // TODO add log
            }
            return
        }
        if err := writer.Flush(); err != nil {
            // TODO add log
            return
        }
    }
}

func (self *MemoryServer) handleCommand(
    reader *bufio.Reader,
    writer *bufio.Writer,
    decoder Decoder,
    encoder Encoder) error {

    event, err := ReadRequest(reader, decoder)
    if err != nil {
        return err
    }

    self.eventHandler(event)
    response := event.RecvResponse()
    if err := WriteEvent(writer, encoder, response); err != nil {
        return err
    }
    return nil
}
<commit_msg>add log for memory comm<commit_after>package comm

import (
    "bufio"
    "errors"
    "fmt"
    ev "github.com/hhkbp2/rafted/event"
    logging "github.com/hhkbp2/rafted/logging"
    ps "github.com/hhkbp2/rafted/persist"
    "github.com/ugorji/go/codec"
    "io"
    "net"
    "sync"
)

const (
    DefaultTransportBufferSize = 16
)

type TransportChunk struct {
    Data     []byte
    SourceCh chan []byte
}

type MemoryServerTransport struct {
    addr       net.Addr
    ConsumeCh  chan *TransportChunk
    ResponseCh chan []byte
    register   *MemoryTransportRegister
}

func NewMemoryServerTransport(
    addr net.Addr, register *MemoryTransportRegister) *MemoryServerTransport {

    return &MemoryServerTransport{
        addr:       addr,
        ConsumeCh:  make(chan *TransportChunk, DefaultTransportBufferSize),
        ResponseCh: make(chan []byte),
        register:   register,
    }
}

func (self *MemoryServerTransport) Open() error {
    self.register.Register(self.addr.String(), self)
    return nil
}

func (self *MemoryServerTransport) ReadChunk() (*TransportChunk, error) {
    chunk := <-self.ConsumeCh
    return chunk, nil
}

func (self *MemoryServerTransport) WriteChunk(chunk *TransportChunk) error {
    self.ConsumeCh <- chunk
    return nil
}

func (self *MemoryServerTransport) Read(p []byte) (int, error) {
    chunk, _ := self.ReadChunk()
    p = chunk.Data
    return len(p), nil
}

func (self *MemoryServerTransport) Write(p []byte) (int, error) {
    self.ResponseCh <- p
    return len(p), nil
}

func (self *MemoryServerTransport) Close() error {
    return self.register.Unregister(self.addr.String())
}

func (self *MemoryServerTransport) Addr() net.Addr {
    return self.addr
}

type MemoryTransportRegister struct {
    sync.RWMutex
    transports map[string]*MemoryServerTransport
}

func NewMemoryTransportRegister() *MemoryTransportRegister {
    return &MemoryTransportRegister{
        transports: make(map[string]*MemoryServerTransport),
    }
}

func (self *MemoryTransportRegister) Register(
    id string, transport *MemoryServerTransport) {

    self.Lock()
    defer self.Unlock()
    self.transports[id] = transport
}

func (self *MemoryTransportRegister) Unregister(id string) error {
    self.Lock()
    defer self.Unlock()
    if _, ok := self.transports[id]; ok {
        delete(self.transports, id)
        return nil
    }
    return errors.New(fmt.Sprintf("no transport for id: %v", id))
}

func (self *MemoryTransportRegister) Get(
    id string) (transport *MemoryServerTransport, ok bool) {

    self.RLock()
    defer self.RUnlock()
    transport, ok = self.transports[id]
    return transport, ok
}

type MemoryTransport struct {
    addr      net.Addr
    consumeCh chan []byte
    register  *MemoryTransportRegister
    peer      *MemoryServerTransport
}

func NewMemoryTransport(addr net.Addr, register *MemoryTransportRegister) *MemoryTransport {
    return &MemoryTransport{
        addr:      addr,
        consumeCh: make(chan []byte, DefaultTransportBufferSize),
        register:  register,
    }
}

func (self *MemoryTransport) PeerAddr() net.Addr {
    return self.addr
}

func (self *MemoryTransport) Open() error {
    if transport, ok := self.register.Get(self.addr.String()); ok {
        self.peer = transport
        return nil
    }
    return errors.New(
        fmt.Sprintf("no server transport for id: %v", self.addr.String()))
}

func (self *MemoryTransport) Close() error {
    // empty body
    return nil
}

func (self *MemoryTransport) Read(b []byte) (int, error) {
    b = <-self.consumeCh
    return len(b), nil
}

func (self *MemoryTransport) Write(b []byte) (int, error) {
    chunk := &TransportChunk{b, self.consumeCh}
    if err := self.peer.WriteChunk(chunk); err != nil {
        return len(b), err
    }
    return len(b), nil
}

type MemoryConnection struct {
    *MemoryTransport
    reader  *bufio.Reader
    writer  *bufio.Writer
    encoder Encoder
    decoder Decoder
}

func NewMemoryConnection(
    addr net.Addr, register *MemoryTransportRegister) *MemoryConnection {

    conn := &MemoryConnection{
        MemoryTransport: NewMemoryTransport(addr, register),
    }
    conn.reader = bufio.NewReader(conn.MemoryTransport)
    conn.writer = bufio.NewWriter(conn.MemoryTransport)
    conn.decoder = codec.NewDecoder(conn.reader, &codec.MsgpackHandle{})
    conn.encoder = codec.NewEncoder(conn.writer, &codec.MsgpackHandle{})
    return conn
}

func (self *MemoryConnection) CallRPC(
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    if err := WriteEvent(self.writer, self.encoder, request); err != nil {
        self.Close()
        return nil, err
    }

    event, err := ReadResponse(self.reader, self.decoder)
    if err != nil {
        self.Close()
        return nil, err
    }
    return event, nil
}

type MemoryClient struct {
    connectionPool     map[string][]*MemoryConnection
    connectionPoolLock sync.Mutex

    poolSize int
    register *MemoryTransportRegister
}

func NewMemoryClient(
    poolSize int, register *MemoryTransportRegister) *MemoryClient {

    return &MemoryClient{
        connectionPool: make(map[string][]*MemoryConnection),
        poolSize:       poolSize,
        register:       register,
    }
}

func (self *MemoryClient) CallRPCTo(
    target net.Addr,
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    connection, err := self.getConnection(target)
    if err != nil {
        return nil, err
    }

    response, err = connection.CallRPC(request)
    if err == nil {
        self.returnConnectionToPool(connection)
    }
    return response, err
}

func (self *MemoryClient) CloseAll(target net.Addr) error {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := target.String()
    connections, ok := self.connectionPool[key]
    if ok {
        var err error
        for _, connection := range connections {
            err = connection.Close()
        }
        delete(self.connectionPool, key)
        return err
    }
    return nil
}

func (self *MemoryClient) getConnectionFromPool(
    target net.Addr) (*MemoryConnection, error) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := target.String()
    connections, ok := self.connectionPool[key]

    if !ok || len(connections) == 0 {
        return nil, errors.New("no connection for this target")
    }

    connection := connections[len(connections)-1]
    self.connectionPool[key] = connections[:len(connections)-1]
    return connection, nil
}

func (self *MemoryClient) returnConnectionToPool(
    connection *MemoryConnection) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := connection.PeerAddr().String()
    connections, ok := self.connectionPool[key]
    if !ok {
        connections = make([]*MemoryConnection, 0)
        self.connectionPool[key] = connections
    }

    if len(connections) < self.poolSize {
        self.connectionPool[key] = append(connections, connection)
    } else {
        connection.Close()
    }
}

func (self *MemoryClient) getConnection(
    target net.Addr) (*MemoryConnection, error) {

    connection, err := self.getConnectionFromPool(target)
    if connection != nil && err == nil {
        return connection, nil
    }

    connection = NewMemoryConnection(target, self.register)
    if err := connection.Open(); err != nil {
        return nil, err
    }

    return connection, nil
}

type MemoryServer struct {
    transport           *MemoryServerTransport
    acceptedConnections map[chan []byte]*MemoryServerTransport

    eventHandler func(ev.RaftRequestEvent)
    register     *MemoryTransportRegister
    logger       logging.Logger
}

func NewMemoryServer(
    bindAddr net.Addr,
    eventHandler func(ev.RaftRequestEvent),
    register *MemoryTransportRegister,
    logger logging.Logger) *MemoryServer {

    transport := NewMemoryServerTransport(bindAddr, register)
    return &MemoryServer{
        transport:    transport,
        eventHandler: eventHandler,
        register:     register,
        logger:       logger,
    }
}

func (self *MemoryServer) Serve() {
    for {
        chunk, err := self.transport.ReadChunk()
        if err != nil {
            self.logger.Error(
                "memory server fails to read first chunk, error: %s", err)
            continue
        }
        transport, ok := self.acceptedConnections[chunk.SourceCh]
        if ok {
            // connection already accepted
            continue
        } else {
            addr := ps.ServerAddr{
                Protocol: self.transport.Addr().Network(),
                IP: fmt.Sprintf(
                    "%s.%#v", self.transport.Addr().String(), chunk.SourceCh),
            }
            transport := NewMemoryServerTransport(&addr, self.register)
            transport.ResponseCh = chunk.SourceCh
            self.acceptedConnections[chunk.SourceCh] = transport
            go self.handleConn(transport)
        }
        transport.WriteChunk(chunk)
    }
}

func (self *MemoryServer) handleConn(
    transport *MemoryServerTransport) {

    defer func() {
        transport.Close()
        delete(self.acceptedConnections, transport.ResponseCh)
    }()
    reader := bufio.NewReader(transport)
    writer := bufio.NewWriter(transport)
    decoder := codec.NewDecoder(reader, &codec.MsgpackHandle{})
    encoder := codec.NewEncoder(writer, &codec.MsgpackHandle{})

    for {
        if err := self.handleCommand(
            reader, writer, decoder, encoder); err != nil {

            if err != io.EOF {
                self.logger.Error(
                    "memory server fails to handle command, error: %s", err)
            }
            return
        }
        if err := writer.Flush(); err != nil {
            self.logger.Error(
                "memory server fails to flush writer, error: %s", err)
            return
        }
    }
}

func (self *MemoryServer) handleCommand(
    reader *bufio.Reader,
    writer *bufio.Writer,
    decoder Decoder,
    encoder Encoder) error {

    event, err := ReadRequest(reader, decoder)
    if err != nil {
        return err
    }

    self.eventHandler(event)
    response := event.RecvResponse()
    if err := WriteEvent(writer, encoder, response); err != nil {
        return err
    }
    return nil
}
<|endoftext|>


<commit_before>package command

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	// enforceIndexRegex is a regular expression which extracts the enforcement error
	enforceIndexRegex = regexp.MustCompile(`\((Enforcing job modify index.*)\)`)
)

type RunCommand struct {
	Meta
	JobGetter
}

func (c *RunCommand) Help() string {
	helpText := `
Usage: nomad run [options] <path>

  Starts running a new job or updates an existing job using
  the specification located at <path>. This is the main command
  used to interact with Nomad.

  If the supplied path is "-", the jobfile is read from stdin. Otherwise
  it is read from the file at the supplied path or downloaded and
  read from URL specified.

  Upon successful job submission, this command will immediately
  enter an interactive monitor. This is useful to watch Nomad's
  internals make scheduling decisions and place the submitted work
  onto nodes. The monitor will end once job placement is done. It
  is safe to exit the monitor early using ctrl+c.

  On successful job submission and scheduling, exit code 0 will be
  returned. If there are job placement issues encountered
  (unsatisfiable constraints, resource exhaustion, etc), then the
  exit code will be 2. Any other errors, including client connection
  issues or internal errors, are indicated by exit code 1.

  If the job has specified the region, the -region flag and NOMAD_REGION
  environment variable are overridden and the the job's region is used.

  The run command will set the vault_token of the job based on the following
  precedence, going from highest to lowest: the -vault-token flag, the
  $VAULT_TOKEN environment variable and finally the value in the job file.

General Options:

  ` + generalOptionsUsage() + `

Run Options:

  -check-index
    If set, the job is only registered or updated if the the passed
    job modify index matches the server side version. If a check-index value of
    zero is passed, the job is only registered if it does not yet exist. If a
    non-zero value is passed, it ensures that the job is being updated from a
    known state. The use of this flag is most common in conjunction with plan
    command.

  -detach
    Return immediately instead of entering monitor mode. After job submission,
    the evaluation ID will be printed to the screen, which can be used to
    examine the evaluation using the eval-status command.

  -verbose
    Display full information.

  -vault-token
    If set, the passed Vault token is stored in the job before sending to the
    Nomad servers. This allows passing the Vault token without storing it in
    the job file. This overrides the token found in $VAULT_TOKEN environment
    variable and that found in the job.

  -output
    Output the JSON that would be submitted to the HTTP API without submitting
    the job.
`
	return strings.TrimSpace(helpText)
}

func (c *RunCommand) Synopsis() string {
	return "Run a new job or update an existing job"
}

func (c *RunCommand) Run(args []string) int {
	var detach, verbose, output bool
	var checkIndexStr, vaultToken string

	flags := c.Meta.FlagSet("run", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&detach, "detach", false, "")
	flags.BoolVar(&verbose, "verbose", false, "")
	flags.BoolVar(&output, "output", false, "")
	flags.StringVar(&checkIndexStr, "check-index", "", "")
	flags.StringVar(&vaultToken, "vault-token", "", "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Truncate the id unless full length is requested
	length := shortId
	if verbose {
		length = fullId
	}

	// Check that we got exactly one argument
	args = flags.Args()
	if len(args) != 1 {
		c.Ui.Error(c.Help())
		return 1
	}

	// Check that we got exactly one node
	args = flags.Args()
	if len(args) != 1 {
		c.Ui.Error(c.Help())
		return 1
	}

	// Get Job struct from Jobfile
	job, err := c.JobGetter.StructJob(args[0])
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error getting job struct: %s", err))
		return 1
	}

	// Initialize any fields that need to be.
	job.Canonicalize()

	// Check that the job is valid
	if err := job.Validate(); err != nil {
		c.Ui.Error(fmt.Sprintf("Error validating job: %v", err))
		return 1
	}

	// Check if the job is periodic.
	periodic := job.IsPeriodic()

	// Parse the Vault token
	if vaultToken == "" {
		// Check the environment variable
		vaultToken = os.Getenv("VAULT_TOKEN")
	}

	if vaultToken != "" {
		job.VaultToken = vaultToken
	}

	// Convert it to something we can use
	apiJob, err := convertStructJob(job)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error converting job: %s", err))
		return 1
	}

	if output {
		req := api.RegisterJobRequest{Job: apiJob}
		buf, err := json.MarshalIndent(req, "", "    ")
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error converting job: %s", err))
			return 1
		}

		c.Ui.Output(string(buf))
		return 0
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// Force the region to be that of the job.
	if r := job.Region; r != "" {
		client.SetRegion(r)
	}

	// Parse the check-index
	checkIndex, enforce, err := parseCheckIndex(checkIndexStr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing check-index value %q: %v", checkIndexStr, err))
		return 1
	}

	// Submit the job
	var evalID string
	if enforce {
		evalID, _, err = client.Jobs().EnforceRegister(apiJob, checkIndex, nil)
	} else {
		evalID, _, err = client.Jobs().Register(apiJob, nil)
	}
	if err != nil {
		if strings.Contains(err.Error(), api.RegisterEnforceIndexErrPrefix) {
			// Format the error specially if the error is due to index
			// enforcement
			matches := enforceIndexRegex.FindStringSubmatch(err.Error())
			if len(matches) == 2 {
				c.Ui.Error(matches[1]) // The matched group
				c.Ui.Error("Job not updated")
				return 1
			}
		}

		c.Ui.Error(fmt.Sprintf("Error submitting job: %s", err))
		return 1
	}

	// Check if we should enter monitor mode
	if detach || periodic {
		c.Ui.Output("Job registration successful")
		if periodic {
			now := time.Now().UTC()
			next := job.Periodic.Next(now)
			c.Ui.Output(fmt.Sprintf("Approximate next launch time: %s (%s from now)",
				formatTime(next), formatTimeDifference(now, next, time.Second)))
		} else {
			c.Ui.Output("Evaluation ID: " + evalID)
		}

		return 0
	}

	// Detach was not specified, so start monitoring
	mon := newMonitor(c.Ui, client, length)
	return mon.monitor(evalID, false)

}

// parseCheckIndex parses the check-index flag and returns the index, whether it
// was set and potentially an error during parsing.
func parseCheckIndex(input string) (uint64, bool, error) {
	if input == "" {
		return 0, false, nil
	}

	u, err := strconv.ParseUint(input, 10, 64)
	return u, true, err
}

// convertStructJob is used to take a *structs.Job and convert it to an *api.Job.
// This function is just a hammer and probably needs to be revisited.
func convertStructJob(in *structs.Job) (*api.Job, error) {
	gob.Register([]map[string]interface{}{})
	gob.Register([]interface{}{})
	var apiJob *api.Job
	buf := new(bytes.Buffer)
	if err := gob.NewEncoder(buf).Encode(in); err != nil {
		return nil, err
	}
	if err := gob.NewDecoder(buf).Decode(&apiJob); err != nil {
		return nil, err
	}
	return apiJob, nil
}
<commit_msg>add import<commit_after>package command

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	// enforceIndexRegex is a regular expression which extracts the enforcement error
	enforceIndexRegex = regexp.MustCompile(`\((Enforcing job modify index.*)\)`)
)

type RunCommand struct {
	Meta
	JobGetter
}

func (c *RunCommand) Help() string {
	helpText := `
Usage: nomad run [options] <path>

  Starts running a new job or updates an existing job using
  the specification located at <path>. This is the main command
  used to interact with Nomad.

  If the supplied path is "-", the jobfile is read from stdin. Otherwise
  it is read from the file at the supplied path or downloaded and
  read from URL specified.

  Upon successful job submission, this command will immediately
  enter an interactive monitor. This is useful to watch Nomad's
  internals make scheduling decisions and place the submitted work
  onto nodes. The monitor will end once job placement is done. It
  is safe to exit the monitor early using ctrl+c.

  On successful job submission and scheduling, exit code 0 will be
  returned. If there are job placement issues encountered
  (unsatisfiable constraints, resource exhaustion, etc), then the
  exit code will be 2. Any other errors, including client connection
  issues or internal errors, are indicated by exit code 1.

  If the job has specified the region, the -region flag and NOMAD_REGION
  environment variable are overridden and the the job's region is used.

  The run command will set the vault_token of the job based on the following
  precedence, going from highest to lowest: the -vault-token flag, the
  $VAULT_TOKEN environment variable and finally the value in the job file.

General Options:

  ` + generalOptionsUsage() + `

Run Options:

  -check-index
    If set, the job is only registered or updated if the the passed
    job modify index matches the server side version. If a check-index value of
    zero is passed, the job is only registered if it does not yet exist. If a
    non-zero value is passed, it ensures that the job is being updated from a
    known state. The use of this flag is most common in conjunction with plan
    command.

  -detach
    Return immediately instead of entering monitor mode. After job submission,
    the evaluation ID will be printed to the screen, which can be used to
    examine the evaluation using the eval-status command.

  -verbose
    Display full information.

  -vault-token
    If set, the passed Vault token is stored in the job before sending to the
    Nomad servers. This allows passing the Vault token without storing it in
    the job file. This overrides the token found in $VAULT_TOKEN environment
    variable and that found in the job.

  -output
    Output the JSON that would be submitted to the HTTP API without submitting
    the job.
`
	return strings.TrimSpace(helpText)
}

func (c *RunCommand) Synopsis() string {
	return "Run a new job or update an existing job"
}

func (c *RunCommand) Run(args []string) int {
	var detach, verbose, output bool
	var checkIndexStr, vaultToken string

	flags := c.Meta.FlagSet("run", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&detach, "detach", false, "")
	flags.BoolVar(&verbose, "verbose", false, "")
	flags.BoolVar(&output, "output", false, "")
	flags.StringVar(&checkIndexStr, "check-index", "", "")
	flags.StringVar(&vaultToken, "vault-token", "", "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Truncate the id unless full length is requested
	length := shortId
	if verbose {
		length = fullId
	}

	// Check that we got exactly one argument
	args = flags.Args()
	if len(args) != 1 {
		c.Ui.Error(c.Help())
		return 1
	}

	// Check that we got exactly one node
	args = flags.Args()
	if len(args) != 1 {
		c.Ui.Error(c.Help())
		return 1
	}

	// Get Job struct from Jobfile
	job, err := c.JobGetter.StructJob(args[0])
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error getting job struct: %s", err))
		return 1
	}

	// Initialize any fields that need to be.
	job.Canonicalize()

	// Check that the job is valid
	if err := job.Validate(); err != nil {
		c.Ui.Error(fmt.Sprintf("Error validating job: %v", err))
		return 1
	}

	// Check if the job is periodic.
	periodic := job.IsPeriodic()

	// Parse the Vault token
	if vaultToken == "" {
		// Check the environment variable
		vaultToken = os.Getenv("VAULT_TOKEN")
	}

	if vaultToken != "" {
		job.VaultToken = vaultToken
	}

	// Convert it to something we can use
	apiJob, err := convertStructJob(job)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error converting job: %s", err))
		return 1
	}

	if output {
		req := api.RegisterJobRequest{Job: apiJob}
		buf, err := json.MarshalIndent(req, "", "    ")
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error converting job: %s", err))
			return 1
		}

		c.Ui.Output(string(buf))
		return 0
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// Force the region to be that of the job.
	if r := job.Region; r != "" {
		client.SetRegion(r)
	}

	// Parse the check-index
	checkIndex, enforce, err := parseCheckIndex(checkIndexStr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing check-index value %q: %v", checkIndexStr, err))
		return 1
	}

	// Submit the job
	var evalID string
	if enforce {
		evalID, _, err = client.Jobs().EnforceRegister(apiJob, checkIndex, nil)
	} else {
		evalID, _, err = client.Jobs().Register(apiJob, nil)
	}
	if err != nil {
		if strings.Contains(err.Error(), api.RegisterEnforceIndexErrPrefix) {
			// Format the error specially if the error is due to index
			// enforcement
			matches := enforceIndexRegex.FindStringSubmatch(err.Error())
			if len(matches) == 2 {
				c.Ui.Error(matches[1]) // The matched group
				c.Ui.Error("Job not updated")
				return 1
			}
		}

		c.Ui.Error(fmt.Sprintf("Error submitting job: %s", err))
		return 1
	}

	// Check if we should enter monitor mode
	if detach || periodic {
		c.Ui.Output("Job registration successful")
		if periodic {
			now := time.Now().UTC()
			next := job.Periodic.Next(now)
			c.Ui.Output(fmt.Sprintf("Approximate next launch time: %s (%s from now)",
				formatTime(next), formatTimeDifference(now, next, time.Second)))
		} else {
			c.Ui.Output("Evaluation ID: " + evalID)
		}

		return 0
	}

	// Detach was not specified, so start monitoring
	mon := newMonitor(c.Ui, client, length)
	return mon.monitor(evalID, false)

}

// parseCheckIndex parses the check-index flag and returns the index, whether it
// was set and potentially an error during parsing.
func parseCheckIndex(input string) (uint64, bool, error) {
	if input == "" {
		return 0, false, nil
	}

	u, err := strconv.ParseUint(input, 10, 64)
	return u, true, err
}

// convertStructJob is used to take a *structs.Job and convert it to an *api.Job.
// This function is just a hammer and probably needs to be revisited.
func convertStructJob(in *structs.Job) (*api.Job, error) {
	gob.Register([]map[string]interface{}{})
	gob.Register([]interface{}{})
	var apiJob *api.Job
	buf := new(bytes.Buffer)
	if err := gob.NewEncoder(buf).Encode(in); err != nil {
		return nil, err
	}
	if err := gob.NewDecoder(buf).Decode(&apiJob); err != nil {
		return nil, err
	}
	return apiJob, nil
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"os"
	"strings"
	"syscall"

	"github.com/ondevice/ondevice/config"
	"github.com/ondevice/ondevice/logg"
)

// option list copied from debian jessie's openssh source package (from ssh.c, line 509)
var sshFlags = sshParseFlags("1246ab:c:e:fgi:kl:m:no:p:qstvxD:L:NR:")

func sshRun(args []string) int {
	sshPath := "/usr/bin/ssh"

	// parse args (to detect the ones before 'user@host')
	target, args := sshParseArgs(args)
	tgtHost, tgtUser := sshParseTarget(target)

	// compose ProxyCommand
	// TODO this will fail miserably if argv[0] or tgtHost contain spaces
	proxyCmd := fmt.Sprintf("-oProxyCommand=%s pipe %%h ssh", os.Args[0])

	// create something like `ssh -oProxyCommand=... user@ondevice:devId <opts`
	a := make([]string, 0, 10)
	a = append(a, sshPath, proxyCmd)
	if tgtUser != "" {
		a = append(a, fmt.Sprintf("%s@ondevice:%s", tgtUser, tgtHost))
	} else {
		a = append(a, fmt.Sprintf("ondevice:%s", tgtHost))
	}
	a = append(a, args...)

	// syscall.Exec will replace this app with ssh (yes, replace it, not just launch)
	// therefore, unless there's an error, this is the last line of code to be executed
	err := syscall.Exec(sshPath, a, os.Environ())
	if err != nil {
		logg.Fatal("Failed to run ", sshPath, ": ", err)
	}

	// nothing here should ever be executed
	logg.Fatal("This should never happen")
	return -1
}

func sshParseArgs(args []string) (string, []string) {
	var target string
	var outArgs []string

	for i := 0; i < len(args); i++ {
		arg := args[i]

		if len(arg) == 0 {
			logg.Fatal("Failed to parse SSH arguments: got an empty argument while looking for '[user@]host'")
		} else if arg == "-" {
			logg.Fatal("Stray '-' SSH argument")
		} else if arg[0] == '-' {
			hasValue, ok := sshFlags[arg[1]]
			if !ok {
				logg.Fatal("Unsupported SSH argument: ", arg)
			}
			if hasValue && len(arg) == 2 {
				// the value's in the next argument, push both to outArgs
				outArgs = append(outArgs, arg, args[i+1])
				i++
			} else if hasValue && len(arg) > 2 {
				// the value's part of arg
				outArgs = append(outArgs, arg)
			} else if !hasValue && len(arg) > 2 { // && !hasValue
				logg.Fatal("Got value for flag that doesn't expect one: ", arg)
			} else if !hasValue && len(arg) == 2 {
				// a simple flag)
				outArgs = append(outArgs, arg)
			} else {
				// yay to defensive programming
				logg.Fatal("this should never happen (fifth state of two binary values)")
			}
		} else {
			// first non-option argument -> extract target and keep the rest as-is
			target = arg
			outArgs = append(outArgs, args[i+1:]...)
			return target, outArgs
		}
	}

	logg.Fatal("Missing SSH target user/host!")
	return "", nil
}

func sshParseTarget(target string) (tgtHost string, tgtUser string) {
	parts := strings.SplitN(target, "@", 2)
	if len(parts) == 1 {
		tgtUser = ""
		tgtHost = parts[0]
	} else if len(parts) == 2 {
		tgtUser = parts[0]
		tgtHost = parts[1]
	} else {
		logg.Fatal("SplitN(..., 2) returned odd number of results: ", len(parts))
	}

	// always use qualified device names
	if !strings.Contains(tgtHost, ".") {
		user, _, err := config.GetClientAuth()
		if err == nil {
			tgtHost = fmt.Sprintf("%s.%s", user, tgtHost)
		}
	}

	return tgtHost, tgtUser
}

func sshParseFlags(flags string) map[byte]bool {
	rc := map[byte]bool{}

	for i := 0; i < len(flags); i++ {
		flag := flags[i]
		hasValue := false
		if flags[i+1] == ':' {
			hasValue = true
			i++
		}
		rc[flag] = hasValue
	}
	return rc
}

// SSHCommand -- implements `ondevice ssh`
var SSHCommand = BaseCommand{
	Arguments: "[ssh-arguments...]",
	ShortHelp: "Connect to your devices using the ssh protocol",
	RunFn:     sshRun,
	LongHelp: `$ ondevice ssh [<user>@]<device> [ssh-arguments...]

Connect to your devices using the 'ssh' command.

This is a relatively thin wrapper around the 'ssh' command.
The main difference to invoking ssh directly is that instead of regular host names you'll have to specify an ondevice deviceId.
The connection is routed through the ondevice.io network.

ondevice ssh will try to parse ssh's arguments, the first non-argument has to be
the user@hostname combo.

See ssh's documentation for further details.

Examples:
- ondevice ssh device1
  simply connect to device1
- ondevice ssh user@device1
  open an SSH connection to device1, logging in as 'user'
- ondevice ssh device1 echo hello world
  run 'echo hello world' on device1
- ondevice ssh device1 -N -L 1234:localhost:80
  Tunnel the HTTP server on device1 to the local port 1234 without opening
  a shell
- ondevice ssh device1 -D 1080
  Starting a SOCKS5 proxy listening on port 1080. It'll redirect all traffic
  to the target host.
`,
}
<commit_msg>commands/ssh.go: updated sshParseArgs() to return opts and args (in preparation of scp support)<commit_after>package command

import (
	"fmt"
	"os"
	"strings"
	"syscall"

	"github.com/ondevice/ondevice/config"
	"github.com/ondevice/ondevice/logg"
)

// option list copied from debian jessie's openssh source package (from ssh.c, line 509)
var sshFlags = sshParseFlags("1246ab:c:e:fgi:kl:m:no:p:qstvxD:L:NR:")

func sshRun(args []string) int {
	sshPath := "/usr/bin/ssh"

	// parse args (to detect the ones before 'user@host')
	args, opts := sshParseArgs(sshFlags, args)
	if len(args) < 1 {
		logg.Fatal("missing target host")
	}

	// first non-option target is the [user@]host.
	tgtHost, tgtUser := sshParseTarget(args[0])
	args = args[1:]

	// compose ProxyCommand
	// TODO this will fail miserably if os.Args[0] contain spaces

	// compose the ssh command
	var a = make([]string, 0, 20)
	// ssh -oProxyCommand=ondevice pipe ssh %h ssh
	a = append(a, sshPath, fmt.Sprintf("-oProxyCommand=%s pipe %%h ssh", os.Args[0]))
	a = append(a, opts...) // ... ssh flags (-L -R -D ...)

	// target [user@]devId (qualified devId with 'ondevice:' prefix)
	if tgtUser != "" {
		a = append(a, fmt.Sprintf("%s@ondevice:%s", tgtUser, tgtHost))
	} else {
		a = append(a, fmt.Sprintf("ondevice:%s", tgtHost))
	}
	a = append(a, args...) // non-option ssh arguments (command to be run on the host)

	// syscall.Exec will replace this app with ssh (yes, replace it, not just launch)
	// therefore, unless there's an error, this is the last line of code to be executed
	err := syscall.Exec(sshPath, a, os.Environ())
	if err != nil {
		logg.Fatal("Failed to run ", sshPath, ": ", err)
	}

	// nothing here should ever be executed
	logg.Fatal("This should never happen")
	return -1
}

// sshParseArgs -- Takes `ondevice ssh` arguments and parses them (into flags/options and other arguments)
func sshParseArgs(flags map[byte]bool, args []string) (outArgs []string, outOpts []string) {
	for i := 0; i < len(args); i++ {
		arg := args[i]

		if len(arg) == 0 {
			logg.Fatal("Failed to parse SSH arguments: got an empty argument while looking for '[user@]host'")
		} else if arg == "-" {
			logg.Fatal("Stray '-' SSH argument")
		} else if arg[0] == '-' {
			hasValue, ok := flags[arg[1]]
			if !ok {
				logg.Fatal("Unsupported SSH argument: ", arg)
			}
			if hasValue && len(arg) == 2 {
				// the value's in the next argument, push both to outArgs
				outOpts = append(outOpts, arg, args[i+1])
				i++
			} else if hasValue && len(arg) > 2 {
				// the value's part of arg
				outOpts = append(outOpts, arg)
			} else if !hasValue && len(arg) > 2 { // && !hasValue
				logg.Fatal("Got value for flag that doesn't expect one: ", arg)
			} else if !hasValue && len(arg) == 2 {
				// a simple flag
				outOpts = append(outOpts, arg)
			} else {
				// yay to defensive programming
				logg.Fatal("this should never happen (fifth state of two binary values)")
			}
		} else {
			// first non-option argument -> we're done
			outArgs = args[i:]
			break
		}
	}

	return
}

func sshParseTarget(target string) (tgtHost string, tgtUser string) {
	parts := strings.SplitN(target, "@", 2)
	if len(parts) == 1 {
		tgtUser = ""
		tgtHost = parts[0]
	} else if len(parts) == 2 {
		tgtUser = parts[0]
		tgtHost = parts[1]
	} else {
		logg.Fatal("SplitN(..., 2) returned odd number of results: ", len(parts))
	}

	// always use qualified device names
	if !strings.Contains(tgtHost, ".") {
		user, _, err := config.GetClientAuth()
		if err == nil {
			tgtHost = fmt.Sprintf("%s.%s", user, tgtHost)
		}
	}

	return tgtHost, tgtUser
}

// sshParseFlags -- takes a getopt-style argument string and returns a map
// of flag characters and whether or not they expect an argument
func sshParseFlags(flags string) map[byte]bool {
	rc := map[byte]bool{}

	for i := 0; i < len(flags); i++ {
		flag := flags[i]
		hasValue := false
		if flags[i+1] == ':' {
			hasValue = true
			i++
		}
		rc[flag] = hasValue
	}
	return rc
}

// SSHCommand -- implements `ondevice ssh`
var SSHCommand = BaseCommand{
	Arguments: "[ssh-arguments...]",
	ShortHelp: "Connect to your devices using the ssh protocol",
	RunFn:     sshRun,
	LongHelp: `$ ondevice ssh [<user>@]<device> [ssh-arguments...]

Connect to your devices using the 'ssh' command.

This is a relatively thin wrapper around the 'ssh' command.
The main difference to invoking ssh directly is that instead of regular host names you'll have to specify an ondevice deviceId.
The connection is routed through the ondevice.io network.

ondevice ssh will try to parse ssh's arguments, the first non-argument has to be
the user@hostname combo.

See ssh's documentation for further details.

Examples:
- ondevice ssh device1
  simply connect to device1
- ondevice ssh user@device1
  open an SSH connection to device1, logging in as 'user'
- ondevice ssh device1 echo hello world
  run 'echo hello world' on device1
- ondevice ssh device1 -N -L 1234:localhost:80
  Tunnel the HTTP server on device1 to the local port 1234 without opening
  a shell
- ondevice ssh device1 -D 1080
  Starting a SOCKS5 proxy listening on port 1080. It'll redirect all traffic
  to the target host.
`,
}
<|endoftext|>


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

package main

import (
	"database/sql"
	"encoding/json"
	"flag"

	"github.com/BurntSushi/toml"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/tidb-tools/pkg/dbutil"
	router "github.com/pingcap/tidb-tools/pkg/table-router"
	"go.uber.org/zap"
)

const (
	percent0   = 0
	percent100 = 100
)

var sourceInstanceMap map[string]interface{} = make(map[string]interface{})

// DBConfig is the config of database, and keep the connection.
type DBConfig struct {
	dbutil.DBConfig

	InstanceID string `toml:"instance-id" json:"instance-id"`

	Conn *sql.DB
}

// Valid returns true if database's config is valide.
func (c *DBConfig) Valid() bool {
	if c.InstanceID == "" {
		log.Error("must specify source database's instance id")
		return false
	}
	sourceInstanceMap[c.InstanceID] = struct{}{}

	return true
}

// CheckTables saves the tables need to check.
type CheckTables struct {
	// schema name
	Schema string `toml:"schema" json:"schema"`

	// table list
	Tables []string `toml:"tables" json:"tables"`

	ExcludeTables []string `toml:"exclude-tables" json:"exclude-tables"`
}

// TableConfig is the config of table.
type TableConfig struct {
	// table's origin information
	TableInstance
	// columns be ignored, will not check this column's data
	IgnoreColumns []string `toml:"ignore-columns"`
	// field should be the primary key, unique key or field with index
	Fields string `toml:"index-fields"`
	// select range, for example: "age > 10 AND age < 20"
	Range string `toml:"range"`
	// set true if comparing sharding tables with target table, should have more than one source tables.
	IsSharding bool `toml:"is-sharding"`
	// saves the source tables's info.
	// may have more than one source for sharding tables.
	// or you want to compare table with different schema and table name.
	// SourceTables can be nil when source and target is one-to-one correspondence.
	SourceTables    []TableInstance `toml:"source-tables"`
	TargetTableInfo *model.TableInfo

	// collation config in mysql/tidb
	Collation string `toml:"collation"`
}

// Valid returns true if table's config is valide.
func (t *TableConfig) Valid() bool {
	if t.Schema == "" || t.Table == "" {
		log.Error("schema and table's name can't be empty")
		return false
	}

	if t.IsSharding {
		if len(t.SourceTables) <= 1 {
			log.Error("must have more than one source tables if comparing sharding tables")
			return false
		}

	} else {
		if len(t.SourceTables) > 1 {
			log.Error("have more than one source table in no sharding mode")
			return false
		}
	}

	for _, sourceTable := range t.SourceTables {
		if !sourceTable.Valid() {
			return false
		}
	}

	return true
}

// TableInstance saves the base information of table.
type TableInstance struct {
	// database's instance id
	InstanceID string `toml:"instance-id" json:"instance-id"`
	// schema name
	Schema string `toml:"schema"`
	// table name
	Table string `toml:"table"`
}

// Valid returns true if table instance's info is valide.
// should be executed after source database's check.
func (t *TableInstance) Valid() bool {
	if t.InstanceID == "" {
		log.Error("must specify the database's instance id for source table")
		return false
	}

	if _, ok := sourceInstanceMap[t.InstanceID]; !ok {
		log.Error("unknown database instance id", zap.String("instance id", t.InstanceID))
		return false
	}

	if t.Schema == "" || t.Table == "" {
		log.Error("schema and table's name can't be empty")
		return false
	}

	return true
}

// Config is the configuration.
type Config struct {
	*flag.FlagSet `json:"-"`

	// log level
	LogLevel string `toml:"log-level" json:"log-level"`

	// source database's config
	SourceDBCfg []DBConfig `toml:"source-db" json:"source-db"`

	// target database's config
	TargetDBCfg DBConfig `toml:"target-db" json:"target-db"`

	// for example, the whole data is [1...100]
	// we can split these data to [1...10], [11...20], ..., [91...100]
	// the [1...10] is a chunk, and it's chunk size is 10
	// size of the split chunk
	ChunkSize int `toml:"chunk-size" json:"chunk-size"`

	// sampling check percent, for example 10 means only check 10% data
	Sample int `toml:"sample-percent" json:"sample-percent"`

	// how many goroutines are created to check data
	CheckThreadCount int `toml:"check-thread-count" json:"check-thread-count"`

	// set false if want to comapre the data directly
	UseChecksum bool `toml:"use-checksum" json:"use-checksum"`

	// set true if just want compare data by checksum, will skip select data when checksum is not equal.
	OnlyUseChecksum bool `toml:"only-use-checksum" json:"only-use-checksum"`

	// the name of the file which saves sqls used to fix different data
	FixSQLFile string `toml:"fix-sql-file" json:"fix-sql-file"`

	// the tables to be checked
	Tables []*CheckTables `toml:"check-tables" json:"check-tables"`

	// TableRules defines table name and database name's conversion relationship between source database and target database
	TableRules []*router.TableRule `toml:"table-rules" json:"table-rules"`

	// the config of table
	TableCfgs []*TableConfig `toml:"table-config" json:"table-config"`

	// ignore check table's struct
	IgnoreStructCheck bool `toml:"ignore-struct-check" json:"ignore-struct-check"`

	// ignore check table's data
	IgnoreDataCheck bool `toml:"ignore-data-check" json:"ignore-data-check"`

	// set true will continue check from the latest checkpoint
	UseCheckpoint bool `toml:"use-checkpoint" json:"use-checkpoint"`

	// config file
	ConfigFile string

	// print version if set true
	PrintVersion bool
}

// NewConfig creates a new config.
func NewConfig() *Config {
	cfg := &Config{}
	cfg.FlagSet = flag.NewFlagSet("diff", flag.ContinueOnError)
	fs := cfg.FlagSet

	fs.StringVar(&cfg.ConfigFile, "config", "", "Config file")
	fs.StringVar(&cfg.LogLevel, "L", "info", "log level: debug, info, warn, error, fatal")
	fs.IntVar(&cfg.ChunkSize, "chunk-size", 1000, "diff check chunk size")
	fs.IntVar(&cfg.Sample, "sample", 100, "the percent of sampling check")
	fs.IntVar(&cfg.CheckThreadCount, "check-thread-count", 1, "how many goroutines are created to check data")
	fs.BoolVar(&cfg.UseChecksum, "use-checksum", true, "set false if want to comapre the data directly")
	fs.StringVar(&cfg.FixSQLFile, "fix-sql-file", "fix.sql", "the name of the file which saves sqls used to fix different data")
	fs.BoolVar(&cfg.PrintVersion, "V", false, "print version of sync_diff_inspector")
	fs.BoolVar(&cfg.IgnoreDataCheck, "ignore-data-check", false, "ignore check table's data")
	fs.BoolVar(&cfg.IgnoreStructCheck, "ignore-struct-check", false, "ignore check table's struct")
	fs.BoolVar(&cfg.UseCheckpoint, "use-checkpoint", true, "set true will continue check from the latest checkpoint")

	return cfg
}

// Parse parses flag definitions from the argument list.
func (c *Config) Parse(arguments []string) error {
	// Parse first to get config file.
	err := c.FlagSet.Parse(arguments)
	if err != nil {
		return errors.Trace(err)
	}

	// Load config file if specified.
	if c.ConfigFile != "" {
		err = c.configFromFile(c.ConfigFile)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// Parse again to replace with command line options.
	err = c.FlagSet.Parse(arguments)
	if err != nil {
		return errors.Trace(err)
	}

	if len(c.FlagSet.Args()) != 0 {
		return errors.Errorf("'%s' is an invalid flag", c.FlagSet.Arg(0))
	}

	return nil
}

func (c *Config) String() string {
	cfg, err := json.Marshal(c)
	if err != nil {
		return "<nil>"
	}
	return string(cfg)
}

// configFromFile loads config from file.
func (c *Config) configFromFile(path string) error {
	_, err := toml.DecodeFile(path, c)
	return errors.Trace(err)
}

func (c *Config) checkConfig() bool {
	if c.Sample > percent100 || c.Sample < percent0 {
		log.Error("sample must be greater than 0 and less than or equal to 100!")
		return false
	}

	if c.CheckThreadCount <= 0 {
		log.Error("check-thcount must greater than 0!")
		return false
	}

	if len(c.SourceDBCfg) == 0 {
		log.Error("must have at least one source database")
		return false
	}

	for i := range c.SourceDBCfg {
		if !c.SourceDBCfg[i].Valid() {
			return false
		}
	}

	if c.TargetDBCfg.InstanceID == "" {
		c.TargetDBCfg.InstanceID = "target"
	}
	if _, ok := sourceInstanceMap[c.TargetDBCfg.InstanceID]; ok {
		log.Error("target has same instance id in source", zap.String("instance id", c.TargetDBCfg.InstanceID))
		return false
	}

	if len(c.Tables) == 0 {
		log.Error("must specify check tables")
		return false
	}

	for _, tableCfg := range c.TableCfgs {
		if !tableCfg.Valid() {
			return false
		}
	}

	if c.OnlyUseChecksum {
		if !c.UseChecksum {
			log.Error("need set use-checksum = true")
			return false
		}
	} else {
		if len(c.FixSQLFile) == 0 {
			log.Warn("fix-sql-file is invalid, will use default value 'fix.sql'")
			c.FixSQLFile = "fix.sql"
		}
	}

	return true
}
<commit_msg>Tiny fix for the Snapshot using the DateTime string like "2020-01-01 20:20:20" (#373)<commit_after>// Copyright 2018 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"strconv"

	"github.com/BurntSushi/toml"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/tidb-tools/pkg/dbutil"
	router "github.com/pingcap/tidb-tools/pkg/table-router"
	"go.uber.org/zap"
)

const (
	percent0   = 0
	percent100 = 100
)

var sourceInstanceMap map[string]interface{} = make(map[string]interface{})

// DBConfig is the config of database, and keep the connection.
type DBConfig struct {
	dbutil.DBConfig

	InstanceID string `toml:"instance-id" json:"instance-id"`

	Conn *sql.DB
}

// Valid returns true if database's config is valide.
func (c *DBConfig) Valid() bool {
	if c.InstanceID == "" {
		log.Error("must specify source database's instance id")
		return false
	}
	sourceInstanceMap[c.InstanceID] = struct{}{}

	return true
}

// CheckTables saves the tables need to check.
type CheckTables struct {
	// schema name
	Schema string `toml:"schema" json:"schema"`

	// table list
	Tables []string `toml:"tables" json:"tables"`

	ExcludeTables []string `toml:"exclude-tables" json:"exclude-tables"`
}

// TableConfig is the config of table.
type TableConfig struct {
	// table's origin information
	TableInstance
	// columns be ignored, will not check this column's data
	IgnoreColumns []string `toml:"ignore-columns"`
	// field should be the primary key, unique key or field with index
	Fields string `toml:"index-fields"`
	// select range, for example: "age > 10 AND age < 20"
	Range string `toml:"range"`
	// set true if comparing sharding tables with target table, should have more than one source tables.
	IsSharding bool `toml:"is-sharding"`
	// saves the source tables's info.
	// may have more than one source for sharding tables.
	// or you want to compare table with different schema and table name.
	// SourceTables can be nil when source and target is one-to-one correspondence.
	SourceTables    []TableInstance `toml:"source-tables"`
	TargetTableInfo *model.TableInfo

	// collation config in mysql/tidb
	Collation string `toml:"collation"`
}

// Valid returns true if table's config is valide.
func (t *TableConfig) Valid() bool {
	if t.Schema == "" || t.Table == "" {
		log.Error("schema and table's name can't be empty")
		return false
	}

	if t.IsSharding {
		if len(t.SourceTables) <= 1 {
			log.Error("must have more than one source tables if comparing sharding tables")
			return false
		}

	} else {
		if len(t.SourceTables) > 1 {
			log.Error("have more than one source table in no sharding mode")
			return false
		}
	}

	for _, sourceTable := range t.SourceTables {
		if !sourceTable.Valid() {
			return false
		}
	}

	return true
}

// TableInstance saves the base information of table.
type TableInstance struct {
	// database's instance id
	InstanceID string `toml:"instance-id" json:"instance-id"`
	// schema name
	Schema string `toml:"schema"`
	// table name
	Table string `toml:"table"`
}

// Valid returns true if table instance's info is valide.
// should be executed after source database's check.
func (t *TableInstance) Valid() bool {
	if t.InstanceID == "" {
		log.Error("must specify the database's instance id for source table")
		return false
	}

	if _, ok := sourceInstanceMap[t.InstanceID]; !ok {
		log.Error("unknown database instance id", zap.String("instance id", t.InstanceID))
		return false
	}

	if t.Schema == "" || t.Table == "" {
		log.Error("schema and table's name can't be empty")
		return false
	}

	return true
}

// Config is the configuration.
type Config struct {
	*flag.FlagSet `json:"-"`

	// log level
	LogLevel string `toml:"log-level" json:"log-level"`

	// source database's config
	SourceDBCfg []DBConfig `toml:"source-db" json:"source-db"`

	// target database's config
	TargetDBCfg DBConfig `toml:"target-db" json:"target-db"`

	// for example, the whole data is [1...100]
	// we can split these data to [1...10], [11...20], ..., [91...100]
	// the [1...10] is a chunk, and it's chunk size is 10
	// size of the split chunk
	ChunkSize int `toml:"chunk-size" json:"chunk-size"`

	// sampling check percent, for example 10 means only check 10% data
	Sample int `toml:"sample-percent" json:"sample-percent"`

	// how many goroutines are created to check data
	CheckThreadCount int `toml:"check-thread-count" json:"check-thread-count"`

	// set false if want to comapre the data directly
	UseChecksum bool `toml:"use-checksum" json:"use-checksum"`

	// set true if just want compare data by checksum, will skip select data when checksum is not equal.
	OnlyUseChecksum bool `toml:"only-use-checksum" json:"only-use-checksum"`

	// the name of the file which saves sqls used to fix different data
	FixSQLFile string `toml:"fix-sql-file" json:"fix-sql-file"`

	// the tables to be checked
	Tables []*CheckTables `toml:"check-tables" json:"check-tables"`

	// TableRules defines table name and database name's conversion relationship between source database and target database
	TableRules []*router.TableRule `toml:"table-rules" json:"table-rules"`

	// the config of table
	TableCfgs []*TableConfig `toml:"table-config" json:"table-config"`

	// ignore check table's struct
	IgnoreStructCheck bool `toml:"ignore-struct-check" json:"ignore-struct-check"`

	// ignore check table's data
	IgnoreDataCheck bool `toml:"ignore-data-check" json:"ignore-data-check"`

	// set true will continue check from the latest checkpoint
	UseCheckpoint bool `toml:"use-checkpoint" json:"use-checkpoint"`

	// config file
	ConfigFile string

	// print version if set true
	PrintVersion bool
}

// NewConfig creates a new config.
func NewConfig() *Config {
	cfg := &Config{}
	cfg.FlagSet = flag.NewFlagSet("diff", flag.ContinueOnError)
	fs := cfg.FlagSet

	fs.StringVar(&cfg.ConfigFile, "config", "", "Config file")
	fs.StringVar(&cfg.LogLevel, "L", "info", "log level: debug, info, warn, error, fatal")
	fs.IntVar(&cfg.ChunkSize, "chunk-size", 1000, "diff check chunk size")
	fs.IntVar(&cfg.Sample, "sample", 100, "the percent of sampling check")
	fs.IntVar(&cfg.CheckThreadCount, "check-thread-count", 1, "how many goroutines are created to check data")
	fs.BoolVar(&cfg.UseChecksum, "use-checksum", true, "set false if want to comapre the data directly")
	fs.StringVar(&cfg.FixSQLFile, "fix-sql-file", "fix.sql", "the name of the file which saves sqls used to fix different data")
	fs.BoolVar(&cfg.PrintVersion, "V", false, "print version of sync_diff_inspector")
	fs.BoolVar(&cfg.IgnoreDataCheck, "ignore-data-check", false, "ignore check table's data")
	fs.BoolVar(&cfg.IgnoreStructCheck, "ignore-struct-check", false, "ignore check table's struct")
	fs.BoolVar(&cfg.UseCheckpoint, "use-checkpoint", true, "set true will continue check from the latest checkpoint")

	return cfg
}

// Parse parses flag definitions from the argument list.
func (c *Config) Parse(arguments []string) error {
	// Parse first to get config file.
	err := c.FlagSet.Parse(arguments)
	if err != nil {
		return errors.Trace(err)
	}

	// Load config file if specified.
	if c.ConfigFile != "" {
		err = c.configFromFile(c.ConfigFile)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// Parse again to replace with command line options.
	err = c.FlagSet.Parse(arguments)
	if err != nil {
		return errors.Trace(err)
	}

	if len(c.FlagSet.Args()) != 0 {
		return errors.Errorf("'%s' is an invalid flag", c.FlagSet.Arg(0))
	}

	return nil
}

func (c *Config) String() string {
	cfg, err := json.Marshal(c)
	if err != nil {
		return "<nil>"
	}
	return string(cfg)
}

// configFromFile loads config from file.
func (c *Config) configFromFile(path string) error {
	_, err := toml.DecodeFile(path, c)
	return errors.Trace(err)
}

func (c *Config) checkConfig() bool {
	if c.Sample > percent100 || c.Sample < percent0 {
		log.Error("sample must be greater than 0 and less than or equal to 100!")
		return false
	}

	if c.CheckThreadCount <= 0 {
		log.Error("check-thcount must greater than 0!")
		return false
	}

	if len(c.SourceDBCfg) == 0 {
		log.Error("must have at least one source database")
		return false
	}

	for i := range c.SourceDBCfg {
		if !c.SourceDBCfg[i].Valid() {
			return false
		}
		if c.SourceDBCfg[i].Snapshot != "" {
			c.SourceDBCfg[i].Snapshot = strconv.Quote(c.SourceDBCfg[i].Snapshot)
		}
	}

	if c.TargetDBCfg.InstanceID == "" {
		c.TargetDBCfg.InstanceID = "target"
	}
	if c.TargetDBCfg.Snapshot != "" {
		c.TargetDBCfg.Snapshot = strconv.Quote(c.TargetDBCfg.Snapshot)
	}
	if _, ok := sourceInstanceMap[c.TargetDBCfg.InstanceID]; ok {
		log.Error("target has same instance id in source", zap.String("instance id", c.TargetDBCfg.InstanceID))
		return false
	}

	if len(c.Tables) == 0 {
		log.Error("must specify check tables")
		return false
	}

	for _, tableCfg := range c.TableCfgs {
		if !tableCfg.Valid() {
			return false
		}
	}

	if c.OnlyUseChecksum {
		if !c.UseChecksum {
			log.Error("need set use-checksum = true")
			return false
		}
	} else {
		if len(c.FixSQLFile) == 0 {
			log.Warn("fix-sql-file is invalid, will use default value 'fix.sql'")
			c.FixSQLFile = "fix.sql"
		}
	}

	return true
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os/exec"
	"flag"
	"bytes"
	"encoding/json"
	"log"
	"io/ioutil"
	"strings"
	"strconv"
	"github.com/mattn/go-jsonpointer"
	"github.com/cmiceli/password-generator-go"
)

type awsPolicy struct {
	GroupName			string		`json:"GroupName"`
	PolicyName		string		`json:"PolicyName"`
	PolicyPath		string		`json:"PolicyPath"`
}

type awsParam struct {
	policy				awsPolicy
	userName			[]string
	awsAccount		string
}

func main() {
	var userNames string
	var param awsParam
	var policyJson string
	flag.StringVar(&userNames, "user-name", "", "user-name1,user-name2")
	flag.StringVar(&param.awsAccount, "account", "", ".aws/config file account")
	flag.StringVar(&policyJson	, "policy-json", "", "policy setting json file path")
	flag.Parse()

	//get json parameter
	param.policy = getPolicyJson(policyJson)
	param.userName = strings.Split(userNames, ",")

	awsAccountGenerate(param)
}

func awsAccountGenerate(param awsParam) {
	for _, v := range param.userName {
		cmdCreateGroupAndPolicy(param.policy.GroupName, param.policy.PolicyName, param.policy.PolicyPath, param.awsAccount)

		fmt.Println("URL: ")
		fmt.Println(cmdAccountAlias(param.awsAccount))

		fmt.Println("userName: ")
		fmt.Println(cmdCreateUser(param.awsAccount, v))
		cmdAddUserToGroup(param.policy.GroupName, v, param.awsAccount)

		fmt.Println("password: ")
		fmt.Println(cmdSetUserPassword(param.awsAccount, v))

		keyId, secretKey := cmdAccessKey(param.awsAccount, v)
		fmt.Println("accessKeyId: ")
		fmt.Println(keyId)
		fmt.Println("secretAccessKey: ")
		fmt.Println(secretKey)

		fmt.Println()
	}

}

func getPolicyJson(policyJsonPath string) (policy awsPolicy) {
	policyJson, err := ioutil.ReadFile(policyJsonPath)

	if err != nil {
		log.Fatalln("policy json read error:",err)
	}

	err = json.Unmarshal(policyJson, &policy)
	if err != nil {
		log.Fatalln("json unmarshal error:",err)
	}

	return
}

func cmdRun(cmd *exec.Cmd) (resultJson interface{}, err error) {
	var commandResponse bytes.Buffer
	cmd.Stdout = &commandResponse
	err = cmd.Run()
	if err != nil {
		fmt.Println(cmd.Args)
		log.Fatalln("cmd.Run:",err)
	}

	if commandResponse.Len() == 0 {
		return
	}

	err = json.Unmarshal([]byte(commandResponse.Bytes()), &resultJson)
	if err != nil {
		fmt.Println(cmd.Args)
		log.Fatalln("json.Unmarshal:",err)
	}

	return
}

func cmdCreateGroupAndPolicy(group string, policy string, path string,  account string) {
	execCommand := exec.Command("aws", "iam", "list-groups", "--profile", account)
	groupsJson, _ := cmdRun(execCommand)

	for i := 0; ; i++ {
		isValid := jsonpointer.Has(groupsJson, "/Groups/" + strconv.Itoa(i) + "/GroupName")
		if isValid == false {
			break
		}

		getGroup, _ := jsonpointer.Get(groupsJson, "/Groups/" + strconv.Itoa(i) + "/GroupName")
		if getGroup == group {
			//fmt.Println("group [" + group + "] is already.")
			return
		}
	}

	execCommand = exec.Command("aws", "iam", "create-group", "--group-name", group, "--profile", account)
	_, err := cmdRun(execCommand)
	if err != nil {
		log.Fatalln("create-group:",err)
	}

	execCommand = exec.Command("aws", "iam", "put-group-policy", "--group-name", group, "--policy-name", policy, "--policy-document", path, "--profile", account)
	_, err = cmdRun(execCommand)
	if err != nil {
		log.Fatalln("put-group-policy:",err)
	}
}

func cmdAddUserToGroup(group string, user string, account string) {
	execCommand := exec.Command("aws", "iam", "add-user-to-group", "--group-name", group, "--user-name", user, "--profile", account)
	_, err := cmdRun(execCommand)

	if err != nil {
		log.Fatalln("add-user-to-group:",err)
	}
}

func cmdAccountAlias(account string) string {
	execCommand :=	exec.Command("aws", "iam", "list-account-aliases", "--profile", account)
	result, _ := cmdRun(execCommand)
	accountAlias, _ := jsonpointer.Get(result, "/AccountAliases")

	if nil == accountAlias {
		execCommand :=	exec.Command("aws", "iam", "create-account-alias", "--account-alias", account, "--profile", account)
		result, _ = cmdRun(execCommand)
		accountAlias, _ = jsonpointer.Get(result, "/AccountAliases")
	} else {
		//fmt.Println("create account alias is already.")
	}

	a := fmt.Sprint(accountAlias)
	return "https://" + a[1:len(a) - 1] + ".signin.aws.amazon.com/console"
}

func cmdGroupPolicy(account string, user string, policy []string) {
}

func cmdCreateUser(account string, user string) interface{} {
	execCommand :=	exec.Command("aws", "iam", "create-user", "--user-name", user, "--profile", account)
	result, _ := cmdRun(execCommand)
	userName, _ := jsonpointer.Get(result, "/User/UserName")
	return userName
}

func cmdSetUserPassword(account string, user string) string {
	password := pwordgen.NewPassword(20)
	execCommand :=	exec.Command("aws", "iam", "create-login-profile", "--user-name", user, "--password", password, "--profile", account)
	_, err := cmdRun(execCommand)
	if err != nil {
		log.Fatalln("cmdSetUserPassword",err)
	}
	return password
}

func cmdAccessKey(account string, user string) (interface{}, interface{}) {
	execCommand :=	exec.Command("aws", "iam", "create-access-key", "--user-name", user, "--profile", account)
	result, _ := cmdRun(execCommand)
	keyId, _ := jsonpointer.Get(result, "/AccessKey/AccessKeyId")
	secretKey, _ := jsonpointer.Get(result, "/AccessKey/SecretAccessKey")
	return keyId, secretKey
}
<commit_msg>delete a func cmdGroupPolicy()<commit_after>package main

import (
	"fmt"
	"os/exec"
	"flag"
	"bytes"
	"encoding/json"
	"log"
	"io/ioutil"
	"strings"
	"strconv"
	"github.com/mattn/go-jsonpointer"
	"github.com/cmiceli/password-generator-go"
)

type awsPolicy struct {
	GroupName			string		`json:"GroupName"`
	PolicyName		string		`json:"PolicyName"`
	PolicyPath		string		`json:"PolicyPath"`
}

type awsParam struct {
	policy				awsPolicy
	userName			[]string
	awsAccount		string
}

func main() {
	var userNames string
	var param awsParam
	var policyJson string
	flag.StringVar(&userNames, "user-name", "", "user-name1,user-name2")
	flag.StringVar(&param.awsAccount, "account", "", ".aws/config file account")
	flag.StringVar(&policyJson	, "policy-json", "", "policy setting json file path")
	flag.Parse()

	//get json parameter
	param.policy = getPolicyJson(policyJson)
	param.userName = strings.Split(userNames, ",")

	awsAccountGenerate(param)
}

func awsAccountGenerate(param awsParam) {
	for _, v := range param.userName {
		cmdCreateGroupAndPolicy(param.policy.GroupName, param.policy.PolicyName, param.policy.PolicyPath, param.awsAccount)

		fmt.Println("URL: ")
		fmt.Println(cmdAccountAlias(param.awsAccount))

		fmt.Println("userName: ")
		fmt.Println(cmdCreateUser(param.awsAccount, v))
		cmdAddUserToGroup(param.policy.GroupName, v, param.awsAccount)

		fmt.Println("password: ")
		fmt.Println(cmdSetUserPassword(param.awsAccount, v))

		keyId, secretKey := cmdAccessKey(param.awsAccount, v)
		fmt.Println("accessKeyId: ")
		fmt.Println(keyId)
		fmt.Println("secretAccessKey: ")
		fmt.Println(secretKey)

		fmt.Println()
	}

}

func getPolicyJson(policyJsonPath string) (policy awsPolicy) {
	policyJson, err := ioutil.ReadFile(policyJsonPath)

	if err != nil {
		log.Fatalln("policy json read error:",err)
	}

	err = json.Unmarshal(policyJson, &policy)
	if err != nil {
		log.Fatalln("json unmarshal error:",err)
	}

	return
}

func cmdRun(cmd *exec.Cmd) (resultJson interface{}, err error) {
	var commandResponse bytes.Buffer
	cmd.Stdout = &commandResponse
	err = cmd.Run()
	if err != nil {
		fmt.Println(cmd.Args)
		log.Fatalln("cmd.Run:",err)
	}

	if commandResponse.Len() == 0 {
		return
	}

	err = json.Unmarshal([]byte(commandResponse.Bytes()), &resultJson)
	if err != nil {
		fmt.Println(cmd.Args)
		log.Fatalln("json.Unmarshal:",err)
	}

	return
}

func cmdCreateGroupAndPolicy(group string, policy string, path string,  account string) {
	execCommand := exec.Command("aws", "iam", "list-groups", "--profile", account)
	groupsJson, _ := cmdRun(execCommand)

	for i := 0; ; i++ {
		isValid := jsonpointer.Has(groupsJson, "/Groups/" + strconv.Itoa(i) + "/GroupName")
		if isValid == false {
			break
		}

		getGroup, _ := jsonpointer.Get(groupsJson, "/Groups/" + strconv.Itoa(i) + "/GroupName")
		if getGroup == group {
			//fmt.Println("group [" + group + "] is already.")
			return
		}
	}

	execCommand = exec.Command("aws", "iam", "create-group", "--group-name", group, "--profile", account)
	_, err := cmdRun(execCommand)
	if err != nil {
		log.Fatalln("create-group:",err)
	}

	execCommand = exec.Command("aws", "iam", "put-group-policy", "--group-name", group, "--policy-name", policy, "--policy-document", path, "--profile", account)
	_, err = cmdRun(execCommand)
	if err != nil {
		log.Fatalln("put-group-policy:",err)
	}
}

func cmdAddUserToGroup(group string, user string, account string) {
	execCommand := exec.Command("aws", "iam", "add-user-to-group", "--group-name", group, "--user-name", user, "--profile", account)
	_, err := cmdRun(execCommand)

	if err != nil {
		log.Fatalln("add-user-to-group:",err)
	}
}

func cmdAccountAlias(account string) string {
	execCommand :=	exec.Command("aws", "iam", "list-account-aliases", "--profile", account)
	result, _ := cmdRun(execCommand)
	accountAlias, _ := jsonpointer.Get(result, "/AccountAliases")

	if nil == accountAlias {
		execCommand :=	exec.Command("aws", "iam", "create-account-alias", "--account-alias", account, "--profile", account)
		result, _ = cmdRun(execCommand)
		accountAlias, _ = jsonpointer.Get(result, "/AccountAliases")
	} else {
		//fmt.Println("create account alias is already.")
	}

	a := fmt.Sprint(accountAlias)
	return "https://" + a[1:len(a) - 1] + ".signin.aws.amazon.com/console"
}

func cmdCreateUser(account string, user string) interface{} {
	execCommand :=	exec.Command("aws", "iam", "create-user", "--user-name", user, "--profile", account)
	result, _ := cmdRun(execCommand)
	userName, _ := jsonpointer.Get(result, "/User/UserName")
	return userName
}

func cmdSetUserPassword(account string, user string) string {
	password := pwordgen.NewPassword(20)
	execCommand :=	exec.Command("aws", "iam", "create-login-profile", "--user-name", user, "--password", password, "--profile", account)
	_, err := cmdRun(execCommand)
	if err != nil {
		log.Fatalln("cmdSetUserPassword",err)
	}
	return password
}

func cmdAccessKey(account string, user string) (interface{}, interface{}) {
	execCommand :=	exec.Command("aws", "iam", "create-access-key", "--user-name", user, "--profile", account)
	result, _ := cmdRun(execCommand)
	keyId, _ := jsonpointer.Get(result, "/AccessKey/AccessKeyId")
	secretKey, _ := jsonpointer.Get(result, "/AccessKey/SecretAccessKey")
	return keyId, secretKey
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cipher_test

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"os"
)

func ExampleNewCBCDecrypter() {
	key := []byte("example key 1234")
	ciphertext, _ := hex.DecodeString("f363f3ccdcb12bb883abf484ba77d9cd7d32b5baecb3d4b1b3e0e4beffdb3ded")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	if len(ciphertext) < aes.BlockSize {
		panic("ciphertext too short")
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	// CBC mode always works in whole blocks.
	if len(ciphertext)%aes.BlockSize != 0 {
		panic("ciphertext is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)

	// CryptBlocks can work in-place if the two arguments are the same.
	mode.CryptBlocks(ciphertext, ciphertext)

	// If the original plaintext lengths are not a multiple of the block
	// size, padding would have to be added when encrypting, which would be
	// removed at this point. For an example, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. However, it's
	// critical to note that ciphertexts must be authenticated (i.e. by
	// using crypto/hmac) before being decrypted in order to avoid creating
	// a padding oracle.

	fmt.Printf("%s\n", ciphertext)
	// Output: exampleplaintext
}

func ExampleNewCBCEncrypter() {
	key := []byte("example key 1234")
	plaintext := []byte("exampleplaintext")

	// CBC mode works on blocks so plaintexts may need to be padded to the
	// next whole block. For an example of such padding, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. Here we'll
	// assume that the plaintext is already of the correct length.
	if len(plaintext)%aes.BlockSize != 0 {
		panic("plaintext is not a multiple of the block size")
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	fmt.Printf("%x\n", ciphertext)
}

func ExampleNewCFBDecrypter() {
	key := []byte("example key 1234")
	ciphertext, _ := hex.DecodeString("22277966616d9bc47177bd02603d08c9a67d5380d0fe8cf3b44438dff7b9")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	if len(ciphertext) < aes.BlockSize {
		panic("ciphertext too short")
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)

	// XORKeyStream can work in-place if the two arguments are the same.
	stream.XORKeyStream(ciphertext, ciphertext)
	fmt.Printf("%s", ciphertext)
	// Output: some plaintext
}

func ExampleNewCFBEncrypter() {
	key := []byte("example key 1234")
	plaintext := []byte("some plaintext")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.
}

func ExampleNewCTR() {
	key := []byte("example key 1234")
	plaintext := []byte("some plaintext")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	// CTR mode is the same for both encryption and decryption, so we can
	// also decrypt that ciphertext with NewCTR.

	plaintext2 := make([]byte, len(plaintext))
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(plaintext2, ciphertext[aes.BlockSize:])

	fmt.Printf("%s\n", plaintext2)
	// Output: some plaintext
}

func ExampleNewOFB() {
	key := []byte("example key 1234")
	plaintext := []byte("some plaintext")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewOFB(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	// OFB mode is the same for both encryption and decryption, so we can
	// also decrypt that ciphertext with NewOFB.

	plaintext2 := make([]byte, len(plaintext))
	stream = cipher.NewOFB(block, iv)
	stream.XORKeyStream(plaintext2, ciphertext[aes.BlockSize:])

	fmt.Printf("%s\n", plaintext2)
	// Output: some plaintext
}

func ExampleStreamReader() {
	key := []byte("example key 1234")

	inFile, err := os.Open("encrypted-file")
	if err != nil {
		panic(err)
	}
	defer inFile.Close()

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// If the key is unique for each ciphertext, then it's ok to use a zero
	// IV.
	var iv [aes.BlockSize]byte
	stream := cipher.NewOFB(block, iv[:])

	outFile, err := os.OpenFile("decrypted-file", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		panic(err)
	}
	defer outFile.Close()

	reader := &cipher.StreamReader{S: stream, R: inFile}
	// Copy the input file to the output file, decrypting as we go.
	if _, err := io.Copy(outFile, reader); err != nil {
		panic(err)
	}

	// Note that this example is simplistic in that it omits any
	// authentication of the encrypted data. If you were actually to use
	// StreamReader in this manner, an attacker could flip arbitrary bits in
	// the output.
}

func ExampleStreamWriter() {
	key := []byte("example key 1234")

	inFile, err := os.Open("plaintext-file")
	if err != nil {
		panic(err)
	}
	defer inFile.Close()

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// If the key is unique for each ciphertext, then it's ok to use a zero
	// IV.
	var iv [aes.BlockSize]byte
	stream := cipher.NewOFB(block, iv[:])

	outFile, err := os.OpenFile("encrypted-file", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		panic(err)
	}
	defer outFile.Close()

	writer := &cipher.StreamWriter{S: stream, W: outFile}
	// Copy the input file to the output file, encrypting as we go.
	if _, err := io.Copy(writer, inFile); err != nil {
		panic(err)
	}

	// Note that this example is simplistic in that it omits any
	// authentication of the encrypted data. It you were actually to use
	// StreamReader in this manner, an attacker could flip arbitrary bits in
	// the decrypted result.
}
<commit_msg>crypto/cipher: fix typo in example comment<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cipher_test

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"os"
)

func ExampleNewCBCDecrypter() {
	key := []byte("example key 1234")
	ciphertext, _ := hex.DecodeString("f363f3ccdcb12bb883abf484ba77d9cd7d32b5baecb3d4b1b3e0e4beffdb3ded")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	if len(ciphertext) < aes.BlockSize {
		panic("ciphertext too short")
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	// CBC mode always works in whole blocks.
	if len(ciphertext)%aes.BlockSize != 0 {
		panic("ciphertext is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)

	// CryptBlocks can work in-place if the two arguments are the same.
	mode.CryptBlocks(ciphertext, ciphertext)

	// If the original plaintext lengths are not a multiple of the block
	// size, padding would have to be added when encrypting, which would be
	// removed at this point. For an example, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. However, it's
	// critical to note that ciphertexts must be authenticated (i.e. by
	// using crypto/hmac) before being decrypted in order to avoid creating
	// a padding oracle.

	fmt.Printf("%s\n", ciphertext)
	// Output: exampleplaintext
}

func ExampleNewCBCEncrypter() {
	key := []byte("example key 1234")
	plaintext := []byte("exampleplaintext")

	// CBC mode works on blocks so plaintexts may need to be padded to the
	// next whole block. For an example of such padding, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. Here we'll
	// assume that the plaintext is already of the correct length.
	if len(plaintext)%aes.BlockSize != 0 {
		panic("plaintext is not a multiple of the block size")
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	fmt.Printf("%x\n", ciphertext)
}

func ExampleNewCFBDecrypter() {
	key := []byte("example key 1234")
	ciphertext, _ := hex.DecodeString("22277966616d9bc47177bd02603d08c9a67d5380d0fe8cf3b44438dff7b9")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	if len(ciphertext) < aes.BlockSize {
		panic("ciphertext too short")
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)

	// XORKeyStream can work in-place if the two arguments are the same.
	stream.XORKeyStream(ciphertext, ciphertext)
	fmt.Printf("%s", ciphertext)
	// Output: some plaintext
}

func ExampleNewCFBEncrypter() {
	key := []byte("example key 1234")
	plaintext := []byte("some plaintext")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.
}

func ExampleNewCTR() {
	key := []byte("example key 1234")
	plaintext := []byte("some plaintext")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	// CTR mode is the same for both encryption and decryption, so we can
	// also decrypt that ciphertext with NewCTR.

	plaintext2 := make([]byte, len(plaintext))
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(plaintext2, ciphertext[aes.BlockSize:])

	fmt.Printf("%s\n", plaintext2)
	// Output: some plaintext
}

func ExampleNewOFB() {
	key := []byte("example key 1234")
	plaintext := []byte("some plaintext")

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewOFB(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	// OFB mode is the same for both encryption and decryption, so we can
	// also decrypt that ciphertext with NewOFB.

	plaintext2 := make([]byte, len(plaintext))
	stream = cipher.NewOFB(block, iv)
	stream.XORKeyStream(plaintext2, ciphertext[aes.BlockSize:])

	fmt.Printf("%s\n", plaintext2)
	// Output: some plaintext
}

func ExampleStreamReader() {
	key := []byte("example key 1234")

	inFile, err := os.Open("encrypted-file")
	if err != nil {
		panic(err)
	}
	defer inFile.Close()

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// If the key is unique for each ciphertext, then it's ok to use a zero
	// IV.
	var iv [aes.BlockSize]byte
	stream := cipher.NewOFB(block, iv[:])

	outFile, err := os.OpenFile("decrypted-file", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		panic(err)
	}
	defer outFile.Close()

	reader := &cipher.StreamReader{S: stream, R: inFile}
	// Copy the input file to the output file, decrypting as we go.
	if _, err := io.Copy(outFile, reader); err != nil {
		panic(err)
	}

	// Note that this example is simplistic in that it omits any
	// authentication of the encrypted data. If you were actually to use
	// StreamReader in this manner, an attacker could flip arbitrary bits in
	// the output.
}

func ExampleStreamWriter() {
	key := []byte("example key 1234")

	inFile, err := os.Open("plaintext-file")
	if err != nil {
		panic(err)
	}
	defer inFile.Close()

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// If the key is unique for each ciphertext, then it's ok to use a zero
	// IV.
	var iv [aes.BlockSize]byte
	stream := cipher.NewOFB(block, iv[:])

	outFile, err := os.OpenFile("encrypted-file", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		panic(err)
	}
	defer outFile.Close()

	writer := &cipher.StreamWriter{S: stream, W: outFile}
	// Copy the input file to the output file, encrypting as we go.
	if _, err := io.Copy(writer, inFile); err != nil {
		panic(err)
	}

	// Note that this example is simplistic in that it omits any
	// authentication of the encrypted data. If you were actually to use
	// StreamReader in this manner, an attacker could flip arbitrary bits in
	// the decrypted result.
}
<|endoftext|>