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

package main

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"strings"

	"github.com/adamdecaf/cert-manage/pkg/cmd"
	"github.com/adamdecaf/cert-manage/pkg/store"
	"github.com/adamdecaf/cert-manage/pkg/ui"
)

const Version = "0.1.1-dev"

var (
	fs = flag.NewFlagSet("flags", flag.ExitOnError)

	// incantations of "--help"
	flagHelp1 = fs.Bool("h", false, "")
	flagHelp2 = fs.Bool("help", false, "")
	flagHelp3 = fs.Bool("-help", false, "") // --help

	// -file is used to specify an input file path
	flagFile = fs.String("file", "", "")

	// -url is used to specify an input URL
	flagURL = fs.String("url", "", "")

	// -app is used for operating on an installed application
	flagApp = fs.String("app", "", "")

	// -ui is used for choosing a different ui
	flagUI = fs.String("ui", ui.DefaultUI(), "")

	// -from is used by 'gen-whitelist' to specify url sources
	flagFrom = fs.String("from", "", "")

	// -out is used by 'gen-whitelist' to specify output file location
	flagOutFile = fs.String("out", "", "")

	// Output
	flagCount  = fs.Bool("count", false, "")
	flagFormat = fs.String("format", ui.DefaultFormat(), "")

	// internal override to show help text
	callForHelp = false
)

func init() {
	fs.Usage = func() {
		fmt.Printf(`Usage of cert-manage version %s
SUB-COMMANDS
  add           Add certificate(s) to a store
                Accepts: -app, -file

  backup        Take a backup of the specified certificate store

  gen-whitelist Create a whitelist from various sources

  list          List the currently installed and trusted certificates

  restore       Revert the certificate trust back to, optionally takes -file <path>

  version       Show the version of cert-manage

  whitelist     Remove trust from certificates which do not match the whitelist in <path>

APPS
  Supported apps: %s

FLAGS
  -app <name>      The name of an application which to perform the given command on. (Examples: chrome, java)
  -file <path>     Local file path
  -from <type(s)>  Which sources to capture urls from. Comma separated list. (Options: browser, chrome, firefox, file)
  -help            Show this help dialog
  -ui <type>       Method of adjusting certificates to be removed/untrusted. (default: %s, options: %s)
  -url <where>     Remote URL to download and use in a command

OUTPUT
  -count  Output the count of certificates instead of each certificate
  -format <format> Change the output format for a given command (default: %s, options: %s)

DEBUGGING
  Alongside command line flags are two environmental varialbes read by cert-manage:
  - DEBUG=1        Enabled debug logging, GODEBUG=x509roots=1 also works and enabled Go's debugging
  - TRACE=<where>  Saves a binary trace file at <where> of the execution
`,
			getVersion(),
			strings.Join(store.GetApps(), ", "),
			ui.DefaultUI(),
			strings.Join(ui.GetUIs(), ", "),
			ui.DefaultFormat(),
			strings.Join(ui.GetFormats(), ", "),
		)
	}
}

func calledHelp() bool {
	return callForHelp || *flagHelp1 || *flagHelp2 || *flagHelp3
}

type command struct {
	fn    func() error
	appfn func(string) error

	help string
}

func trace() *cmd.Trace {
	trace, err := cmd.NewTrace(os.Getenv("TRACE"))
	if err != nil {
		panic(err)
	}
	err = trace.Start()
	if err != nil {
		panic(err)
	}
	return trace
}

func main() {
	t := trace()
	defer func() {
		err := t.Stop()
		if err != nil {
			panic(err)
		}
	}()

	// Just show help if there aren't enough arguments to do anything
	if len(os.Args) < 2 {
		fs.Usage()
		return
	}
	fs.Parse(os.Args[1:])
	if len(os.Args) == 2 && calledHelp() {
		fs.Usage()
		return
	}
	fs.Parse(os.Args[2:]) // reparse

	// Lift config options into a higher-level
	cfg := &ui.Config{
		Count:   *flagCount,
		Format:  *flagFormat,
		Outfile: *flagOutFile,
		UI:      *flagUI,
	}

	// Build up sub-commands
	commands := make(map[string]*command, 0)
	commands["add"] = &command{
		fn: func() error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.AddCertsFromFile(*flagFile)
		},
		appfn: func(a string) error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.AddCertsToAppFromFile(a, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage add -file <path> [-app <name>]

  Add a certificate to the platform store
    cert-manage add -file <path>

  Add a certificate to an application's store
    cert-manage add -file <path> -app <name>

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["backup"] = &command{
		fn: func() error {
			return cmd.BackupForPlatform()
		},
		appfn: func(a string) error {
			return cmd.BackupForApp(a)
		},
		help: fmt.Sprintf(`Usage: cert-manage backup [-app <name>]

  Backup a certificate store. This can be done for the platform or a given app.

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["gen-whitelist"] = &command{
		fn: func() error {
			if *flagOutFile == "" || (*flagFrom == "" && *flagFile == "") {
				callForHelp = true
				return nil
			}
			return cmd.GenerateWhitelist(*flagOutFile, *flagFrom, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage gen-whitelist -out <where> [-file <file>] [-from <type>]

  Generate a whitelist and write it to the filesystem. (At wherever -out points to.)

  Also, you can pass -file to read a newline delimited file of URL's.
    cert-manage gen-whitelist -file <path> -out whitelist.json

  Generate a whitelist from browser history
    cert-manage gen-whitelist -from firefox -out whitelist.json

  Generate a whitelist from all browsers on a computer
    cert-manage gen-whitelist -from browsers -out whitelist.json

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["list"] = &command{
		fn: func() error {
			if *flagFile != "" {
				return cmd.ListCertsFromFile(*flagFile, cfg)
			}
			if *flagURL != "" {
				return cmd.ListCertsFromURL(*flagURL, cfg)
			}
			return cmd.ListCertsForPlatform(cfg)
		},
		appfn: func(a string) error {
			return cmd.ListCertsForApp(a, cfg)
		},
		help: fmt.Sprintf(`Usage: cert-manage list [options]

  List certificates currently installed on the platform or application.

  List certificates from an application
    cert-mange list -app firefox

  List certificates from a file
    cert-mange list -file <path>

  List certificates from a URL
    cert-manage list -url <endpoint>

FORMATTING

  Change the output format (Default: %s, Options: %s)
    cert-manage list -format openssl

  Only show the count of certificates found
    cert-manage list -count
    cert-manage list -app java -count
    cert-manage list -file <path> -count

  Show the certificates on a local webpage (Default: %s, Options: %s)
    cert-manage list -ui web

APPS
  Supported apps: %s`,
			ui.DefaultFormat(),
			strings.Join(ui.GetFormats(), ", "),
			ui.DefaultUI(),
			strings.Join(ui.GetUIs(), ", "),
			strings.Join(store.GetApps(), ", ")),
	}
	commands["restore"] = &command{
		fn: func() error {
			return cmd.RestoreForPlatform(*flagFile)
		},
		appfn: func(a string) error {
			return cmd.RestoreForApp(a, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage restore [-app <name>] [-file <path>]

  Restore certificates from the latest backup
    cert-manage restore

  Restore certificates for the platform from a file
    cert-manage restore -file <path>

  Restore certificates for an application from the latest backup
    cert-manage restore -app java

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["whitelist"] = &command{
		fn: func() error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.WhitelistForPlatform(*flagFile)
		},
		appfn: func(a string) error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.WhitelistForApp(a, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage whitelist [-app <name>] -file <path>

  Remove untrusted certificates from a store for the platform
    cert-manage whitelist -file whitelist.json

  Remove untrusted certificates in an app
    cert-manage whitelist -file whitelist.json -app java

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["version"] = &command{
		fn: func() error {
			fmt.Printf("%s\n", getVersion())
			return nil
		},
		appfn: func(_ string) error {
			return nil
		},
		help: getVersion(),
	}

	// Run whatever function we've got here..
	c, ok := commands[strings.ToLower(os.Args[1])]
	if !ok && calledHelp() { // sub-command wasn't found
		fs.Usage()
		os.Exit(1)
	}
	if ok && calledHelp() {
		fmt.Println(c.help)
		os.Exit(1)
	}

	// sub-command found, try and exec something off it
	if flagApp != nil && *flagApp != "" {
		err := c.appfn(*flagApp)
		if err != nil {
			fmt.Printf("ERROR: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}
	err := c.fn()
	if err != nil {
		fmt.Printf("ERROR: %v\n", err)
		os.Exit(1)
	}

	// some flags can set callForHelp, so let's check again
	if calledHelp() {
		fmt.Println(c.help)
		os.Exit(1)
	}
}

func getVersion() string {
	return fmt.Sprintf("%s (Go: %s)", Version, runtime.Version())
}
<commit_msg>docs: drop app exanples from help text<commit_after>// Copyright 2018 Adam Shannon
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"strings"

	"github.com/adamdecaf/cert-manage/pkg/cmd"
	"github.com/adamdecaf/cert-manage/pkg/store"
	"github.com/adamdecaf/cert-manage/pkg/ui"
)

const Version = "0.1.1-dev"

var (
	fs = flag.NewFlagSet("flags", flag.ExitOnError)

	// incantations of "--help"
	flagHelp1 = fs.Bool("h", false, "")
	flagHelp2 = fs.Bool("help", false, "")
	flagHelp3 = fs.Bool("-help", false, "") // --help

	// -file is used to specify an input file path
	flagFile = fs.String("file", "", "")

	// -url is used to specify an input URL
	flagURL = fs.String("url", "", "")

	// -app is used for operating on an installed application
	flagApp = fs.String("app", "", "")

	// -ui is used for choosing a different ui
	flagUI = fs.String("ui", ui.DefaultUI(), "")

	// -from is used by 'gen-whitelist' to specify url sources
	flagFrom = fs.String("from", "", "")

	// -out is used by 'gen-whitelist' to specify output file location
	flagOutFile = fs.String("out", "", "")

	// Output
	flagCount  = fs.Bool("count", false, "")
	flagFormat = fs.String("format", ui.DefaultFormat(), "")

	// internal override to show help text
	callForHelp = false
)

func init() {
	fs.Usage = func() {
		fmt.Printf(`Usage of cert-manage version %s
SUB-COMMANDS
  add           Add certificate(s) to a store
                Accepts: -app, -file

  backup        Take a backup of the specified certificate store

  gen-whitelist Create a whitelist from various sources

  list          List the currently installed and trusted certificates

  restore       Revert the certificate trust back to, optionally takes -file <path>

  version       Show the version of cert-manage

  whitelist     Remove trust from certificates which do not match the whitelist in <path>

APPS
  Supported apps: %s

FLAGS
  -app <name>      The name of an application which to perform the given command on.
  -file <path>     Local file path
  -from <type(s)>  Which sources to capture urls from. Comma separated list. (Options: browser, chrome, firefox, file)
  -help            Show this help dialog
  -ui <type>       Method of adjusting certificates to be removed/untrusted. (default: %s, options: %s)
  -url <where>     Remote URL to download and use in a command

OUTPUT
  -count  Output the count of certificates instead of each certificate
  -format <format> Change the output format for a given command (default: %s, options: %s)

DEBUGGING
  Alongside command line flags are two environmental varialbes read by cert-manage:
  - DEBUG=1        Enabled debug logging, GODEBUG=x509roots=1 also works and enabled Go's debugging
  - TRACE=<where>  Saves a binary trace file at <where> of the execution
`,
			getVersion(),
			strings.Join(store.GetApps(), ", "),
			ui.DefaultUI(),
			strings.Join(ui.GetUIs(), ", "),
			ui.DefaultFormat(),
			strings.Join(ui.GetFormats(), ", "),
		)
	}
}

func calledHelp() bool {
	return callForHelp || *flagHelp1 || *flagHelp2 || *flagHelp3
}

type command struct {
	fn    func() error
	appfn func(string) error

	help string
}

func trace() *cmd.Trace {
	trace, err := cmd.NewTrace(os.Getenv("TRACE"))
	if err != nil {
		panic(err)
	}
	err = trace.Start()
	if err != nil {
		panic(err)
	}
	return trace
}

func main() {
	t := trace()
	defer func() {
		err := t.Stop()
		if err != nil {
			panic(err)
		}
	}()

	// Just show help if there aren't enough arguments to do anything
	if len(os.Args) < 2 {
		fs.Usage()
		return
	}
	fs.Parse(os.Args[1:])
	if len(os.Args) == 2 && calledHelp() {
		fs.Usage()
		return
	}
	fs.Parse(os.Args[2:]) // reparse

	// Lift config options into a higher-level
	cfg := &ui.Config{
		Count:   *flagCount,
		Format:  *flagFormat,
		Outfile: *flagOutFile,
		UI:      *flagUI,
	}

	// Build up sub-commands
	commands := make(map[string]*command, 0)
	commands["add"] = &command{
		fn: func() error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.AddCertsFromFile(*flagFile)
		},
		appfn: func(a string) error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.AddCertsToAppFromFile(a, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage add -file <path> [-app <name>]

  Add a certificate to the platform store
    cert-manage add -file <path>

  Add a certificate to an application's store
    cert-manage add -file <path> -app <name>

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["backup"] = &command{
		fn: func() error {
			return cmd.BackupForPlatform()
		},
		appfn: func(a string) error {
			return cmd.BackupForApp(a)
		},
		help: fmt.Sprintf(`Usage: cert-manage backup [-app <name>]

  Backup a certificate store. This can be done for the platform or a given app.

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["gen-whitelist"] = &command{
		fn: func() error {
			if *flagOutFile == "" || (*flagFrom == "" && *flagFile == "") {
				callForHelp = true
				return nil
			}
			return cmd.GenerateWhitelist(*flagOutFile, *flagFrom, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage gen-whitelist -out <where> [-file <file>] [-from <type>]

  Generate a whitelist and write it to the filesystem. (At wherever -out points to.)

  Also, you can pass -file to read a newline delimited file of URL's.
    cert-manage gen-whitelist -file <path> -out whitelist.json

  Generate a whitelist from browser history
    cert-manage gen-whitelist -from firefox -out whitelist.json

  Generate a whitelist from all browsers on a computer
    cert-manage gen-whitelist -from browsers -out whitelist.json

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["list"] = &command{
		fn: func() error {
			if *flagFile != "" {
				return cmd.ListCertsFromFile(*flagFile, cfg)
			}
			if *flagURL != "" {
				return cmd.ListCertsFromURL(*flagURL, cfg)
			}
			return cmd.ListCertsForPlatform(cfg)
		},
		appfn: func(a string) error {
			return cmd.ListCertsForApp(a, cfg)
		},
		help: fmt.Sprintf(`Usage: cert-manage list [options]

  List certificates currently installed on the platform or application.

  List certificates from an application
    cert-mange list -app firefox

  List certificates from a file
    cert-mange list -file <path>

  List certificates from a URL
    cert-manage list -url <endpoint>

FORMATTING

  Change the output format (Default: %s, Options: %s)
    cert-manage list -format openssl

  Only show the count of certificates found
    cert-manage list -count
    cert-manage list -app java -count
    cert-manage list -file <path> -count

  Show the certificates on a local webpage (Default: %s, Options: %s)
    cert-manage list -ui web

APPS
  Supported apps: %s`,
			ui.DefaultFormat(),
			strings.Join(ui.GetFormats(), ", "),
			ui.DefaultUI(),
			strings.Join(ui.GetUIs(), ", "),
			strings.Join(store.GetApps(), ", ")),
	}
	commands["restore"] = &command{
		fn: func() error {
			return cmd.RestoreForPlatform(*flagFile)
		},
		appfn: func(a string) error {
			return cmd.RestoreForApp(a, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage restore [-app <name>] [-file <path>]

  Restore certificates from the latest backup
    cert-manage restore

  Restore certificates for the platform from a file
    cert-manage restore -file <path>

  Restore certificates for an application from the latest backup
    cert-manage restore -app java

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["whitelist"] = &command{
		fn: func() error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.WhitelistForPlatform(*flagFile)
		},
		appfn: func(a string) error {
			if *flagFile == "" {
				callForHelp = true
				return nil
			}
			return cmd.WhitelistForApp(a, *flagFile)
		},
		help: fmt.Sprintf(`Usage: cert-manage whitelist [-app <name>] -file <path>

  Remove untrusted certificates from a store for the platform
    cert-manage whitelist -file whitelist.json

  Remove untrusted certificates in an app
    cert-manage whitelist -file whitelist.json -app java

APPS
  Supported apps: %s`,
			strings.Join(store.GetApps(), ", ")),
	}
	commands["version"] = &command{
		fn: func() error {
			fmt.Printf("%s\n", getVersion())
			return nil
		},
		appfn: func(_ string) error {
			return nil
		},
		help: getVersion(),
	}

	// Run whatever function we've got here..
	c, ok := commands[strings.ToLower(os.Args[1])]
	if !ok && calledHelp() { // sub-command wasn't found
		fs.Usage()
		os.Exit(1)
	}
	if ok && calledHelp() {
		fmt.Println(c.help)
		os.Exit(1)
	}

	// sub-command found, try and exec something off it
	if flagApp != nil && *flagApp != "" {
		err := c.appfn(*flagApp)
		if err != nil {
			fmt.Printf("ERROR: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}
	err := c.fn()
	if err != nil {
		fmt.Printf("ERROR: %v\n", err)
		os.Exit(1)
	}

	// some flags can set callForHelp, so let's check again
	if calledHelp() {
		fmt.Println(c.help)
		os.Exit(1)
	}
}

func getVersion() string {
	return fmt.Sprintf("%s (Go: %s)", Version, runtime.Version())
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "os"
  "os/exec"
  "syscall"
  "time"
  "bytes"
  "encoding/json"
  "net"
  "strings"
)

const (
  BIND_ADDR   = "0.0.0.0:20000"
  BUFFER_SIZE = 1024
)

type Command struct {
  Command    string         `json:"command"`
  ExitStatus int            `json:"exit_status"`
  Output     string         `json:"output"`
  TimeStart  time.Time      `json:"time_start"`
  TimeFinish time.Time      `json:"time_finish"`
  Duration   time.Duration  `json:"duration"`
}

func (cmd *Command) Run(command string) {
  var output bytes.Buffer
  shell := exec.Command("bash", "-c", command)

  shell.Stdout  = &output
  shell.Stderr  = &output
  cmd.Command   = command
  cmd.TimeStart = time.Now()

  err := shell.Start()
  if (err != nil) {
    fmt.Println("Execution failed:", err)
  }

  err = shell.Wait()

  cmd.TimeFinish = time.Now()
  cmd.Duration   = cmd.TimeFinish.Sub(cmd.TimeStart)

  if msg, ok := err.(*exec.ExitError); ok {
    cmd.ExitStatus = msg.Sys().(syscall.WaitStatus).ExitStatus()
  } else {
    cmd.ExitStatus = 0
  }

  cmd.Output = string(output.Bytes())
}

func (cmd *Command) ToJson() (s string) {
  buff, err := json.Marshal(cmd)

  if (err != nil) {
    s = ""
    return
  }

  return string(buff)
}

func (cmd *Command) Print() {
  fmt.Println("Command:   ", cmd.Command)
  fmt.Println("ExitStatus:", cmd.ExitStatus)
  fmt.Println("Duration:  ", cmd.Duration)
  fmt.Println("Output:    ", cmd.Output)
}

func (cmd *Command) Success() bool {
  return cmd.ExitStatus == 0
}

func Exec(str string) *Command {
  var command *Command
  command = new(Command)
  command.Run(str)

  return command
}

func HandleConnection(socket net.Conn) {
  buffer   := make([]byte, BUFFER_SIZE)
  num, err := socket.Read(buffer)

  if err != nil {
    fmt.Println("Read error:", err.Error())
    return
  }

  cmd    := strings.TrimSpace(string(buffer[0:num]))
  result := Exec(cmd)

  _, err = socket.Write([]byte(result.Output))
  socket.Close()
}

func main() {
  fmt.Printf("Starting server on %s\n", BIND_ADDR)

  server, err := net.Listen("tcp", BIND_ADDR)
  if err != nil {
    fmt.Println("Error:", err.Error())
    os.Exit(1)
  }

  for {
    socket, err := server.Accept()
    if err != nil {
      fmt.Println("Accept error:", err.Error())
      return
    }

    go HandleConnection(socket)
  }
}<commit_msg>Wait for commands<commit_after>package main

import (
  "fmt"
  "os"
  "os/exec"
  "syscall"
  "time"
  "bytes"
  "encoding/json"
  "net"
  "strings"
)

const (
  BIND_ADDR   = "0.0.0.0:20000"
  BUFFER_SIZE = 1024
)

type Command struct {
  Command    string         `json:"command"`
  ExitStatus int            `json:"exit_status"`
  Output     string         `json:"output"`
  TimeStart  time.Time      `json:"time_start"`
  TimeFinish time.Time      `json:"time_finish"`
  Duration   time.Duration  `json:"duration"`
}

func (cmd *Command) Run(command string) {
  var output bytes.Buffer
  shell := exec.Command("bash", "-c", command)

  shell.Stdout  = &output
  shell.Stderr  = &output
  cmd.Command   = command
  cmd.TimeStart = time.Now()

  err := shell.Start()
  if (err != nil) {
    fmt.Println("Execution failed:", err)
  }

  err = shell.Wait()

  cmd.TimeFinish = time.Now()
  cmd.Duration   = cmd.TimeFinish.Sub(cmd.TimeStart)

  if msg, ok := err.(*exec.ExitError); ok {
    cmd.ExitStatus = msg.Sys().(syscall.WaitStatus).ExitStatus()
  } else {
    cmd.ExitStatus = 0
  }

  cmd.Output = string(output.Bytes())
}

func (cmd *Command) ToJson() (s string) {
  buff, err := json.Marshal(cmd)

  if (err != nil) {
    s = ""
    return
  }

  return string(buff)
}

func (cmd *Command) Print() {
  fmt.Println("Command:   ", cmd.Command)
  fmt.Println("ExitStatus:", cmd.ExitStatus)
  fmt.Println("Duration:  ", cmd.Duration)
  fmt.Println("Output:    ", cmd.Output)
}

func (cmd *Command) Success() bool {
  return cmd.ExitStatus == 0
}

func Exec(str string) *Command {
  var command *Command
  command = new(Command)
  command.Run(str)

  return command
}

func HandleConnection(socket net.Conn) {
  buffer := make([]byte, BUFFER_SIZE)

  for {
    num, err := socket.Read(buffer)

    if err != nil {
      fmt.Println("Read error:", err.Error())
      break
    }

    cmd := strings.TrimSpace(string(buffer[0:num]))

    if (len(cmd) == 0) {
      continue
    }

    if (cmd == "!done") {
      break
    }

    fmt.Println("Executing:", cmd)
    result := Exec(cmd)
    _, err = socket.Write([]byte(result.ToJson()))
  }

  fmt.Println("Client connection closed")
  socket.Close()
}

func main() {
  fmt.Printf("Starting server on %s\n", BIND_ADDR)

  server, err := net.Listen("tcp", BIND_ADDR)
  if err != nil {
    fmt.Println("Error:", err.Error())
    os.Exit(1)
  }

  for {
    socket, err := server.Accept()
    if err != nil {
      fmt.Println("Accept error:", err.Error())
      return
    }

    go HandleConnection(socket)
  }
}<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"

	flag "github.com/ogier/pflag"
	"github.com/yarbelk/refasta/formats"
	"github.com/yarbelk/refasta/sequence"
)

type FakeReadCloser struct {
	io.Reader
}

func (f FakeReadCloser) Close() error {
	return nil
}

type FakeWriteCloser struct {
	io.Writer
}

type TNTContext struct {
	Title    string
	Outgroup string
}

func (f FakeWriteCloser) Close() error {
	return nil
}

func getOutputFilePointer(filename string) (io.WriteCloser, error) {
	if filename == "--" {
		return FakeWriteCloser{os.Stdout}, nil
	}
	return os.Create(filename)
}

func getInputFilePointer(filename string) (io.ReadCloser, error) {
	if filename == "--" {
		return FakeReadCloser{os.Stdin}, nil
	}
	return os.Open(filename)
}

func IsDirectory(path string) (bool, error) {
	fileInfo, err := os.Stat(path)
	return fileInfo.IsDir(), err
}

func isFormat(file, format string) bool {
	switch path.Ext(file) {
	case ".fas", ".fasta":
		return format == formats.FASTA_FORMAT
	default:
		return false
	}
}

func dirInput(dir, format string, recurse bool) ([]string, error) {
	absDir, err := filepath.Abs(dir)
	if err != nil {
		return nil, err
	}
	files, err := ioutil.ReadDir(absDir)
	if err != nil {
		return nil, err
	}
	filteredFiles := make([]string, 0, 10)
	for _, file := range files {
		if file.IsDir() && recurse {
			f, err := dirInput(filepath.Join(absDir, file.Name()), format, true)
			if err != nil {
				return nil, err
			}
			filteredFiles = append(filteredFiles, f...)
		} else if isFormat(file.Name(), format) {
			filteredFiles = append(filteredFiles, filepath.Join(absDir, file.Name()))
		}
	}
	return filteredFiles, nil
}

func isDirectory(path string) (bool, error) {
	fileInfo, err := os.Stat(path)
	return fileInfo.IsDir(), err
}

func handleFastaInput(input string) ([]sequence.Sequence, error) {
	var files []string
	var sequences []sequence.Sequence

	if isDir, err := isDirectory(input); isDir && err == nil {
		files, err = dirInput(input, formats.FASTA_FORMAT, true)
		if err != nil {
			// Some error in walking the directory tree
			return nil, err
		}
	} else {
		files = []string{input}
	}
	fmt.Println("files", files)

	for _, file := range files {
		ext := path.Ext(file)
		geneName := file[:len(file)-len(ext)]
		fasta := formats.Fasta{SpeciesFromID: true}
		fd, err := os.Open(file)
		if err != nil {
			// probably an Access Control issue, or race condition
			return nil, err
		}
		defer fd.Close()
		err = fasta.Parse(fd, geneName)
		if err != nil {
			// Some parsing error...
			return nil, err
		}
		sequences = append(sequences, fasta.Sequences...)
	}
	return sequences, nil
}

func handleFastaOutput(sequences []sequence.Sequence, output string) error {
	fasta := formats.Fasta{}
	fmt.Println(sequences)
	fasta.AddSequence(sequences...)
	fd, err := os.Create(output)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Issue opening output file;\n%s", err.Error())
	}
	defer fd.Close()
	return fasta.WriteSequences(fd)
}

func handleTNTOutput(context TNTContext, sequences []sequence.Sequence, output string) error {
	tnt := formats.TNT{Title: context.Title}
	fmt.Println(sequences)
	tnt.AddSequence(sequences...)
	tnt.SetOutgroup(context.Outgroup)
	fd, err := os.Create(output)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Issue opening output file;\n%s", err.Error())
	}
	defer fd.Close()
	return tnt.WriteSequences(fd)
}

func main() {
	inputFormat := flag.StringP("input-format", "f", formats.FASTA_FORMAT, "intput format, must be 'fasta'")
	outputFormat := flag.StringP("outpu-format", "F", formats.TNT_FORMAT, "output format, must be [fasta|tnt]")
	input := flag.StringP("input-file", "i", "--", "input file, it must be a valid input, or '--', or blank.  if blank. or '--', will read from stdin")
	output := flag.StringP("output-file", "o", "--", "output file, it must be a valid input, or '--', or blank.  if blank. or '--', will write to stdout")
	tntTitle := flag.StringP("tnt-title", "t", "", "title for TNT output")
	outgroup := flag.String("outgroup", "", "outgroup for TNT")
	flag.Parse()

	var sequences []sequence.Sequence
	var err error

	switch *inputFormat {
	case formats.FASTA_FORMAT:
		fmt.Fprintf(os.Stderr, "intput format is fasta; parsing\n")
		sequences, err = handleFastaInput(*input)
	default:
		fmt.Fprintf(os.Stderr, "Unknown intput format '%s'", inputFormat)
		os.Exit(1)
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err.Error())
		os.Exit(1)
	}

	switch *outputFormat {
	case formats.FASTA_FORMAT:
		fmt.Fprintf(os.Stderr, "Output format is fasta; serializing\n")
		if err := handleFastaOutput(sequences, *output); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err.Error())
			os.Exit(1)
		}
	case formats.TNT_FORMAT:
		fmt.Fprintf(os.Stderr, "Output format is TNT; serializing\n")
		context := TNTContext{
			Title:    *tntTitle,
			Outgroup: *outgroup,
		}
		if err := handleTNTOutput(context, sequences, *output); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err.Error())
			os.Exit(1)
		}
	default:
		fmt.Fprintf(os.Stderr, "Unknown output format '%s'", inputFormat)
		os.Exit(1)
	}

}
<commit_msg>close open files faster, remove print debugging in main<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"

	flag "github.com/ogier/pflag"
	"github.com/yarbelk/refasta/formats"
	"github.com/yarbelk/refasta/sequence"
)

type FakeReadCloser struct {
	io.Reader
}

func (f FakeReadCloser) Close() error {
	return nil
}

type FakeWriteCloser struct {
	io.Writer
}

type TNTContext struct {
	Title    string
	Outgroup string
}

func (f FakeWriteCloser) Close() error {
	return nil
}

func getOutputFilePointer(filename string) (io.WriteCloser, error) {
	if filename == "--" {
		return FakeWriteCloser{os.Stdout}, nil
	}
	return os.Create(filename)
}

func getInputFilePointer(filename string) (io.ReadCloser, error) {
	if filename == "--" {
		return FakeReadCloser{os.Stdin}, nil
	}
	return os.Open(filename)
}

func IsDirectory(path string) (bool, error) {
	fileInfo, err := os.Stat(path)
	return fileInfo.IsDir(), err
}

func isFormat(file, format string) bool {
	switch path.Ext(file) {
	case ".fas", ".fasta":
		return format == formats.FASTA_FORMAT
	default:
		return false
	}
}

func dirInput(dir, format string, recurse bool) ([]string, error) {
	absDir, err := filepath.Abs(dir)
	if err != nil {
		return nil, err
	}
	files, err := ioutil.ReadDir(absDir)
	if err != nil {
		return nil, err
	}
	filteredFiles := make([]string, 0, 10)
	for _, file := range files {
		if file.IsDir() && recurse {
			f, err := dirInput(filepath.Join(absDir, file.Name()), format, true)
			if err != nil {
				return nil, err
			}
			filteredFiles = append(filteredFiles, f...)
		} else if isFormat(file.Name(), format) {
			filteredFiles = append(filteredFiles, filepath.Join(absDir, file.Name()))
		}
	}
	return filteredFiles, nil
}

func isDirectory(path string) (bool, error) {
	fileInfo, err := os.Stat(path)
	return fileInfo.IsDir(), err
}

func handleFastaInput(input string) ([]sequence.Sequence, error) {
	var files []string
	var sequences []sequence.Sequence

	if isDir, err := isDirectory(input); isDir && err == nil {
		files, err = dirInput(input, formats.FASTA_FORMAT, true)
		if err != nil {
			// Some error in walking the directory tree
			return nil, err
		}
	} else {
		files = []string{input}
	}

	for _, file := range files {
		ext := path.Ext(file)
		geneName := file[:len(file)-len(ext)]
		err := func() error {
			fasta := formats.Fasta{SpeciesFromID: true}
			fd, err := os.Open(file)
			if err != nil {
				// probably an Access Control issue, or race condition
				return err
			}
			defer fd.Close()
			err = fasta.Parse(fd, geneName)
			if err != nil {
				// Some parsing error...
				return err
			}
			sequences = append(sequences, fasta.Sequences...)
			return nil
		}()
		if err != nil {
			return nil, err
		}
	}
	return sequences, nil
}

func handleFastaOutput(sequences []sequence.Sequence, output string) error {
	fasta := formats.Fasta{}
	fasta.AddSequence(sequences...)
	fd, err := os.Create(output)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Issue opening output file;\n%s", err.Error())
	}
	defer fd.Close()
	return fasta.WriteSequences(fd)
}

func handleTNTOutput(context TNTContext, sequences []sequence.Sequence, output string) error {
	tnt := formats.TNT{Title: context.Title}
	tnt.AddSequence(sequences...)
	tnt.SetOutgroup(context.Outgroup)
	fd, err := os.Create(output)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Issue opening output file;\n%s", err.Error())
	}
	defer fd.Close()
	return tnt.WriteSequences(fd)
}

func main() {
	inputFormat := flag.StringP("input-format", "f", formats.FASTA_FORMAT, "intput format, must be 'fasta'")
	outputFormat := flag.StringP("outpu-format", "F", formats.TNT_FORMAT, "output format, must be [fasta|tnt]")
	input := flag.StringP("input-file", "i", "--", "input file, it must be a valid input, or '--', or blank.  if blank. or '--', will read from stdin")
	output := flag.StringP("output-file", "o", "--", "output file, it must be a valid input, or '--', or blank.  if blank. or '--', will write to stdout")
	tntTitle := flag.StringP("tnt-title", "t", "", "title for TNT output")
	outgroup := flag.String("outgroup", "", "outgroup for TNT")
	flag.Parse()

	var sequences []sequence.Sequence
	var err error

	switch *inputFormat {
	case formats.FASTA_FORMAT:
		fmt.Fprintf(os.Stderr, "intput format is fasta; parsing\n")
		sequences, err = handleFastaInput(*input)
	default:
		fmt.Fprintf(os.Stderr, "Unknown intput format '%s'", inputFormat)
		os.Exit(1)
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err.Error())
		os.Exit(1)
	}

	switch *outputFormat {
	case formats.FASTA_FORMAT:
		fmt.Fprintf(os.Stderr, "Output format is fasta; serializing\n")
		if err := handleFastaOutput(sequences, *output); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err.Error())
			os.Exit(1)
		}
	case formats.TNT_FORMAT:
		fmt.Fprintf(os.Stderr, "Output format is TNT; serializing\n")
		context := TNTContext{
			Title:    *tntTitle,
			Outgroup: *outgroup,
		}
		if err := handleTNTOutput(context, sequences, *output); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err.Error())
			os.Exit(1)
		}
	default:
		fmt.Fprintf(os.Stderr, "Unknown output format '%s'", inputFormat)
		os.Exit(1)
	}

}
<|endoftext|>


<commit_before>/*
Command csvsplit splits a .csv into multiple, smaller files.

Resulting files will be saved as 1.csv, 2.csv, etc.  in the currect directory, unless the -output flag is used.

Install

Requires Go to be installed first, https://golang.org/doc/install.

	$ go get github.com/JeffPaine/csvsplit

Flags

Basic usage: csvsplit -records <number of records> <file>

	-records Number of records per file
	-output Output filename / path (optional)
	-headers Number of header lines in the input file to add to each ouput file (optional, default=0)

Examples

Split file.csv into files with 300 records a piece.
	$ csvplit -records 300 file.csv

Accept csv data from stdin.
	$ cat file.csv | csvsplit -records 20

Split file.csv into files with 40 records a piece and two header lines (preserved in all files).
	$ csvplit -records 40 -headers 2 file.csv

You can use the -output flag to customize the resulting filenames.
The below will generate custom_filename-001.csv, custom_filename-002.csv, etc..
	$ csvsplit -records 20 -output custom_filename- file.csv

Split file.csv into files with 37 records a piece into the subfolder 'stuff'.
	$ csvplit -records 37 -output stuff/ file.csv
*/
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
)

var (
	flagRecords = flag.Int("records", 0, "The number of records per output file")
	flagOutput  = flag.String("output", "", "Filename / path of the output file (leave blank for current directory)")
	flagHeaders = flag.Int("headers", 0, "Number of header lines in the input file to preserve in each output file")
)

func main() {
	flag.Parse()

	// Sanity check command line flags.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "usage: csvsplit [options] -records <number of records> <file>")
		flag.PrintDefaults()
		os.Exit(1)
	}
	if *flagRecords < 1 {
		fmt.Fprintln(os.Stderr, "-records must be > 1")
		flag.Usage()
	}
	if *flagHeaders < 0 {
		fmt.Fprintln(os.Stderr, "-headers must be > 0")
		flag.Usage()
	}
	if *flagHeaders >= *flagRecords {
		fmt.Fprintln(os.Stderr, "-headers must be >= -records")
		flag.Usage()
	}

	// Get input from a given file or stdin
	var r *csv.Reader
	if len(flag.Args()) == 1 {
		f, err := os.Open(flag.Args()[0])
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		r = csv.NewReader(f)
	} else {
		r = csv.NewReader(os.Stdin)
	}

	// Read the input .csv file line by line. Save to a new file after reaching
	// the amount of records prescribed by the -records flag.
	var recs [][]string
	count := 1
	for {
		record, err := r.Read()
		if err == io.EOF {
			save(&recs, count)
			break
		} else if err != nil {
			log.Fatal(err)
		}

		recs = append(recs, record)
		if len(recs) == *flagRecords {
			save(&recs, count)
			// Reset records to include just the header lines (if any)
			recs = recs[:*flagHeaders]
			count++
		}
	}
}

// save() saves the given *[][]string of csv data to a .csv file. Files are named
// sequentially in the form of 1.csv, 2.csv, etc.
func save(recs *[][]string, c int) {
	name := fmt.Sprintf("%v%d%v", *flagOutput, c, ".csv")

	// Make sure we don't overwrite existing files
	if _, err := os.Stat(name); err == nil {
		log.Fatal("file exists: ", name)
	}

	// If a directory is specified, make sure that directory exists
	if filepath.Dir(*flagOutput) != "." {
		_, err := os.Stat(filepath.Dir(*flagOutput))
		if err != nil {
			log.Fatal("no such directory:", *flagOutput)
		}
	}

	f, err := os.Create(name)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	w := csv.NewWriter(f)
	w.WriteAll(*recs)
}
<commit_msg>docs formatting<commit_after>/*
Command csvsplit splits a .csv into multiple, smaller files.

Resulting files will be saved as 1.csv, 2.csv, etc.  in the currect directory, unless the -output flag is used.

Install

Requires Go to be installed first, https://golang.org/doc/install.

	$ go get github.com/JeffPaine/csvsplit

Flags

Basic usage: csvsplit -records <number of records> <file>

	-records
Number of records per file

	-output
Output filename / path (optional)

	-headers
Number of header lines in the input file to add to each ouput file (optional, default=0)

Examples

Split file.csv into files with 300 records a piece.
	$ csvplit -records 300 file.csv

Accept csv data from stdin.
	$ cat file.csv | csvsplit -records 20

Split file.csv into files with 40 records a piece and two header lines (preserved in all files).
	$ csvplit -records 40 -headers 2 file.csv

You can use the -output flag to customize the resulting filenames.
The below will generate custom_filename-001.csv, custom_filename-002.csv, etc..
	$ csvsplit -records 20 -output custom_filename- file.csv

Split file.csv into files with 37 records a piece into the subfolder 'stuff'.
	$ csvplit -records 37 -output stuff/ file.csv
*/
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
)

var (
	flagRecords = flag.Int("records", 0, "The number of records per output file")
	flagOutput  = flag.String("output", "", "Filename / path of the output file (leave blank for current directory)")
	flagHeaders = flag.Int("headers", 0, "Number of header lines in the input file to preserve in each output file")
)

func main() {
	flag.Parse()

	// Sanity check command line flags.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "usage: csvsplit [options] -records <number of records> <file>")
		flag.PrintDefaults()
		os.Exit(1)
	}
	if *flagRecords < 1 {
		fmt.Fprintln(os.Stderr, "-records must be > 1")
		flag.Usage()
	}
	if *flagHeaders < 0 {
		fmt.Fprintln(os.Stderr, "-headers must be > 0")
		flag.Usage()
	}
	if *flagHeaders >= *flagRecords {
		fmt.Fprintln(os.Stderr, "-headers must be >= -records")
		flag.Usage()
	}

	// Get input from a given file or stdin
	var r *csv.Reader
	if len(flag.Args()) == 1 {
		f, err := os.Open(flag.Args()[0])
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		r = csv.NewReader(f)
	} else {
		r = csv.NewReader(os.Stdin)
	}

	// Read the input .csv file line by line. Save to a new file after reaching
	// the amount of records prescribed by the -records flag.
	var recs [][]string
	count := 1
	for {
		record, err := r.Read()
		if err == io.EOF {
			save(&recs, count)
			break
		} else if err != nil {
			log.Fatal(err)
		}

		recs = append(recs, record)
		if len(recs) == *flagRecords {
			save(&recs, count)
			// Reset records to include just the header lines (if any)
			recs = recs[:*flagHeaders]
			count++
		}
	}
}

// save() saves the given *[][]string of csv data to a .csv file. Files are named
// sequentially in the form of 1.csv, 2.csv, etc.
func save(recs *[][]string, c int) {
	name := fmt.Sprintf("%v%d%v", *flagOutput, c, ".csv")

	// Make sure we don't overwrite existing files
	if _, err := os.Stat(name); err == nil {
		log.Fatal("file exists: ", name)
	}

	// If a directory is specified, make sure that directory exists
	if filepath.Dir(*flagOutput) != "." {
		_, err := os.Stat(filepath.Dir(*flagOutput))
		if err != nil {
			log.Fatal("no such directory:", *flagOutput)
		}
	}

	f, err := os.Create(name)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	w := csv.NewWriter(f)
	w.WriteAll(*recs)
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2016 Maciej Borzecki

// 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 main

import (
	"github.com/codegangsta/cli"
	"log"
	"os"
)

func doImport(c *cli.Context) {
	log.Printf("import")

	srcfile := c.Args().First()
	log.Fatal(runImport(srcfile))
}

func doDaemon(c *cli.Context) {
	log.Printf("starting daemon")
	log.Fatal(runDaemon())
}

func main() {
	log.Printf("starting up")

	app := cli.NewApp()
	app.Name = "q3stats"
	app.HideHelp = true
	app.HideVersion = true
	app.Commands = []cli.Command{
		{
			Name:   "import",
			Usage:  "import match logs",
			Action: doImport,
		},
		{
			Name:   "daemon",
			Usage:  "run daemon",
			Action: doDaemon,
		},
	}

	app.Run(os.Args)
}
<commit_msg>main: remove startup log<commit_after>// The MIT License (MIT)

// Copyright (c) 2016 Maciej Borzecki

// 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 main

import (
	"github.com/codegangsta/cli"
	"log"
	"os"
)

func doImport(c *cli.Context) {
	log.Printf("import")

	srcfile := c.Args().First()
	log.Fatal(runImport(srcfile))
}

func doDaemon(c *cli.Context) {
	log.Printf("starting daemon")
	log.Fatal(runDaemon())
}

func main() {
	app := cli.NewApp()
	app.Name = "q3stats"
	app.HideHelp = true
	app.HideVersion = true
	app.Commands = []cli.Command{
		{
			Name:   "import",
			Usage:  "import match logs",
			Action: doImport,
		},
		{
			Name:   "daemon",
			Usage:  "run daemon",
			Action: doDaemon,
		},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"sync"
)

// TODO: more checks

func processFile(filename string, report chan *Warning) int {
	fm, err := NewFileMap(filename)
	check(err)

	count := 0

	for _, term := range ShouldNotExist() {
		for _, warn := range term.Match(fm) {
			report <- warn
			count++
		}
	}

	return count
}

func main() {
	report := make(chan *Warning, 64)
	wg := sync.WaitGroup{}

	for _, filename := range os.Args[1:] {
		wg.Add(1)
		go func(fn string) {
			defer wg.Done()
			processFile(fn, report)
		}(filename)
	}

	go func() {
		wg.Wait()
		close(report)
	}()

	// TODO: better output
	// TODO: select output format on command line - should include JSON fmt from: http://steelbrain.me/linter/types/linter-message-v2.html
	// TODO: max warnings per file
	// TODO: max warnings overall
	// TODO: option to sort warnings before output
	for warn := range report {
		fmt.Printf("%s:%d:%d:warning: %s\n",
			warn.Filename, warn.Row+1, warn.Col+1, warn.Msg,
		)
	}
}
<commit_msg>Start verbose option for extra info on stderr<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sync"
)

// TODO: more checks

func processFile(filename string, report chan *Warning) int {
	fm, err := NewFileMap(filename)
	check(err)

	count := 0

	for _, term := range ShouldNotExist() {
		for _, warn := range term.Match(fm) {
			report <- warn
			count++
		}
	}

	return count
}

func main() {
	// If it should always be printed, we use log. If it should only be printed
	// verbose=true, then we use verb
	var verbose = true // TODO: from command line
	var verb *log.Logger
	if verbose {
		verb = log.New(os.Stderr, "", 0)
	} else {
		verb = log.New(ioutil.Discard, "", 0)
	}

	verb.Printf("Verbose mode: ON\n")

	wg := sync.WaitGroup{}
	report := make(chan *Warning, 64)

	for _, filename := range os.Args[1:] {
		verb.Printf("FILE: %s\n", filename)
		wg.Add(1)
		go func(fn string) {
			defer wg.Done()
			processFile(fn, report)
		}(filename)
	}

	go func() {
		wg.Wait()
		close(report)
	}()

	// TODO: better output
	// TODO: select output format on command line - should include JSON fmt from: http://steelbrain.me/linter/types/linter-message-v2.html
	// TODO: max warnings per file
	// TODO: max warnings overall
	// TODO: option to sort warnings before output
	for warn := range report {
		fmt.Printf("%s:%d:%d:warning: %s\n",
			warn.Filename, warn.Row+1, warn.Col+1, warn.Msg,
		)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/jlaffaye/ftp"
	"golang.org/x/crypto/bcrypt"
)

var (
	rpcListenPort = flag.Int("rpc-listen-port", 7800, "Specify a port number for JSON-RPC server to listen to. Possible values: 1024-65535")
	rpcSecret     = flag.String("rpc-secret", "", "Set RPC secret authorization token (required)")

	n = flag.Int("n", 4, "Number of connections to use when downloading single file. Possible values: 1-100")
	o = flag.String("o", "", "Output directory (optional, default value is the current working directory)")
	p = flag.Int("p", 1, "Number of files to download in parallel when mirroring directories. Possible values: 1-10")

	// Info is used for logging information.
	Info = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)

	// Error is used for logging errors.
	Error = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)

	errMissingURL           = errors.New("No URL specified in a request")
	errProtocolMismatch     = errors.New("Only FTP downloads are supported")
	errInvalidRequestFormat = errors.New("Invalid request format")
	errTokenMismatch        = errors.New("Secret token does not match")
	errUnauthorized         = errors.New("Missing or invalid credentials")
)

// Request represents single request for mirroring one FTP directory or a file.
type Request struct {
	Path     string `json:"path"`
	Username string `json:"username"`
	Password string `json:"password"`
	Secret   string `json:"secret"`
}

// Response represents response to a client with ID for a created job or error message in case of error.
type Response struct {
	ID      string `json:"id"`
	Message string `json:"message"`
}

// Handler implements http.Handler interface and processes download requests sequentially.
type Handler struct {
	Jobs        chan *Job
	HashedToken []byte
}

// Job is single download request with associated ID and LFTP command.
type Job struct {
	ID      string
	Command *exec.Cmd
}

func (request *Request) extractURL() (*url.URL, error) {
	if request.Path == "" {
		return nil, errMissingURL
	}

	url, err := url.Parse(request.Path)

	if err != nil {
		return nil, fmt.Errorf("Invalid URL: %s", request.Path)
	}

	if url.Scheme != "ftp" {
		return nil, errProtocolMismatch
	}

	return url, nil
}

func makeLftpCmd(path string) string {
	escaped := "/"

	if path != "" {
		escaped = strings.Replace(path, "\"", "\\\"", -1)
	}

	if strings.HasSuffix(path, "/") {
		return fmt.Sprintf("mirror --parallel=%d --use-pget-n=%d \"%s\" && exit", *p, *n, escaped)
	}

	return fmt.Sprintf("pget -n %d \"%s\" && exit", *n, escaped)
}

func makeCmd(url *url.URL, username, password string) *exec.Cmd {
	lftpCmd := makeLftpCmd(url.Path)
	var args []string

	if username != "" && password != "" {
		args = []string{"--user", username, "--password", password, "-e", lftpCmd, url.Host}
	} else {
		args = []string{"-e", lftpCmd, url.Host}
	}

	cmd := exec.Command("lftp", args...)

	cmd.Dir = *o
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	return cmd
}

func generateID() string {
	b := make([]byte, 32)

	if _, err := rand.Read(b); err != nil {
		panic("Random number generator failed")
	}

	return base64.StdEncoding.EncodeToString(b)
}

func (handler *Handler) processRequest(r *http.Request) (string, error) {
	id := generateID()
	Info.Printf("Received download request %s from %s\n", id, r.RemoteAddr)

	var request Request
	decoder := json.NewDecoder(r.Body)

	if err := decoder.Decode(&request); err != nil {
		return "", errInvalidRequestFormat
	}

	if err := bcrypt.CompareHashAndPassword(handler.HashedToken, []byte(request.Secret)); err != nil {
		return "", errTokenMismatch
	}

	Info.Printf("Download request %s has URL %s\n", id, request.Path)
	url, err := request.extractURL()

	if err != nil {
		return "", err
	}

	host, port, err := net.SplitHostPort(url.Host)

	if err != nil {
		host, port = url.Host, strconv.Itoa(21)
	}

	conn, err := ftp.DialTimeout(net.JoinHostPort(host, port), 5*time.Second)

	if err != nil {
		return "", fmt.Errorf("Unable to connect to FTP server at %s", url.Host)
	}

	if request.Username != "" && request.Password != "" {
		err = conn.Login(request.Username, request.Password)
	} else {
		err = conn.Login("anonymous", "anonymous")
	}

	if err != nil {
		return "", errUnauthorized
	}

	conn.Logout()

	cmd := makeCmd(url, request.Username, request.Password)
	job := Job{ID: id, Command: cmd}

	go func() {
		handler.Jobs <- &job
	}()

	return id, nil
}

func (handler *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	id, err := handler.processRequest(r)

	if err == nil {
		json.NewEncoder(w).Encode(Response{ID: id})
		return
	}

	Error.Printf("Invalid request %s: %s\n", id, err)
	status := http.StatusBadRequest

	if err == errUnauthorized {
		status = http.StatusUnauthorized
	}

	w.WriteHeader(status)
	json.NewEncoder(w).Encode(Response{Message: err.Error()})
}

func (handler *Handler) worker() {
	for job := range handler.Jobs {
		Info.Printf("Begin LFTP output for request %s", job.ID)
		err := job.Command.Run()
		Info.Printf("End LFTP output for request %s", job.ID)

		if err != nil {
			Error.Printf("Failed to execute request %s with error: %v\n", job.ID, err)
		} else {
			Info.Printf("Request %s completed", job.ID)
		}
	}
}

func getOutputDir(dir string) (string, error) {
	var err error

	if dir == "" {
		if dir, err = os.Getwd(); err != nil {
			return "", err
		}
	}

	abs, err := filepath.Abs(dir)

	if err != nil {
		return "", err
	}

	file, err := os.Stat(abs)

	if err != nil {
		return "", err
	}

	if !file.IsDir() {
		return "", fmt.Errorf("%s is not a directory", abs)
	}

	return abs, nil
}

func main() {
	flag.Parse()

	if (*rpcListenPort < 1024 || *rpcListenPort > 65535) || *rpcSecret == "" {
		flag.Usage()
		os.Exit(1)
	}

	if *n < 1 || *n > 100 || *p < 1 || *p > 10 {
		flag.Usage()
		os.Exit(1)
	}

	if dir, err := getOutputDir(*o); err != nil {
		log.Fatal(err)
	} else {
		*o = dir
	}

	hashedToken, err := bcrypt.GenerateFromPassword([]byte(*rpcSecret), bcrypt.DefaultCost)

	if err != nil {
		log.Fatal("bcrypt failed to generate hashed token")
	}

	if _, err := exec.LookPath("lftp"); err != nil {
		log.Fatal("LFTP not found")
	}

	handler := &Handler{
		Jobs:        make(chan *Job, 10),
		HashedToken: hashedToken,
	}

	http.Handle("/jsonrpc", handler)
	go handler.worker()

	Info.Printf("Starting LFTP server on port %d\n", *rpcListenPort)
	Info.Printf("Output directory is %s\n", *o)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *rpcListenPort), nil))
}
<commit_msg>Remove ID from invalid request message (ID is not available at that time)<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/jlaffaye/ftp"
	"golang.org/x/crypto/bcrypt"
)

var (
	rpcListenPort = flag.Int("rpc-listen-port", 7800, "Specify a port number for JSON-RPC server to listen to. Possible values: 1024-65535")
	rpcSecret     = flag.String("rpc-secret", "", "Set RPC secret authorization token (required)")

	n = flag.Int("n", 4, "Number of connections to use when downloading single file. Possible values: 1-100")
	o = flag.String("o", "", "Output directory (optional, default value is the current working directory)")
	p = flag.Int("p", 1, "Number of files to download in parallel when mirroring directories. Possible values: 1-10")

	// Info is used for logging information.
	Info = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)

	// Error is used for logging errors.
	Error = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)

	errMissingURL           = errors.New("No URL specified in a request")
	errProtocolMismatch     = errors.New("Only FTP downloads are supported")
	errInvalidRequestFormat = errors.New("Invalid request format")
	errTokenMismatch        = errors.New("Secret token does not match")
	errUnauthorized         = errors.New("Missing or invalid credentials")
)

// Request represents single request for mirroring one FTP directory or a file.
type Request struct {
	Path     string `json:"path"`
	Username string `json:"username"`
	Password string `json:"password"`
	Secret   string `json:"secret"`
}

// Response represents response to a client with ID for a created job or error message in case of error.
type Response struct {
	ID      string `json:"id"`
	Message string `json:"message"`
}

// Handler implements http.Handler interface and processes download requests sequentially.
type Handler struct {
	Jobs        chan *Job
	HashedToken []byte
}

// Job is single download request with associated ID and LFTP command.
type Job struct {
	ID      string
	Command *exec.Cmd
}

func (request *Request) extractURL() (*url.URL, error) {
	if request.Path == "" {
		return nil, errMissingURL
	}

	url, err := url.Parse(request.Path)

	if err != nil {
		return nil, fmt.Errorf("Invalid URL: %s", request.Path)
	}

	if url.Scheme != "ftp" {
		return nil, errProtocolMismatch
	}

	return url, nil
}

func makeLftpCmd(path string) string {
	escaped := "/"

	if path != "" {
		escaped = strings.Replace(path, "\"", "\\\"", -1)
	}

	if strings.HasSuffix(path, "/") {
		return fmt.Sprintf("mirror --parallel=%d --use-pget-n=%d \"%s\" && exit", *p, *n, escaped)
	}

	return fmt.Sprintf("pget -n %d \"%s\" && exit", *n, escaped)
}

func makeCmd(url *url.URL, username, password string) *exec.Cmd {
	lftpCmd := makeLftpCmd(url.Path)
	var args []string

	if username != "" && password != "" {
		args = []string{"--user", username, "--password", password, "-e", lftpCmd, url.Host}
	} else {
		args = []string{"-e", lftpCmd, url.Host}
	}

	cmd := exec.Command("lftp", args...)

	cmd.Dir = *o
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	return cmd
}

func generateID() string {
	b := make([]byte, 32)

	if _, err := rand.Read(b); err != nil {
		panic("Random number generator failed")
	}

	return base64.StdEncoding.EncodeToString(b)
}

func (handler *Handler) processRequest(r *http.Request) (string, error) {
	id := generateID()
	Info.Printf("Received download request %s from %s\n", id, r.RemoteAddr)

	var request Request
	decoder := json.NewDecoder(r.Body)

	if err := decoder.Decode(&request); err != nil {
		return "", errInvalidRequestFormat
	}

	if err := bcrypt.CompareHashAndPassword(handler.HashedToken, []byte(request.Secret)); err != nil {
		return "", errTokenMismatch
	}

	Info.Printf("Download request %s has URL %s\n", id, request.Path)
	url, err := request.extractURL()

	if err != nil {
		return "", err
	}

	host, port, err := net.SplitHostPort(url.Host)

	if err != nil {
		host, port = url.Host, strconv.Itoa(21)
	}

	conn, err := ftp.DialTimeout(net.JoinHostPort(host, port), 5*time.Second)

	if err != nil {
		return "", fmt.Errorf("Unable to connect to FTP server at %s", url.Host)
	}

	if request.Username != "" && request.Password != "" {
		err = conn.Login(request.Username, request.Password)
	} else {
		err = conn.Login("anonymous", "anonymous")
	}

	if err != nil {
		return "", errUnauthorized
	}

	conn.Logout()

	cmd := makeCmd(url, request.Username, request.Password)
	job := Job{ID: id, Command: cmd}

	go func() {
		handler.Jobs <- &job
	}()

	return id, nil
}

func (handler *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	id, err := handler.processRequest(r)

	if err == nil {
		json.NewEncoder(w).Encode(Response{ID: id})
		return
	}

	Error.Printf("Invalid request received: %s\n", err)
	status := http.StatusBadRequest

	if err == errUnauthorized {
		status = http.StatusUnauthorized
	}

	w.WriteHeader(status)
	json.NewEncoder(w).Encode(Response{Message: err.Error()})
}

func (handler *Handler) worker() {
	for job := range handler.Jobs {
		Info.Printf("Begin LFTP output for request %s", job.ID)
		err := job.Command.Run()
		Info.Printf("End LFTP output for request %s", job.ID)

		if err != nil {
			Error.Printf("Failed to execute request %s with error: %v\n", job.ID, err)
		} else {
			Info.Printf("Request %s completed", job.ID)
		}
	}
}

func getOutputDir(dir string) (string, error) {
	var err error

	if dir == "" {
		if dir, err = os.Getwd(); err != nil {
			return "", err
		}
	}

	abs, err := filepath.Abs(dir)

	if err != nil {
		return "", err
	}

	file, err := os.Stat(abs)

	if err != nil {
		return "", err
	}

	if !file.IsDir() {
		return "", fmt.Errorf("%s is not a directory", abs)
	}

	return abs, nil
}

func main() {
	flag.Parse()

	if (*rpcListenPort < 1024 || *rpcListenPort > 65535) || *rpcSecret == "" {
		flag.Usage()
		os.Exit(1)
	}

	if *n < 1 || *n > 100 || *p < 1 || *p > 10 {
		flag.Usage()
		os.Exit(1)
	}

	if dir, err := getOutputDir(*o); err != nil {
		log.Fatal(err)
	} else {
		*o = dir
	}

	hashedToken, err := bcrypt.GenerateFromPassword([]byte(*rpcSecret), bcrypt.DefaultCost)

	if err != nil {
		log.Fatal("bcrypt failed to generate hashed token")
	}

	if _, err := exec.LookPath("lftp"); err != nil {
		log.Fatal("LFTP not found")
	}

	handler := &Handler{
		Jobs:        make(chan *Job, 10),
		HashedToken: hashedToken,
	}

	http.Handle("/jsonrpc", handler)
	go handler.worker()

	Info.Printf("Starting LFTP server on port %d\n", *rpcListenPort)
	Info.Printf("Output directory is %s\n", *o)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *rpcListenPort), nil))
}
<|endoftext|>


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

package main

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

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

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

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

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

/* This reads both HTTP requests and HTTP responses in two separate streams */
func (h *TCPStream) runOut() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64
	for {
		req, err := http.ReadRequest(buf)
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			return
		} else if err != nil {
			log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
		} else {
			// bodyBytes := tcpreader.DiscardBytesToEOF(req.Body)
			req.Body.Close()
			err = h.storage.SentRequest(h.bidikey, reqID, time.Now(), req)
			if err != nil {
				log.Println("Error storing request", err)
			}

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

func (h *TCPStream) runIn() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64
	for {

		/* Don't start reading a response if no data is available */
		_, err := buf.Peek(1)
		if err == io.EOF {
			return
		}
		/* Data available, read response */
		resp, err := http.ReadResponse(buf, nil) // TODO: request
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			return
		} else if err != nil {
			log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
		} else {
			_, err := tcpreader.DiscardBytesToFirstError(resp.Body)
			if err != nil && err != io.EOF {
				log.Println(err)
			}
			resp.Body.Close()
			err = h.storage.ReceivedResponse(h.bidikey, reqID, time.Now(), resp)
			if err != nil {
				log.Println("Error storing response", err)
			}
			reqID++
			// fmt.Print(".")
			//log.Println("Received response from stream", h.netFlow, h.tcpFlow,
			//	":", resp, "with", bodyBytes, "bytes in response body")
		}

		/* Match the response with the next request */
	}

}

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

func (h *httpStreamFactory) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {

	/* First the outgoing stream, then the incoming stream */
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

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

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

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

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

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

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

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

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

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

type Assembler struct {
	assembler *tcpassembly.Assembler
}

func NewAssembler(streamPool *tcpassembly.StreamPool) *Assembler {
	return &Assembler{tcpassembly.NewAssembler(streamPool)}
}

/* PFF, a lot of abstractions that make things more difficult then they should be! */
func (a *Assembler) AssembleWithTimestamp(netFlow gopacket.Flow, t *layers.TCP,
	timestamp time.Time) {
	a.assembler.AssembleWithTimestamp(netFlow, t, timestamp)
}

/* Wait for a couple of seconds T, just enough to get the events handled by the
   main function. */
func wait_for_responses_to_arrive(t time.Duration) (timeout chan bool) {
	timeout = make(chan bool, 1)
	go func() {
		time.Sleep(t * time.Second)
		timeout <- true
	}()
	return timeout
}

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

	log.Printf("starting capture on interface %q", *iface)

	// Setup packet capture
	handle, err := pcap.OpenLive(*iface, 1600, true, pcap.BlockForever)
	if err != nil {
		panic(err)
	} else if err := handle.SetBPFFilter("tcp and port 80"); err != nil {
		panic(err)
	}

	/* TODO: we need a storage layer per goroutine! */
	// Set up storage layer
	storage, err := NewStorage()
	if err != nil {
		panic(err)
	}

	// Set up assembly
	streamFactory := &httpStreamFactory{bidiStreams: make(map[uint64]*BidiStream),
		storage: storage}
	streamPool := tcpassembly.NewStreamPool(streamFactory)
	assembler := NewAssembler(streamPool)

	// Setup CTRL-C handler
	ctrlc := make(chan os.Signal, 1)
	signal.Notify(ctrlc, os.Interrupt)

	pid := uint32(0)
	var cmd_done chan error
	var start_time time.Time
	if *launchCmd != "" {
		s := *launchCmd
		//		n := strings.Index(*launchCmd, " ")
		//		fmt.Println("Launching ", s[0:n], " --- ", s[n:])
		//		cmd := exec.Command(s[0:n], strings.Split(s[n:], " "))
		args := strings.Split(s, " ")
		cmd := exec.Command(args[0], args[1:]...)
		start_time = time.Now()
		err := cmd.Start()
		if err != nil {
			panic(err)
		}
		cmd_done = make(chan error, 1)
		go func() {
			cmd_done <- cmd.Wait()
		}()
		pid = uint32(cmd.Process.Pid)
		fmt.Println("PID: ", cmd.Process.Pid)
	}

	netDescSource := NewOSXNetDescSource()
	descriptors := netDescSource.Descriptors()

	log.Println("reading in packets. Press CTRL-C to end and report.")
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets := packetSource.Packets()

	var timeout chan bool

	for {
		select {
		case netDesc := <-descriptors:
			if netDesc.Pid == pid {
				fmt.Println("event received ", netDesc)
			}
		case packet := <-packets:
			if *logAllPackets {
				log.Println(packet)
			}

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

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

			log.Println("Waiting for the remaining responses to arrive.")

			/* Wait for a couple of seconds, just enough to get the events
			   handled by the main function. */
			timeout = wait_for_responses_to_arrive(2)

		case <-ctrlc:
			/* Don't wait. */
			timeout = wait_for_responses_to_arrive(0)

		case <-timeout:
			reporting, err := NewReporting()
			if err != nil {
				panic(err)
			}

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

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

package main

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

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

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

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

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

/* This reads both HTTP requests and HTTP responses in two separate streams */
func (h *TCPStream) runOut() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64
	for {
		req, err := http.ReadRequest(buf)
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			return
		} else if err != nil {
			log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
		} else {
			// bodyBytes := tcpreader.DiscardBytesToEOF(req.Body)
			req.Body.Close()
			err = h.storage.SentRequest(h.bidikey, reqID, time.Now(), req)
			if err != nil {
				log.Println("Error storing request", err)
			}

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

func (h *TCPStream) runIn() {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64
	for {

		/* Don't start reading a response if no data is available */
		_, err := buf.Peek(1)
		if err == io.EOF {
			return
		}
		/* Data available, read response */
		resp, err := http.ReadResponse(buf, nil) // TODO: request
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			return
		} else if err != nil {
			log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
		} else {
			_, err := tcpreader.DiscardBytesToFirstError(resp.Body)
			if err != nil && err != io.EOF {
				log.Println(err)
			}
			resp.Body.Close()
			err = h.storage.ReceivedResponse(h.bidikey, reqID, time.Now(), resp)
			if err != nil {
				log.Println("Error storing response", err)
			}
			reqID++
			// fmt.Print(".")
			//log.Println("Received response from stream", h.netFlow, h.tcpFlow,
			//	":", resp, "with", bodyBytes, "bytes in response body")
		}

		/* Match the response with the next request */
	}

}

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

func (h *httpStreamFactory) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {

	/* First the outgoing stream, then the incoming stream */
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

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

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

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

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

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

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

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

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

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

type Assembler struct {
	assembler *tcpassembly.Assembler
}

func NewAssembler(streamPool *tcpassembly.StreamPool) *Assembler {
	return &Assembler{tcpassembly.NewAssembler(streamPool)}
}

/* PFF, a lot of abstractions that make things more difficult then they should be! */
func (a *Assembler) AssembleWithTimestamp(netFlow gopacket.Flow, t *layers.TCP,
	timestamp time.Time) {
	a.assembler.AssembleWithTimestamp(netFlow, t, timestamp)
}

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

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

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

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

	log.Printf("starting capture on interface %q", *iface)

	/* TODO: we need a storage layer per goroutine! */
	// Set up storage layer
	storage, err := NewStorage()
	if err != nil {
		panic(err)
	}

	// Set up assembly
	streamFactory := &httpStreamFactory{bidiStreams: make(map[uint64]*BidiStream),
		storage: storage}
	streamPool := tcpassembly.NewStreamPool(streamFactory)
	assembler := NewAssembler(streamPool)

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

	// Setup channel that reports all socket kernel events (Mac OS X only)
	netDescSource := NewOSXNetDescSource()
	descriptors := netDescSource.Descriptors()

	// Setup packet capture
	handle, err := pcap.OpenLive(*iface, 1600, true, pcap.BlockForever)
	if err != nil {
		panic(err)
	} else if err := handle.SetBPFFilter("tcp and port 80"); err != nil {
		panic(err)
	}
	log.Println("reading in packets. Press CTRL-C to end and report.")
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets := packetSource.Packets()

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

	var timeout chan bool
	for {
		select {
		case netDesc := <-descriptors:
			if netDesc.Pid == pid {
				fmt.Println("event received ", netDesc)
			}
		case packet := <-packets:
			if *logAllPackets {
				log.Println(packet)
			}

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

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

			log.Println("Waiting for the remaining responses to arrive.")

			/* Wait for a couple of seconds, just enough to get the events
			   handled by the main function. */
			timeout = create_timeout_channel(2)

		case <-ctrlc:
			/* Don't wait. */
			timeout = create_timeout_channel(0)

		case <-timeout:
			reporting, err := NewReporting()
			if err != nil {
				panic(err)
			}

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

			os.Exit(0)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/codegangsta/cli"
	"os"
)

func main() {
	app := cli.NewApp()
	app.Name = "dmlivewiki"
	app.Usage = "dmlivewiki helper"
	app.Author = `Qais "qaisjp" Patankar`
	app.Email = "me@qaisjp.com"
	app.Version = "1.0.6"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "force, f",
			Usage: "skip confirmation",
		},
		cli.BoolFlag{
			Name:  "delete",
			Usage: "instead of creating files, delete files",
		},
		cli.BoolFlag{
			Name:  "single, s",
			Usage: "parse the directory given, not the subdirectories",
		},
	}

	app.Commands = []cli.Command{
		{
			Name:   "checksum",
			Usage:  "perform a checksum of directories",
			Action: performChecksum,
		},
		{
			Name:   "generate",
			Usage:  "generate dirname.txt Infofile's for the passed directory",
			Action: generateInformation,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "tour",
					Usage: "required: the tour name for this directory",
				},
				cli.StringFlag{
					Name:  "tour-file",
					Usage: "file with list of tracks with alternate vocals",
				},
			},
		},
		{
			Name:   "wiki",
			Usage:  "generate dirname.wiki Wikifile's for the passed directory",
			Action: generateWikifiles,
		},
	}

	app.Run(os.Args)
}

var informationTemplate = `{{.Artist}}
{{.Date}}
{{.Album}}
{{.Tour}}

Lineage:

Notes:

This source is considered Source 1 for this date:
https://www.depechemode-live.com/wiki/{{wikiescape .Date}}_{{wikiescape .Album}}/Source_1

Track list:

{{range .Tracks}}{{.Prefix}}{{printf "%02d" .Index}}. [{{.Duration}}] {{.Title}}{{if .HasAlternateLeadVocalist}} (*){{end}}
{{end}}Total time: {{.Duration}}

Torrent downloaded from https://www.depechemode-live.com
`

var wikiTemplate = `== Notes ==

{{.Notes}}

== Listen ==

You can listen to this entire recording below.

<html5media>https://media.depechemode-live.com/stream/{{.FolderName}}/complete.m4a</html5media>

== Track list ==

{{range .Tracks}}#[{{.Duration}}] <sm2>https://media.depechemode-live.com/stream/{{.FolderName}}/{{printf "%02d" .Index}}.m4a</sm2> [[{{.Name}}]]{{if .HasAlternateLeadVocalist}} (*){{end}}
{{end}}*Total time: {{.Duration}}

== Lineage ==

{{.Lineage}}

== Download ==

*[https://depechemode-live.com/torrents/{{.FolderName}}.torrent Download via torrent] - FLAC {{.BPS}}-bit {{.SampleRate}} - {{.Size}}

[[Category:Audience recordings]]
[[Category:Source]]
[[Category:Streamable]]
`

var wikiRegex = `(?:.|[\r\n])+[\r\n]+Lineage: ((?:.|[\r\n]+)*)[\r\n]+Notes: ((?:.|[\r\n]+)*)[\r\n]+This source is considered(?:.|[\r\n]+)*Track list:[\r\n]+[\r\n]+((?:.|[\r\n]+)*)[\r\n]+Total time: (.*)`
<commit_msg>Remove trailing newline in text files<commit_after>package main

import (
	"github.com/codegangsta/cli"
	"os"
)

func main() {
	app := cli.NewApp()
	app.Name = "dmlivewiki"
	app.Usage = "dmlivewiki helper"
	app.Author = `Qais "qaisjp" Patankar`
	app.Email = "me@qaisjp.com"
	app.Version = "1.0.6"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "force, f",
			Usage: "skip confirmation",
		},
		cli.BoolFlag{
			Name:  "delete",
			Usage: "instead of creating files, delete files",
		},
		cli.BoolFlag{
			Name:  "single, s",
			Usage: "parse the directory given, not the subdirectories",
		},
	}

	app.Commands = []cli.Command{
		{
			Name:   "checksum",
			Usage:  "perform a checksum of directories",
			Action: performChecksum,
		},
		{
			Name:   "generate",
			Usage:  "generate dirname.txt Infofile's for the passed directory",
			Action: generateInformation,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "tour",
					Usage: "required: the tour name for this directory",
				},
				cli.StringFlag{
					Name:  "tour-file",
					Usage: "file with list of tracks with alternate vocals",
				},
			},
		},
		{
			Name:   "wiki",
			Usage:  "generate dirname.wiki Wikifile's for the passed directory",
			Action: generateWikifiles,
		},
	}

	app.Run(os.Args)
}

var informationTemplate = `{{.Artist}}
{{.Date}}
{{.Album}}
{{.Tour}}

Lineage:

Notes:

This source is considered Source 1 for this date:
https://www.depechemode-live.com/wiki/{{wikiescape .Date}}_{{wikiescape .Album}}/Source_1

Track list:

{{range .Tracks}}{{.Prefix}}{{printf "%02d" .Index}}. [{{.Duration}}] {{.Title}}{{if .HasAlternateLeadVocalist}} (*){{end}}
{{end}}Total time: {{.Duration}}

Torrent downloaded from https://www.depechemode-live.com`

var wikiTemplate = `== Notes ==

{{.Notes}}

== Listen ==

You can listen to this entire recording below.

<html5media>https://media.depechemode-live.com/stream/{{.FolderName}}/complete.m4a</html5media>

== Track list ==

{{range .Tracks}}#[{{.Duration}}] <sm2>https://media.depechemode-live.com/stream/{{.FolderName}}/{{printf "%02d" .Index}}.m4a</sm2> [[{{.Name}}]]{{if .HasAlternateLeadVocalist}} (*){{end}}
{{end}}*Total time: {{.Duration}}

== Lineage ==

{{.Lineage}}

== Download ==

*[https://depechemode-live.com/torrents/{{.FolderName}}.torrent Download via torrent] - FLAC {{.BPS}}-bit {{.SampleRate}} - {{.Size}}

[[Category:Audience recordings]]
[[Category:Source]]
[[Category:Streamable]]
`

var wikiRegex = `(?:.|[\r\n])+[\r\n]+Lineage: ((?:.|[\r\n]+)*)[\r\n]+Notes: ((?:.|[\r\n]+)*)[\r\n]+This source is considered(?:.|[\r\n]+)*Track list:[\r\n]+[\r\n]+((?:.|[\r\n]+)*)[\r\n]+Total time: (.*)`
<|endoftext|>


<commit_before>// csvsplit: Split a .csv into multiple files.
// https://github.com/JeffPaine/csvsplit
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
)

var (
	flagRecords = flag.Int("records", 0, "The number of records per output file")
	flagOutput  = flag.String("output", "", "Filename / path of the output file (leave blank for current directory)")
	flagHeaders = flag.Int("headers", 0, "Number of header lines in the input file to preserve in each output file")
)

func main() {
	flag.Parse()

	// Sanity check command line flags.
	checkFlags()

	// Get input from a given file or stdin
	var r *csv.Reader
	if len(flag.Args()) == 1 {
		f, err := os.Open(flag.Args()[0])
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		r = csv.NewReader(f)
	} else {
		r = csv.NewReader(os.Stdin)
	}

	// Read the input .csv file line by line. Save to a new file after reaching
	// the amount of records prescribed by the -records flag.
	var recs [][]string
	count := 1
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		} else if err != nil {
			log.Fatal(err)
		}

		recs = append(recs, record)
		if len(recs) == *flagRecords {
			save(&recs, count)
			// Reset records to include just the header lines (if any)
			recs = recs[:*flagHeaders]
			count++
		}
	}
	if len(recs) > 0 {
		save(&recs, count)
	}
}

// save() saves the given *[][]string of csv data to a .csv file. Files are named
// sequentially in the form of 1.csv, 2.csv, etc.
func save(recs *[][]string, c int) {
	name := fmt.Sprintf("%v%d%v", *flagOutput, c, ".csv")

	// Make sure we don't overwrite existing files
	if _, err := os.Stat(name); err == nil {
		log.Fatal("file exists: ", name)
	}

	// If a directory is specified, make sure that directory exists
	if filepath.Dir(*flagOutput) != "." {
		_, err := os.Stat(filepath.Dir(*flagOutput))
		if err != nil {
			log.Fatal("no such directory:", *flagOutput)
		}
	}

	f, err := os.Create(name)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	w := csv.NewWriter(f)
	w.WriteAll(*recs)
}

// checkFlags checks our command line flags for basic sanity.
func checkFlags() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "usage: csvsplit [options] -records <number of records> <file>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	if *flagRecords < 1 {
		fmt.Fprintln(os.Stderr, "-records must be > 1")
		flag.Usage()
	}

	if *flagHeaders < 0 {
		fmt.Fprintln(os.Stderr, "-headers must be > 0")
		flag.Usage()
	}

	if *flagHeaders >= *flagRecords {
		fmt.Fprintln(os.Stderr, "-headers must be >= -records")
		flag.Usage()
	}
}
<commit_msg>light refactor<commit_after>// csvsplit: Split a .csv into multiple files.
// https://github.com/JeffPaine/csvsplit
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
)

var (
	flagRecords = flag.Int("records", 0, "The number of records per output file")
	flagOutput  = flag.String("output", "", "Filename / path of the output file (leave blank for current directory)")
	flagHeaders = flag.Int("headers", 0, "Number of header lines in the input file to preserve in each output file")
)

func main() {
	flag.Parse()

	// Sanity check command line flags.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "usage: csvsplit [options] -records <number of records> <file>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	if *flagRecords < 1 {
		fmt.Fprintln(os.Stderr, "-records must be > 1")
		flag.Usage()
	}

	if *flagHeaders < 0 {
		fmt.Fprintln(os.Stderr, "-headers must be > 0")
		flag.Usage()
	}

	if *flagHeaders >= *flagRecords {
		fmt.Fprintln(os.Stderr, "-headers must be >= -records")
		flag.Usage()
	}

	// Get input from a given file or stdin
	var r *csv.Reader
	if len(flag.Args()) == 1 {
		f, err := os.Open(flag.Args()[0])
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		r = csv.NewReader(f)
	} else {
		r = csv.NewReader(os.Stdin)
	}

	// Read the input .csv file line by line. Save to a new file after reaching
	// the amount of records prescribed by the -records flag.
	var recs [][]string
	count := 1
	for {
		record, err := r.Read()
		if err == io.EOF {
			save(&recs, count)
			break
		} else if err != nil {
			log.Fatal(err)
		}

		recs = append(recs, record)
		if len(recs) == *flagRecords {
			save(&recs, count)
			// Reset records to include just the header lines (if any)
			recs = recs[:*flagHeaders]
			count++
		}
	}
}

// save() saves the given *[][]string of csv data to a .csv file. Files are named
// sequentially in the form of 1.csv, 2.csv, etc.
func save(recs *[][]string, c int) {
	name := fmt.Sprintf("%v%d%v", *flagOutput, c, ".csv")

	// Make sure we don't overwrite existing files
	if _, err := os.Stat(name); err == nil {
		log.Fatal("file exists: ", name)
	}

	// If a directory is specified, make sure that directory exists
	if filepath.Dir(*flagOutput) != "." {
		_, err := os.Stat(filepath.Dir(*flagOutput))
		if err != nil {
			log.Fatal("no such directory:", *flagOutput)
		}
	}

	f, err := os.Create(name)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	w := csv.NewWriter(f)
	w.WriteAll(*recs)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/volatiletech/abcweb/cmd"
)

const abcwebVersion = "3.1.0"

func main() {
	// Too much happens between here and cobra's argument handling, for
	// something so simple. Just do it immediately.
	if len(os.Args) > 1 && os.Args[1] == "--version" {
		fmt.Println("ABCWeb v" + abcwebVersion)
		return
	}

	cmd.Execute()
}
<commit_msg>fix abcweb version<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/volatiletech/abcweb/cmd"
)

const abcwebVersion = "3.0.2"

func main() {
	// Too much happens between here and cobra's argument handling, for
	// something so simple. Just do it immediately.
	if len(os.Args) > 1 && os.Args[1] == "--version" {
		fmt.Println("ABCWeb v" + abcwebVersion)
		return
	}

	cmd.Execute()
}
<|endoftext|>