<commit_before>/*                          _       _
 *__      _____  __ ___   ___  __ _| |_ ___
 *\ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
 * \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
 *  \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
 *
 * Copyright © 2016 - 2018 Weaviate. All rights reserved.
 * LICENSE: https://github.com/creativesoftwarefdn/weaviate/blob/develop/LICENSE.md
 * AUTHOR: Bob van Luijt (bob@kub.design)
 * See www.creativesoftwarefdn.org for details
 * Contact: @CreativeSofwFdn / bob@kub.design
 */

// Package graphqlapi provides the graphql endpoint for Weaviate
package graphqlapi

import (
	"fmt"

	dbconnector "github.com/creativesoftwarefdn/weaviate/database/connectors"
	"github.com/graphql-go/graphql"
)

var dbConnector dbconnector.DatabaseConnector

// Build the GraphQL schema based on
// 1) the static query structure (e.g. Get)
// 2) the (dynamic) database schema from Weaviate

func (g *graphQL) buildGraphqlSchema() error {
	rootFieldsObject, err := g.assembleFullSchema()

	if err != nil {
		return fmt.Errorf("could not build GraphQL schema, because: %v", err)
	}

	schemaObject := graphql.ObjectConfig{
		Name:        "WeaviateObj",
		Fields:      rootFieldsObject,
		Description: "Location of the root query",
	}

	// Run grahpql.NewSchema in a sub-closure, so that we can recover from panics.
	// We need to use panics to return errors deep inside the dynamic generation of the GraphQL schema,
	// inside the FieldThunks. There is _no_ way to bubble up an error besides panicking.
	func() {
		defer func() {
			if r := recover(); r != nil {
				var ok bool
				err, ok = r.(error) // can't shadow err here; we need the err from outside the function closure.

				if !ok {
					err = fmt.Errorf("%v", err)
				}
			}
		}()

		g.weaviateGraphQLSchema, err = graphql.NewSchema(graphql.SchemaConfig{
			Query: graphql.NewObject(schemaObject),
		})
	}()

	if err != nil {
		return fmt.Errorf("could not build GraphQL schema, because: %v", err)
	}
	return nil
}

func (g *graphQL) assembleFullSchema() (graphql.Fields, error) {
	// This map is used to store all the Thing and Action Objects, so that we can use them in references.
	getActionsAndThings := make(map[string]*graphql.Object)
	// this map is used to store all the Filter InputObjects, so that we can use them in references.
	filterOptions := make(map[string]*graphql.InputObject)

	localGetActions, err := g.genActionClassFieldsFromSchema(&getActionsAndThings)

	if err != nil {
		return nil, fmt.Errorf("failed to generate action fields from schema for local Get because: %v", err)
	}

	localGetThings, err := g.genThingClassFieldsFromSchema(&getActionsAndThings)

	// No things, nor actions. Bail out!
	if localGetActions == nil && localGetThings == nil {
		return nil, fmt.Errorf("There are not any Actions or Things classes defined yet.")
	}

	if err != nil {
		return nil, fmt.Errorf("failed to generate thing fields from schema for local Get because: %v", err)
	}

	var localGetMetaActions *graphql.Object

	if localGetActions != nil {
		classParentTypeIsAction := true
		localGetMetaActions, err = genMetaClassFieldsFromSchema(g.databaseSchema.ActionSchema.Schema.Classes, classParentTypeIsAction)
		if err != nil {
			return nil, fmt.Errorf("failed to generate action fields from schema for local MetaGet because: %v", err)
		}
	}

	var localGetMetaThings *graphql.Object
	if localGetThings != nil {
		classParentTypeIsAction := false
		localGetMetaThings, err = genMetaClassFieldsFromSchema(g.databaseSchema.ThingSchema.Schema.Classes, classParentTypeIsAction)
		if err != nil {
			return nil, fmt.Errorf("failed to generate thing fields from schema for local MetaGet because: %v", err)
		}
	}

	localGetObject := g.genThingsAndActionsFieldsForWeaviateLocalGetObj(localGetActions, localGetThings)

	localGetMetaObject := genThingsAndActionsFieldsForWeaviateLocalGetMetaObj(localGetMetaActions, localGetMetaThings)

	localGetAndGetMetaObject := genGetAndGetMetaFields(localGetObject, localGetMetaObject, filterOptions)

	localField := &graphql.Field{
		Type:        localGetAndGetMetaObject,
		Description: "Query a local Weaviate instance",
		Resolve: func(p graphql.ResolveParams) (interface{}, error) {
			result, err := dbConnector.GetGraph(p)
			return result, err
		},
	}

	rootFields := graphql.Fields{
		"Local":   localField,
		"Network": nil,
	}

	return rootFields, nil
}

// generate the static parts of the schema
func (g *graphQL) genThingsAndActionsFieldsForWeaviateLocalGetObj(localGetActions *graphql.Object, localGetThings *graphql.Object) *graphql.Object {
	getThingsAndActionFields := graphql.Fields{}

	if localGetActions != nil {
		getThingsAndActionFields["Actions"] = &graphql.Field{
			Name:        "WeaviateLocalGetActions",
			Description: "Get Actions on the Local Weaviate",
			Type:        localGetActions,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	if localGetThings != nil {
		getThingsAndActionFields["Things"] = &graphql.Field{
			Name:        "WeaviateLocalGetThings",
			Description: "Get Things on the Local Weaviate",
			Type:        localGetThings,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	getThingsAndActionFieldsObject := graphql.ObjectConfig{
		Name:        "WeaviateLocalGetObj",
		Fields:      getThingsAndActionFields,
		Description: "Type of Get function to get Things or Actions on the Local Weaviate",
	}

	return graphql.NewObject(getThingsAndActionFieldsObject)
}

func genThingsAndActionsFieldsForWeaviateLocalGetMetaObj(localGetMetaActions *graphql.Object, localGetMetaThings *graphql.Object) *graphql.Object {
	getMetaThingsAndActionFields := graphql.Fields{}

	if localGetMetaActions != nil {
		getMetaThingsAndActionFields["Actions"] = &graphql.Field{
			Name:        "WeaviateLocalGetMetaActions",
			Description: "Get Meta information about Actions on the Local Weaviate",
			Type:        localGetMetaActions,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	if localGetMetaThings != nil {
		getMetaThingsAndActionFields["Things"] = &graphql.Field{
			Name:        "WeaviateLocalGetMetaThings",
			Description: "Get Meta information about Things on the Local Weaviate",
			Type:        localGetMetaThings,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	getMetaThingsAndActionFieldsObject := graphql.ObjectConfig{
		Name:        "WeaviateLocalGetMetaObj",
		Fields:      getMetaThingsAndActionFields,
		Description: "Type of Get function to get meta information about Things or Actions on the Local Weaviate",
	}

	return graphql.NewObject(getMetaThingsAndActionFieldsObject)
}

func genGetAndGetMetaFields(localGetObject *graphql.Object, localGetMetaObject *graphql.Object, filterOptions map[string]*graphql.InputObject) *graphql.Object {
	filterFields := genFilterFields(filterOptions)
	getAndGetMetaFields := graphql.Fields{

		"Get": &graphql.Field{
			Name:        "WeaviateLocalGet",
			Type:        localGetObject,
			Description: "Get Things or Actions on the local weaviate",
			Args: graphql.FieldConfigArgument{
				"where": &graphql.ArgumentConfig{
					Description: "Filter options for the Get search, to convert the data to the filter input",
					Type: graphql.NewInputObject(
						graphql.InputObjectConfig{
							Name:        "WeaviateLocalGetWhereInpObj",
							Fields:      filterFields,
							Description: "Filter options for the Get search, to convert the data to the filter input",
						},
					),
				},
			},
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		},

		"GetMeta": &graphql.Field{
			Name:        "WeaviateLocalGetMeta",
			Type:        localGetMetaObject,
			Description: "Query to Get Meta information about the data in the local Weaviate instance",
			Args: graphql.FieldConfigArgument{
				"where": &graphql.ArgumentConfig{
					Description: "Filter options for the GetMeta search, to convert the data to the filter input",
					Type: graphql.NewInputObject(
						graphql.InputObjectConfig{
							Name:        "WeaviateLocalGetMetaWhereInpObj",
							Fields:      filterFields,
							Description: "Filter options for the GetMeta search, to convert the data to the filter input",
						},
					),
				},
			},
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		},
	}

	weaviateLocalObject := &graphql.ObjectConfig{
		Name:        "WeaviateLocalObj",
		Fields:      getAndGetMetaFields,
		Description: "Type of query on the local Weaviate",
	}

	return graphql.NewObject(*weaviateLocalObject)
}
<commit_msg>Simplify error handling during schema generation<commit_after>/*                          _       _
 *__      _____  __ ___   ___  __ _| |_ ___
 *\ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
 * \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
 *  \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
 *
 * Copyright © 2016 - 2018 Weaviate. All rights reserved.
 * LICENSE: https://github.com/creativesoftwarefdn/weaviate/blob/develop/LICENSE.md
 * AUTHOR: Bob van Luijt (bob@kub.design)
 * See www.creativesoftwarefdn.org for details
 * Contact: @CreativeSofwFdn / bob@kub.design
 */

// Package graphqlapi provides the graphql endpoint for Weaviate
package graphqlapi

import (
	"fmt"

	dbconnector "github.com/creativesoftwarefdn/weaviate/database/connectors"
	"github.com/graphql-go/graphql"
)

var dbConnector dbconnector.DatabaseConnector

// Build the GraphQL schema based on
// 1) the static query structure (e.g. Get)
// 2) the (dynamic) database schema from Weaviate

func (g *graphQL) buildGraphqlSchema() error {
	rootFieldsObject, err := g.assembleFullSchema()

	if err != nil {
		return fmt.Errorf("could not build GraphQL schema, because: %v", err)
	}

	schemaObject := graphql.ObjectConfig{
		Name:        "WeaviateObj",
		Fields:      rootFieldsObject,
		Description: "Location of the root query",
	}

	// Run grahpql.NewSchema in a sub-closure, so that we can recover from panics.
	// We need to use panics to return errors deep inside the dynamic generation of the GraphQL schema,
	// inside the FieldThunks. There is _no_ way to bubble up an error besides panicking.
	func() {
		defer func() {
			if r := recover(); r != nil {
				err = fmt.Errorf("%v", r)
			}
		}()

		g.weaviateGraphQLSchema, err = graphql.NewSchema(graphql.SchemaConfig{
			Query: graphql.NewObject(schemaObject),
		})
	}()

	if err != nil {
		return fmt.Errorf("could not build GraphQL schema, because: %v", err)
	}
	return nil
}

func (g *graphQL) assembleFullSchema() (graphql.Fields, error) {
	// This map is used to store all the Thing and Action Objects, so that we can use them in references.
	getActionsAndThings := make(map[string]*graphql.Object)
	// this map is used to store all the Filter InputObjects, so that we can use them in references.
	filterOptions := make(map[string]*graphql.InputObject)

	localGetActions, err := g.genActionClassFieldsFromSchema(&getActionsAndThings)

	if err != nil {
		return nil, fmt.Errorf("failed to generate action fields from schema for local Get because: %v", err)
	}

	localGetThings, err := g.genThingClassFieldsFromSchema(&getActionsAndThings)

	// No things, nor actions. Bail out!
	if localGetActions == nil && localGetThings == nil {
		return nil, fmt.Errorf("There are not any Actions or Things classes defined yet.")
	}

	if err != nil {
		return nil, fmt.Errorf("failed to generate thing fields from schema for local Get because: %v", err)
	}

	var localGetMetaActions *graphql.Object

	if localGetActions != nil {
		classParentTypeIsAction := true
		localGetMetaActions, err = genMetaClassFieldsFromSchema(g.databaseSchema.ActionSchema.Schema.Classes, classParentTypeIsAction)
		if err != nil {
			return nil, fmt.Errorf("failed to generate action fields from schema for local MetaGet because: %v", err)
		}
	}

	var localGetMetaThings *graphql.Object
	if localGetThings != nil {
		classParentTypeIsAction := false
		localGetMetaThings, err = genMetaClassFieldsFromSchema(g.databaseSchema.ThingSchema.Schema.Classes, classParentTypeIsAction)
		if err != nil {
			return nil, fmt.Errorf("failed to generate thing fields from schema for local MetaGet because: %v", err)
		}
	}

	localGetObject := g.genThingsAndActionsFieldsForWeaviateLocalGetObj(localGetActions, localGetThings)

	localGetMetaObject := genThingsAndActionsFieldsForWeaviateLocalGetMetaObj(localGetMetaActions, localGetMetaThings)

	localGetAndGetMetaObject := genGetAndGetMetaFields(localGetObject, localGetMetaObject, filterOptions)

	localField := &graphql.Field{
		Type:        localGetAndGetMetaObject,
		Description: "Query a local Weaviate instance",
		Resolve: func(p graphql.ResolveParams) (interface{}, error) {
			result, err := dbConnector.GetGraph(p)
			return result, err
		},
	}

	rootFields := graphql.Fields{
		"Local":   localField,
		"Network": nil,
	}

	return rootFields, nil
}

// generate the static parts of the schema
func (g *graphQL) genThingsAndActionsFieldsForWeaviateLocalGetObj(localGetActions *graphql.Object, localGetThings *graphql.Object) *graphql.Object {
	getThingsAndActionFields := graphql.Fields{}

	if localGetActions != nil {
		getThingsAndActionFields["Actions"] = &graphql.Field{
			Name:        "WeaviateLocalGetActions",
			Description: "Get Actions on the Local Weaviate",
			Type:        localGetActions,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	if localGetThings != nil {
		getThingsAndActionFields["Things"] = &graphql.Field{
			Name:        "WeaviateLocalGetThings",
			Description: "Get Things on the Local Weaviate",
			Type:        localGetThings,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	getThingsAndActionFieldsObject := graphql.ObjectConfig{
		Name:        "WeaviateLocalGetObj",
		Fields:      getThingsAndActionFields,
		Description: "Type of Get function to get Things or Actions on the Local Weaviate",
	}

	return graphql.NewObject(getThingsAndActionFieldsObject)
}

func genThingsAndActionsFieldsForWeaviateLocalGetMetaObj(localGetMetaActions *graphql.Object, localGetMetaThings *graphql.Object) *graphql.Object {
	getMetaThingsAndActionFields := graphql.Fields{}

	if localGetMetaActions != nil {
		getMetaThingsAndActionFields["Actions"] = &graphql.Field{
			Name:        "WeaviateLocalGetMetaActions",
			Description: "Get Meta information about Actions on the Local Weaviate",
			Type:        localGetMetaActions,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	if localGetMetaThings != nil {
		getMetaThingsAndActionFields["Things"] = &graphql.Field{
			Name:        "WeaviateLocalGetMetaThings",
			Description: "Get Meta information about Things on the Local Weaviate",
			Type:        localGetMetaThings,
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		}
	}

	getMetaThingsAndActionFieldsObject := graphql.ObjectConfig{
		Name:        "WeaviateLocalGetMetaObj",
		Fields:      getMetaThingsAndActionFields,
		Description: "Type of Get function to get meta information about Things or Actions on the Local Weaviate",
	}

	return graphql.NewObject(getMetaThingsAndActionFieldsObject)
}

func genGetAndGetMetaFields(localGetObject *graphql.Object, localGetMetaObject *graphql.Object, filterOptions map[string]*graphql.InputObject) *graphql.Object {
	filterFields := genFilterFields(filterOptions)
	getAndGetMetaFields := graphql.Fields{

		"Get": &graphql.Field{
			Name:        "WeaviateLocalGet",
			Type:        localGetObject,
			Description: "Get Things or Actions on the local weaviate",
			Args: graphql.FieldConfigArgument{
				"where": &graphql.ArgumentConfig{
					Description: "Filter options for the Get search, to convert the data to the filter input",
					Type: graphql.NewInputObject(
						graphql.InputObjectConfig{
							Name:        "WeaviateLocalGetWhereInpObj",
							Fields:      filterFields,
							Description: "Filter options for the Get search, to convert the data to the filter input",
						},
					),
				},
			},
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		},

		"GetMeta": &graphql.Field{
			Name:        "WeaviateLocalGetMeta",
			Type:        localGetMetaObject,
			Description: "Query to Get Meta information about the data in the local Weaviate instance",
			Args: graphql.FieldConfigArgument{
				"where": &graphql.ArgumentConfig{
					Description: "Filter options for the GetMeta search, to convert the data to the filter input",
					Type: graphql.NewInputObject(
						graphql.InputObjectConfig{
							Name:        "WeaviateLocalGetMetaWhereInpObj",
							Fields:      filterFields,
							Description: "Filter options for the GetMeta search, to convert the data to the filter input",
						},
					),
				},
			},
			Resolve: func(p graphql.ResolveParams) (interface{}, error) {
				result, err := dbConnector.GetGraph(p)
				return result, err
			},
		},
	}

	weaviateLocalObject := &graphql.ObjectConfig{
		Name:        "WeaviateLocalObj",
		Fields:      getAndGetMetaFields,
		Description: "Type of query on the local Weaviate",
	}

	return graphql.NewObject(*weaviateLocalObject)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strconv"
)

//go:generate stringer -type=instruction
type instruction int

const (
	turnOn instruction = iota
	turnOff
	toggle
)

type coordinates struct {
	xStart int
	yStart int
	xEnd   int
	yEnd   int
}

var lineRegexp = regexp.MustCompile("^(turn on|turn off|toggle) (\\d+),(\\d+) through (\\d+),(\\d+)$")

func makeLights(xSize, ySize uint) [][]bool {
	lights := make([][]bool, ySize)
	allLights := make([]bool, xSize*ySize)
	for i := range lights {
		lights[i], allLights = allLights[:xSize], allLights[xSize:]
	}

	return lights
}

func parseLine(s string) (instruction, coordinates, error) {
	var err error

	m := lineRegexp.FindStringSubmatch(s)
	if m == nil {
		err = fmt.Errorf("invalid instruction %q", s)
		return 0, coordinates{}, err
	}

	var inst instruction
	switch m[1] {
	case "turn on":
		inst = turnOn
	case "turn off":
		inst = turnOff
	case "toggle":
		inst = toggle
	}

	var coords coordinates
	coords.xStart, err = strconv.Atoi(m[2])
	if err != nil {
		return 0, coordinates{}, err
	}
	coords.yStart, err = strconv.Atoi(m[3])
	if err != nil {
		return 0, coordinates{}, err
	}
	coords.xEnd, err = strconv.Atoi(m[4])
	if err != nil {
		return 0, coordinates{}, err
	}
	coords.yEnd, err = strconv.Atoi(m[5])
	if err != nil {
		return 0, coordinates{}, err
	}

	if coords.xStart > coords.xEnd || coords.yStart > coords.yEnd {
		err := fmt.Errorf("invalid coordinates %d", coords)
		return 0, coordinates{}, err
	}

	return inst, coords, nil
}

func processLine(l [][]bool, s string) error {
	inst, coords, err := parseLine(s)
	if err != nil {
		return err
	}

	if coords.yStart > len(l) || coords.yEnd > len(l) ||
		coords.xStart > len(l[0]) || coords.xEnd > len(l[0]) {
		err := fmt.Errorf("invalid coordinates %d", coords)
		return err
	}

	for y := coords.yStart; y <= coords.yEnd; y++ {
		for x := coords.xStart; x <= coords.xEnd; x++ {
			switch inst {
			case turnOn:
				l[x][y] = true
			case turnOff:
				l[x][y] = false
			case toggle:
				l[x][y] = !l[x][y]
			}
		}
	}

	return nil
}

func run() int {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "%s filename\n", os.Args[0])
		return 1
	}

	f, err := os.Open(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return 1
	}
	defer f.Close()

	lights := makeLights(1000, 1000)
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		s := scanner.Text()

		err := processLine(lights, s)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return 1
		}
	}

	var lightCount int
	for _, row := range lights {
		for _, l := range row {
			if l {
				lightCount++
			}
		}
	}

	fmt.Printf("lights: %d\n", lightCount)

	return 0
}

func main() {
	os.Exit(run())
}
<commit_msg>day6 part 2 complete<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strconv"
)

//go:generate stringer -type=instruction
type instruction int

const (
	turnOn instruction = iota
	turnOff
	toggle
)

type coordinates struct {
	xStart int
	yStart int
	xEnd   int
	yEnd   int
}

var lineRegexp = regexp.MustCompile("^(turn on|turn off|toggle) (\\d+),(\\d+) through (\\d+),(\\d+)$")

func makeLights(xSize, ySize uint) [][]bool {
	lights := make([][]bool, ySize)
	allLights := make([]bool, xSize*ySize)
	for i := range lights {
		lights[i], allLights = allLights[:xSize], allLights[xSize:]
	}

	return lights
}

func makeLights2(xSize, ySize uint) [][]int {
	lights := make([][]int, ySize)
	allLights := make([]int, xSize*ySize)
	for i := range lights {
		lights[i], allLights = allLights[:xSize], allLights[xSize:]
	}

	return lights
}

func parseLine(s string) (instruction, coordinates, error) {
	var err error

	m := lineRegexp.FindStringSubmatch(s)
	if m == nil {
		err = fmt.Errorf("invalid instruction %q", s)
		return 0, coordinates{}, err
	}

	var inst instruction
	switch m[1] {
	case "turn on":
		inst = turnOn
	case "turn off":
		inst = turnOff
	case "toggle":
		inst = toggle
	}

	var coords coordinates
	coords.xStart, err = strconv.Atoi(m[2])
	if err != nil {
		return 0, coordinates{}, err
	}
	coords.yStart, err = strconv.Atoi(m[3])
	if err != nil {
		return 0, coordinates{}, err
	}
	coords.xEnd, err = strconv.Atoi(m[4])
	if err != nil {
		return 0, coordinates{}, err
	}
	coords.yEnd, err = strconv.Atoi(m[5])
	if err != nil {
		return 0, coordinates{}, err
	}

	if coords.xStart > coords.xEnd || coords.yStart > coords.yEnd {
		err := fmt.Errorf("invalid coordinates %d", coords)
		return 0, coordinates{}, err
	}

	return inst, coords, nil
}

func processLine(l [][]bool, s string) error {
	inst, coords, err := parseLine(s)
	if err != nil {
		return err
	}

	if coords.yStart > len(l) || coords.yEnd > len(l) ||
		coords.xStart > len(l[0]) || coords.xEnd > len(l[0]) {
		err := fmt.Errorf("invalid coordinates %d", coords)
		return err
	}

	for y := coords.yStart; y <= coords.yEnd; y++ {
		for x := coords.xStart; x <= coords.xEnd; x++ {
			switch inst {
			case turnOn:
				l[x][y] = true
			case turnOff:
				l[x][y] = false
			case toggle:
				l[x][y] = !l[x][y]
			}
		}
	}

	return nil
}

func processLine2(l [][]int, s string) error {
	inst, coords, err := parseLine(s)
	if err != nil {
		return err
	}

	if coords.yStart > len(l) || coords.yEnd > len(l) ||
		coords.xStart > len(l[0]) || coords.xEnd > len(l[0]) {
		err := fmt.Errorf("invalid coordinates %d", coords)
		return err
	}

	for y := coords.yStart; y <= coords.yEnd; y++ {
		for x := coords.xStart; x <= coords.xEnd; x++ {
			switch inst {
			case turnOn:
				l[x][y] += 1
			case turnOff:
				if l[x][y] > 0 {
					l[x][y] -= 1
				}
			case toggle:
				l[x][y] += 2
			}
		}
	}

	return nil
}

func run() int {
	var err error
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "%s filename\n", os.Args[0])
		return 1
	}

	f, err := os.Open(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return 1
	}
	defer f.Close()

	lights := makeLights(1000, 1000)
	lights2 := makeLights2(1000, 1000)
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		s := scanner.Text()

		err = processLine(lights, s)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return 1
		}

		err = processLine2(lights2, s)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return 1
		}
	}

	var lightCount int
	for _, row := range lights {
		for _, l := range row {
			if l {
				lightCount++
			}
		}
	}

	var brightness int
	for _, row := range lights2 {
		for _, l := range row {
			brightness += l
		}
	}

	fmt.Printf("lights: %d\n", lightCount)
	fmt.Printf("brightness: %d\n", brightness)

	return 0
}

func main() {
	os.Exit(run())
}
<|endoftext|>


<commit_before>package db

import (
	"context"
	"database/sql/driver"
	"io"
	"net"
	"net/url"
	"sync"

	"golang.org/x/crypto/ssh"
	"golang.org/x/net/proxy"
)

type connectorFunc func(context.Context) (driver.Conn, error)

func (f connectorFunc) Connect(ctx context.Context) (driver.Conn, error) {
	return f(ctx)
}

func (f connectorFunc) Driver() driver.Driver {
	panic("not supported")
}

func NewSSHProxyConnector(jumpboxAddr string, config *ssh.ClientConfig, sqlDriver driver.Driver, uri string) (driver.Connector, error) {
	if jumpboxAddr == "99.99.99.99:22" {
		// Under testing if this branch is taken, skip tunneling
		return connectorFunc(func(_ context.Context) (driver.Conn, error) {
			return sqlDriver.Open(uri)
		}), nil
	}
	u, err := url.Parse(uri)
	if err != nil {
		return nil, err
	}
	remoteHost := u.Host
	client, err := ssh.Dial("tcp", jumpboxAddr, config)
	if err != nil {
		return nil, err
	}
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				return // BUG: expose this error somehow, detect Temporary errors.
			}
			go proxyConn(conn, client, remoteHost)
		}
	}()
	u.Host = l.Addr().String()
	newURI := u.String()

	return connectorFunc(func(_ context.Context) (driver.Conn, error) {
		return sqlDriver.Open(newURI)
	}), nil
}

func proxyConn(conn net.Conn, dialer proxy.Dialer, addr string) {
	defer conn.Close()
	target, err := dialer.Dial("tcp", addr)
	if err != nil {
		return //BUG: expose this error somehow
	}
	defer target.Close()
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		io.Copy(conn, target)
		wg.Done()
	}()
	go func() {
		io.Copy(target, conn)
		wg.Done()
	}()
	wg.Wait()
}
<commit_msg>Document NewSSHProxyConnector<commit_after>package db

import (
	"context"
	"database/sql/driver"
	"io"
	"net"
	"net/url"
	"sync"

	"golang.org/x/crypto/ssh"
	"golang.org/x/net/proxy"
)

type connectorFunc func(context.Context) (driver.Conn, error)

func (f connectorFunc) Connect(ctx context.Context) (driver.Conn, error) {
	return f(ctx)
}

func (f connectorFunc) Driver() driver.Driver {
	panic("not supported")
}

// NewSSHProxyConnector is used to proxy a database connection through a ssh jumpbox.
// uri is the arguement that will be passed to sqlDriver.Open.
// uri must be a valid URI.
// The return value can be used by sql.OpenDB to obtain a *sql.DB
func NewSSHProxyConnector(jumpboxAddr string, config *ssh.ClientConfig, sqlDriver driver.Driver, uri string) (driver.Connector, error) {
	if jumpboxAddr == "99.99.99.99:22" {
		// Under testing if this branch is taken, skip tunneling
		return connectorFunc(func(_ context.Context) (driver.Conn, error) {
			return sqlDriver.Open(uri)
		}), nil
	}
	u, err := url.Parse(uri)
	if err != nil {
		return nil, err
	}
	remoteHost := u.Host
	client, err := ssh.Dial("tcp", jumpboxAddr, config)
	if err != nil {
		return nil, err
	}
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				return // BUG: expose this error somehow, detect Temporary errors.
			}
			go proxyConn(conn, client, remoteHost)
		}
	}()
	u.Host = l.Addr().String()
	newURI := u.String()

	return connectorFunc(func(_ context.Context) (driver.Conn, error) {
		return sqlDriver.Open(newURI)
	}), nil
}

func proxyConn(conn net.Conn, dialer proxy.Dialer, addr string) {
	defer conn.Close()
	target, err := dialer.Dial("tcp", addr)
	if err != nil {
		return //BUG: expose this error somehow
	}
	defer target.Close()
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		io.Copy(conn, target)
		wg.Done()
	}()
	go func() {
		io.Copy(target, conn)
		wg.Done()
	}()
	wg.Wait()
}
<|endoftext|>


<commit_before>package dbf

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
)

type StructCrack struct {
	Id            string `json:"id"`
	Type          string `json:"type"`
	Generator     string `json:"generator"`
	Cracker       string `json:"cracker"`
	Cmd_generator string `json:"cmd_generator"`
	Cmd_cracker   string `json:"cmd_cracker"`
	Target        string `json:"target"`
	Has_dep       bool   `json:"has_dep"`
}

func processCrack(task *StructCrackTask, crackInfoPath *string) bool {
	var vendorPath, cmdJsonStr string
	var cmdArg []string
	var resultStatus int

	resultStatus = -1
	taskPath := getPath(_PATH_TASK) + task.Platform + PATH_SEPARATOR

	defer func() {
		resultByte, err := ioutil.ReadFile(taskPath + "result")
		if err != nil {
			Log.Printf("%s\n", err)
			resultByte = nil
			resultStatus = -2
		}

		fmt.Printf("Sending result of crack #%s...\n", task.Crack_id)

		if sendResult(`[{"crack_id":"`+task.Crack_id+`","start":"`+task.Start+`","offset":"`+task.Offset+`","result":"`+string(resultByte)+`","status":"`+strconv.Itoa(resultStatus)+`"}]`) == true {
			fmt.Printf("Removing task info of crack #%s (%s)...\n", task.Crack_id, task.Platform)
			err = os.RemoveAll(taskPath)
			if err != nil {
				Log.Printf("%s\n", err)
			}
		}
	}()

	crackJson, err := ioutil.ReadFile(*crackInfoPath)
	if err != nil {
		Log.Printf("%s\n", err)
		resultStatus = -3
		return false
	}

	var crack StructCrack
	err = json.Unmarshal(crackJson, &crack)
	if err != nil {
		Log.Printf("%s\n", err)
		resultStatus = -4
		return false
	}

	/* Process crack */
	// Check cracker
	if crack.Cracker != "" {
		vendorPath = getPath(_PATH_VENDOR) + _VENDOR_TYPE_CRACKER + PATH_SEPARATOR + crack.Cracker + PATH_SEPARATOR + task.Platform + PATH_SEPARATOR + _VENDOR_TYPE_CRACKER + extExecutable
		if checkVendor(_VENDOR_TYPE_CRACKER, &crack.Cracker, &task.Platform, &vendorPath) == false {
			resultStatus = -5
			return false
		}
	}

	// Check hashfile
	*crackInfoPath = filepath.Dir(*crackInfoPath) + PATH_SEPARATOR + "hashfile"
	if _, err := os.Stat(*crackInfoPath); err != nil {
		if os.IsNotExist(err) { // Does not exist, so create it
			err = ioutil.WriteFile(*crackInfoPath, []byte(crack.Target), 0664)
			if err != nil {
				Log.Printf("%s\n", err) // Error in creating
				resultStatus = -6
				return false
			}
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			resultStatus = -7
			return false
		}
	}

	generatorReplacer := strings.NewReplacer("START", task.Start, "OFFSET", task.Offset, `"IN_FILE"`, strconv.Quote(taskPath+"file.fifo"))
	crackerReplacer := strings.NewReplacer(`"HASH_FILE"`, strconv.Quote(*crackInfoPath), `"OUT_FILE"`, strconv.Quote(taskPath+"result"), `"IN_FILE"`, strconv.Quote(taskPath+"file.fifo"))

	if crack.Type == "embed" { // Embeded
		cmdJsonStr = generatorReplacer.Replace(crack.Cmd_cracker)
		cmdJsonStr = crackerReplacer.Replace(cmdJsonStr)
		err = json.Unmarshal([]byte(cmdJsonStr), &cmdArg)
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -8
			return false
		}

		fmt.Printf("Performing crack #%s...\n", task.Crack_id)

		err = exec.Command(vendorPath, cmdArg...).Run()
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -9
			return false
		} else {
			resultStatus = 0
			return true
		}
	} else { // Not embeded
		// Prepare cracker
		cmdJsonStr = crackerReplacer.Replace(crack.Cmd_cracker)
		err = json.Unmarshal([]byte(cmdJsonStr), &cmdArg)
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -10
			return false
		}
		execCracker := exec.Command(vendorPath, cmdArg...)

		// Check generator
		vendorPath = getPath(_PATH_VENDOR) + _VENDOR_TYPE_GENERATOR + PATH_SEPARATOR + crack.Generator + PATH_SEPARATOR + task.Platform + PATH_SEPARATOR + _VENDOR_TYPE_GENERATOR + extExecutable
		if checkVendor(_VENDOR_TYPE_GENERATOR, &crack.Generator, &task.Platform, &vendorPath) == false {
			resultStatus = -11
			return false
		}

		// Prepare generator
		cmdJsonStr = generatorReplacer.Replace(crack.Cmd_generator)
		if strings.Contains(cmdJsonStr, "DEP_GEN") {
			// Check if dependency exists in crack location
			*crackInfoPath = filepath.Dir(*crackInfoPath) + PATH_SEPARATOR + "dep" + PATH_SEPARATOR + "dep-gen"
			if _, err := os.Stat(*crackInfoPath); err == nil { // dep-gen file exists in crack location and is accessible
				cmdJsonStr = strings.Replace(cmdJsonStr, `"DEP_GEN"`, strconv.Quote(*crackInfoPath), -1)
			} else { // Check if dependency exists in generator location
				vendorPath = filepath.Dir(vendorPath) + PATH_SEPARATOR + "dep-gen"
				if _, err := os.Stat(vendorPath); err == nil { // dep-gen file exists in generator location and is accessible
					cmdJsonStr = strings.Replace(cmdJsonStr, `"DEP_GEN"`, strconv.Quote(vendorPath), -1)
				} else {
					resultStatus = -12
					return false
				}
				vendorPath = filepath.Dir(vendorPath) + PATH_SEPARATOR + _VENDOR_TYPE_GENERATOR + extExecutable // Rename back to generator executable
			}
		}
		err = json.Unmarshal([]byte(cmdJsonStr), &cmdArg)
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -13
			return false
		}
		execGenerator := exec.Command(vendorPath, cmdArg...)

		fmt.Printf("Performing crack #%s...\n", task.Crack_id)

		if crack.Type == "infile" {
			err = exec.Command("mkfifo", taskPath+"file.fifo").Run()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -14
				return false
			}

			err = execGenerator.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -15
				return false
			}

			err = execCracker.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -16
				return false
			}

			errG := execGenerator.Wait()
			errC := execCracker.Wait()
			if (errG != nil) || (errC != nil) {

				resultStatus = -16

				if errG != nil {
					Log.Printf("%s\n", errG)
					resultStatus += -1
				} else if errC != nil {
					Log.Printf("%s\n", errC)
					resultStatus += -2
				}

				// Last resultStatus: -19

				return false
			} else {
				resultStatus = 0
				return true
			}
		} else { // Stdin
			r, err := execGenerator.StdoutPipe()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -20
				return false
			}
			execCracker.Stdin = r

			err = execGenerator.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -21
				return false
			}

			err = execCracker.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -22
				return false
			}

			errC := execCracker.Wait()
			errR := r.Close()
			execGenerator.Process.Kill()
			errG := execGenerator.Wait()
			if (errG != nil) || (errR != nil) || (errC != nil) {
				resultStatus = -22

				if errG != nil {
					Log.Printf("%s\n", errG)
					resultStatus += -1
				} else if errR != nil {
					Log.Printf("%s\n", errR)
					resultStatus += -2
				} else if errC != nil {
					Log.Printf("%s\n", errC)
					resultStatus += -4
				}

				// Last resultStatus: -29

				return false
			} else {
				resultStatus = 0
				return true
			}
		}
	}

	return true
}
<commit_msg>- Do not send error if result file is not found. - Do not check cracker exit code (It may send error if the crack was not found!)<commit_after>package dbf

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
)

type StructCrack struct {
	Id            string `json:"id"`
	Type          string `json:"type"`
	Generator     string `json:"generator"`
	Cracker       string `json:"cracker"`
	Cmd_generator string `json:"cmd_generator"`
	Cmd_cracker   string `json:"cmd_cracker"`
	Target        string `json:"target"`
	Has_dep       bool   `json:"has_dep"`
}

func processCrack(task *StructCrackTask, crackInfoPath *string) bool {
	var vendorPath, cmdJsonStr string
	var cmdArg []string
	var resultStatus int

	resultStatus = -1
	taskPath := getPath(_PATH_TASK) + task.Platform + PATH_SEPARATOR

	defer func() {
		resultByte, err := ioutil.ReadFile(taskPath + "result")
		if err != nil {
			resultByte = nil

			if os.IsNotExist(err) {
				resultStatus = 0 // Was not cracked
			} else {
				Log.Printf("%s\n", err)
				resultStatus = -2
			}
		}

		fmt.Printf("Sending result of crack #%s...\n", task.Crack_id)

		if sendResult(`[{"crack_id":"`+task.Crack_id+`","start":"`+task.Start+`","offset":"`+task.Offset+`","result":"`+string(resultByte)+`","status":"`+strconv.Itoa(resultStatus)+`"}]`) == true {
			fmt.Printf("Removing task info of crack #%s (%s)...\n", task.Crack_id, task.Platform)
			err = os.RemoveAll(taskPath)
			if err != nil {
				Log.Printf("%s\n", err)
			}
		}
	}()

	crackJson, err := ioutil.ReadFile(*crackInfoPath)
	if err != nil {
		Log.Printf("%s\n", err)
		resultStatus = -3
		return false
	}

	var crack StructCrack
	err = json.Unmarshal(crackJson, &crack)
	if err != nil {
		Log.Printf("%s\n", err)
		resultStatus = -4
		return false
	}

	/* Process crack */
	// Check cracker
	if crack.Cracker != "" {
		vendorPath = getPath(_PATH_VENDOR) + _VENDOR_TYPE_CRACKER + PATH_SEPARATOR + crack.Cracker + PATH_SEPARATOR + task.Platform + PATH_SEPARATOR + _VENDOR_TYPE_CRACKER + extExecutable
		if checkVendor(_VENDOR_TYPE_CRACKER, &crack.Cracker, &task.Platform, &vendorPath) == false {
			resultStatus = -5
			return false
		}
	}

	// Check hashfile
	*crackInfoPath = filepath.Dir(*crackInfoPath) + PATH_SEPARATOR + "hashfile"
	if _, err := os.Stat(*crackInfoPath); err != nil {
		if os.IsNotExist(err) { // Does not exist, so create it
			err = ioutil.WriteFile(*crackInfoPath, []byte(crack.Target), 0664)
			if err != nil {
				Log.Printf("%s\n", err) // Error in creating
				resultStatus = -6
				return false
			}
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			resultStatus = -7
			return false
		}
	}

	generatorReplacer := strings.NewReplacer("START", task.Start, "OFFSET", task.Offset, `"IN_FILE"`, strconv.Quote(taskPath+"file.fifo"))
	crackerReplacer := strings.NewReplacer(`"HASH_FILE"`, strconv.Quote(*crackInfoPath), `"OUT_FILE"`, strconv.Quote(taskPath+"result"), `"IN_FILE"`, strconv.Quote(taskPath+"file.fifo"))

	if crack.Type == "embed" { // Embeded
		cmdJsonStr = generatorReplacer.Replace(crack.Cmd_cracker)
		cmdJsonStr = crackerReplacer.Replace(cmdJsonStr)
		err = json.Unmarshal([]byte(cmdJsonStr), &cmdArg)
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -8
			return false
		}

		fmt.Printf("Performing crack #%s...\n", task.Crack_id)

		err = exec.Command(vendorPath, cmdArg...).Run()
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -9
			return false
		} else {
			resultStatus = 0
			return true
		}
	} else { // Not embeded
		// Prepare cracker
		cmdJsonStr = crackerReplacer.Replace(crack.Cmd_cracker)
		err = json.Unmarshal([]byte(cmdJsonStr), &cmdArg)
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -10
			return false
		}
		execCracker := exec.Command(vendorPath, cmdArg...)

		// Check generator
		vendorPath = getPath(_PATH_VENDOR) + _VENDOR_TYPE_GENERATOR + PATH_SEPARATOR + crack.Generator + PATH_SEPARATOR + task.Platform + PATH_SEPARATOR + _VENDOR_TYPE_GENERATOR + extExecutable
		if checkVendor(_VENDOR_TYPE_GENERATOR, &crack.Generator, &task.Platform, &vendorPath) == false {
			resultStatus = -11
			return false
		}

		// Prepare generator
		cmdJsonStr = generatorReplacer.Replace(crack.Cmd_generator)
		if strings.Contains(cmdJsonStr, "DEP_GEN") {
			// Check if dependency exists in crack location
			*crackInfoPath = filepath.Dir(*crackInfoPath) + PATH_SEPARATOR + "dep" + PATH_SEPARATOR + "dep-gen"
			if _, err := os.Stat(*crackInfoPath); err == nil { // dep-gen file exists in crack location and is accessible
				cmdJsonStr = strings.Replace(cmdJsonStr, `"DEP_GEN"`, strconv.Quote(*crackInfoPath), -1)
			} else { // Check if dependency exists in generator location
				vendorPath = filepath.Dir(vendorPath) + PATH_SEPARATOR + "dep-gen"
				if _, err := os.Stat(vendorPath); err == nil { // dep-gen file exists in generator location and is accessible
					cmdJsonStr = strings.Replace(cmdJsonStr, `"DEP_GEN"`, strconv.Quote(vendorPath), -1)
				} else {
					resultStatus = -12
					return false
				}
				vendorPath = filepath.Dir(vendorPath) + PATH_SEPARATOR + _VENDOR_TYPE_GENERATOR + extExecutable // Rename back to generator executable
			}
		}
		err = json.Unmarshal([]byte(cmdJsonStr), &cmdArg)
		if err != nil {
			Log.Printf("%s\n", err)
			resultStatus = -13
			return false
		}
		execGenerator := exec.Command(vendorPath, cmdArg...)

		fmt.Printf("Performing crack #%s...\n", task.Crack_id)

		if crack.Type == "infile" {
			err = exec.Command("mkfifo", taskPath+"file.fifo").Run()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -14
				return false
			}

			err = execGenerator.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -15
				return false
			}

			err = execCracker.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -16
				return false
			}

			errG := execGenerator.Wait()
			errC := execCracker.Wait()
			if (errG != nil) || (errC != nil) {

				resultStatus = -16

				if errG != nil {
					Log.Printf("%s\n", errG)
					resultStatus += -1
				} else if errC != nil {
					Log.Printf("%s\n", errC)
					resultStatus += -2
				}

				// Last resultStatus: -19

				return false
			} else {
				resultStatus = 0
				return true
			}
		} else { // Stdin
			r, err := execGenerator.StdoutPipe()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -20
				return false
			}
			execCracker.Stdin = r

			err = execGenerator.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -21
				return false
			}

			err = execCracker.Start()
			if err != nil {
				Log.Printf("%s\n", err)
				resultStatus = -22
				return false
			}

			execCracker.Wait()
			errR := r.Close()
			errG := execGenerator.Process.Release()
			//			errG := execGenerator.Wait()
			if (errR != nil) || (errG != nil) {
				resultStatus = -22

				if errG != nil {
					Log.Printf("%s\n", errG)
					resultStatus += -1
				} else if errR != nil {
					Log.Printf("%s\n", errR)
					resultStatus += -2
				}

				// Last resultStatus: -25

				return false
			} else {
				resultStatus = 0
				return true
			}
		}
	}

	return true
}
<|endoftext|>


<commit_before>package dbng

import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"

	"github.com/Masterminds/squirrel"
	"github.com/concourse/atc/db/migrations"
	"github.com/concourse/atc/dbng/migration"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/lib/pq"
)

type Conn interface {
	Bus() NotificationsBus
	EncryptionStrategy() EncryptionStrategy

	Begin() (Tx, error)
	Driver() driver.Driver
	Exec(query string, args ...interface{}) (sql.Result, error)
	Ping() error
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner
	SetMaxIdleConns(n int)
	SetMaxOpenConns(n int)
	Stats() sql.DBStats

	Close() error
}

type Tx interface {
	Commit() error
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner
	Rollback() error
	Stmt(stmt *sql.Stmt) *sql.Stmt
}

func Open(logger lager.Logger, sqlDriver string, sqlDataSource string, newKey *EncryptionKey, oldKey *EncryptionKey) (Conn, error) {
	for {
		sqlDb, err := migration.Open(sqlDriver, sqlDataSource, migrations.Migrations)
		if err != nil {
			if strings.Contains(err.Error(), " dial ") {
				logger.Error("failed-to-open-db-retrying", err)
				time.Sleep(5 * time.Second)
				continue
			}

			return nil, err
		}

		var strategy EncryptionStrategy
		if newKey != nil {
			strategy = newKey
		} else {
			strategy = NewNoEncryption()
		}

		switch {
		case oldKey != nil && newKey == nil:
			err = decryptToPlaintext(logger.Session("decrypt"), sqlDb, oldKey)
		case oldKey != nil && newKey != nil:
			err = encryptWithNewKey(logger.Session("rotate"), sqlDb, newKey, oldKey)
		}
		if err != nil {
			return nil, err
		}

		if newKey != nil {
			err = encryptPlaintext(logger.Session("encrypt"), sqlDb, newKey)
			if err != nil {
				return nil, err
			}
		}

		listener := pq.NewListener(sqlDataSource, time.Second, time.Minute, nil)

		return &db{
			DB: sqlDb,

			bus:        NewNotificationsBus(listener, sqlDb),
			encryption: strategy,
		}, nil
	}
}

var encryptedColumns = map[string]string{
	"teams":          "auth",
	"resources":      "config",
	"jobs":           "config",
	"resource_types": "config",
	"pipelines":      "config",
}

func encryptPlaintext(logger lager.Logger, sqlDB *sql.DB, key *EncryptionKey) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				id  int
				val sql.NullString
			)

			err := rows.Scan(&id, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			if !val.Valid {
				continue
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			encrypted, nonce, err := key.Encrypt([]byte(val.String))
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = $2
				WHERE id = $3
			`, encrypted, nonce, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			encryptedRows++
		}

		if encryptedRows > 0 {
			tLog.Info("encrypted-existing-plaintext-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

func decryptToPlaintext(logger lager.Logger, sqlDB *sql.DB, oldKey *EncryptionKey) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, nonce, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		decryptedRows := 0

		for rows.Next() {
			var (
				id         int
				val, nonce string
			)

			err := rows.Scan(&id, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				rLog.Error("failed-to-decrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = NULL
				WHERE id = $2
			`, decrypted, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			decryptedRows++
		}

		if decryptedRows > 0 {
			tLog.Info("decrypted-existing-encrypted-data", lager.Data{
				"rows": decryptedRows,
			})
		}
	}

	return nil
}

var ErrEncryptedWithUnknownKey = errors.New("row encrypted with neither old nor new key")

func encryptWithNewKey(logger lager.Logger, sqlDB *sql.DB, newKey *EncryptionKey, oldKey *EncryptionKey) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, nonce, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				id         int
				val, nonce string
			)

			err := rows.Scan(&id, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				_, err = newKey.Decrypt(val, &nonce)
				if err == nil {
					rLog.Debug("already-encrypted-with-new-key")
					continue
				}

				logger.Error("failed-to-decrypt-with-either-key", err)
				return ErrEncryptedWithUnknownKey
			}

			encrypted, newNonce, err := newKey.Encrypt(decrypted)
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = $2
				WHERE id = $3
			`, encrypted, newNonce, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}
		}

		if encryptedRows > 0 {
			tLog.Info("re-encrypted-existing-encrypted-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

type db struct {
	*sql.DB

	bus        NotificationsBus
	encryption EncryptionStrategy
}

func (db *db) Bus() NotificationsBus {
	return db.bus
}

func (db *db) EncryptionStrategy() EncryptionStrategy {
	return db.encryption
}

func (db *db) Close() error {
	var errs error
	dbErr := db.DB.Close()
	if dbErr != nil {
		errs = multierror.Append(errs, dbErr)
	}

	busErr := db.bus.Close()
	if busErr != nil {
		errs = multierror.Append(errs, busErr)
	}

	return errs
}

func (db *db) Begin() (Tx, error) {
	tx, err := db.DB.Begin()
	if err != nil {
		return nil, err
	}

	return &dbTx{tx}, nil
}

// to conform to squirrel.Runner interface
func (db *db) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	return db.DB.QueryRow(query, args...)
}

type dbTx struct {
	*sql.Tx
}

// to conform to squirrel.Runner interface
func (tx *dbTx) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	return tx.Tx.QueryRow(query, args...)
}

type nonOneRowAffectedError struct {
	RowsAffected int64
}

func (err nonOneRowAffectedError) Error() string {
	return fmt.Sprintf("expected 1 row to be updated; got %d", err.RowsAffected)
}
<commit_msg>fix log for re-encrypting<commit_after>package dbng

import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"

	"github.com/Masterminds/squirrel"
	"github.com/concourse/atc/db/migrations"
	"github.com/concourse/atc/dbng/migration"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/lib/pq"
)

type Conn interface {
	Bus() NotificationsBus
	EncryptionStrategy() EncryptionStrategy

	Begin() (Tx, error)
	Driver() driver.Driver
	Exec(query string, args ...interface{}) (sql.Result, error)
	Ping() error
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner
	SetMaxIdleConns(n int)
	SetMaxOpenConns(n int)
	Stats() sql.DBStats

	Close() error
}

type Tx interface {
	Commit() error
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner
	Rollback() error
	Stmt(stmt *sql.Stmt) *sql.Stmt
}

func Open(logger lager.Logger, sqlDriver string, sqlDataSource string, newKey *EncryptionKey, oldKey *EncryptionKey) (Conn, error) {
	for {
		sqlDb, err := migration.Open(sqlDriver, sqlDataSource, migrations.Migrations)
		if err != nil {
			if strings.Contains(err.Error(), " dial ") {
				logger.Error("failed-to-open-db-retrying", err)
				time.Sleep(5 * time.Second)
				continue
			}

			return nil, err
		}

		var strategy EncryptionStrategy
		if newKey != nil {
			strategy = newKey
		} else {
			strategy = NewNoEncryption()
		}

		switch {
		case oldKey != nil && newKey == nil:
			err = decryptToPlaintext(logger.Session("decrypt"), sqlDb, oldKey)
		case oldKey != nil && newKey != nil:
			err = encryptWithNewKey(logger.Session("rotate"), sqlDb, newKey, oldKey)
		}
		if err != nil {
			return nil, err
		}

		if newKey != nil {
			err = encryptPlaintext(logger.Session("encrypt"), sqlDb, newKey)
			if err != nil {
				return nil, err
			}
		}

		listener := pq.NewListener(sqlDataSource, time.Second, time.Minute, nil)

		return &db{
			DB: sqlDb,

			bus:        NewNotificationsBus(listener, sqlDb),
			encryption: strategy,
		}, nil
	}
}

var encryptedColumns = map[string]string{
	"teams":          "auth",
	"resources":      "config",
	"jobs":           "config",
	"resource_types": "config",
	"pipelines":      "config",
}

func encryptPlaintext(logger lager.Logger, sqlDB *sql.DB, key *EncryptionKey) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				id  int
				val sql.NullString
			)

			err := rows.Scan(&id, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			if !val.Valid {
				continue
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			encrypted, nonce, err := key.Encrypt([]byte(val.String))
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = $2
				WHERE id = $3
			`, encrypted, nonce, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			encryptedRows++
		}

		if encryptedRows > 0 {
			tLog.Info("encrypted-existing-plaintext-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

func decryptToPlaintext(logger lager.Logger, sqlDB *sql.DB, oldKey *EncryptionKey) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, nonce, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		decryptedRows := 0

		for rows.Next() {
			var (
				id         int
				val, nonce string
			)

			err := rows.Scan(&id, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				rLog.Error("failed-to-decrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = NULL
				WHERE id = $2
			`, decrypted, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			decryptedRows++
		}

		if decryptedRows > 0 {
			tLog.Info("decrypted-existing-encrypted-data", lager.Data{
				"rows": decryptedRows,
			})
		}
	}

	return nil
}

var ErrEncryptedWithUnknownKey = errors.New("row encrypted with neither old nor new key")

func encryptWithNewKey(logger lager.Logger, sqlDB *sql.DB, newKey *EncryptionKey, oldKey *EncryptionKey) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, nonce, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				id         int
				val, nonce string
			)

			err := rows.Scan(&id, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				_, err = newKey.Decrypt(val, &nonce)
				if err == nil {
					rLog.Debug("already-encrypted-with-new-key")
					continue
				}

				logger.Error("failed-to-decrypt-with-either-key", err)
				return ErrEncryptedWithUnknownKey
			}

			encrypted, newNonce, err := newKey.Encrypt(decrypted)
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = $2
				WHERE id = $3
			`, encrypted, newNonce, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			encryptedRows++
		}

		if encryptedRows > 0 {
			tLog.Info("re-encrypted-existing-encrypted-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

type db struct {
	*sql.DB

	bus        NotificationsBus
	encryption EncryptionStrategy
}

func (db *db) Bus() NotificationsBus {
	return db.bus
}

func (db *db) EncryptionStrategy() EncryptionStrategy {
	return db.encryption
}

func (db *db) Close() error {
	var errs error
	dbErr := db.DB.Close()
	if dbErr != nil {
		errs = multierror.Append(errs, dbErr)
	}

	busErr := db.bus.Close()
	if busErr != nil {
		errs = multierror.Append(errs, busErr)
	}

	return errs
}

func (db *db) Begin() (Tx, error) {
	tx, err := db.DB.Begin()
	if err != nil {
		return nil, err
	}

	return &dbTx{tx}, nil
}

// to conform to squirrel.Runner interface
func (db *db) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	return db.DB.QueryRow(query, args...)
}

type dbTx struct {
	*sql.Tx
}

// to conform to squirrel.Runner interface
func (tx *dbTx) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	return tx.Tx.QueryRow(query, args...)
}

type nonOneRowAffectedError struct {
	RowsAffected int64
}

func (err nonOneRowAffectedError) Error() string {
	return fmt.Sprintf("expected 1 row to be updated; got %d", err.RowsAffected)
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 The btcsuite developers
// Copyright (c) 2015-2017 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime"
	"runtime/pprof"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrwallet/chain"
	"github.com/decred/dcrwallet/internal/prompt"
	"github.com/decred/dcrwallet/internal/zero"
	ldr "github.com/decred/dcrwallet/loader"
	"github.com/decred/dcrwallet/rpc/legacyrpc"
	"github.com/decred/dcrwallet/rpc/rpcserver"
	"github.com/decred/dcrwallet/version"
	"github.com/decred/dcrwallet/wallet"
)

var (
	cfg *config
)

func main() {
	// Create a context that is cancelled when a shutdown request is received
	// through an interrupt signal or an RPC request.
	ctx := withShutdownCancel(context.Background())
	go shutdownListener()

	// Run the wallet until permanent failure or shutdown is requested.
	if err := run(ctx); err != nil && err != context.Canceled {
		os.Exit(1)
	}
}

// done returns whether the context's Done channel was closed due to
// cancellation or exceeded deadline.
func done(ctx context.Context) bool {
	select {
	case <-ctx.Done():
		return true
	default:
		return false
	}
}

// run is the main startup and teardown logic performed by the main package.  It
// is responsible for parsing the config, starting RPC servers, loading and
// syncing the wallet (if necessary), and stopping all started services when the
// context is cancelled.
func run(ctx context.Context) error {
	// Load configuration and parse command line.  This function also
	// initializes logging and configures it accordingly.
	tcfg, _, err := loadConfig(ctx)
	if err != nil {
		return err
	}
	cfg = tcfg
	defer func() {
		if logRotator != nil {
			logRotator.Close()
		}
	}()

	// Show version at startup.
	log.Infof("Version %s (Go version %s)", version.String(), runtime.Version())

	// Read IPC messages from the read end of a pipe created and passed by the
	// parent process, if any.  When this pipe is closed, shutdown is
	// initialized.
	if cfg.PipeRx != nil {
		go serviceControlPipeRx(uintptr(*cfg.PipeRx))
	}
	if cfg.PipeTx != nil {
		go serviceControlPipeTx(uintptr(*cfg.PipeTx))
	} else {
		go drainOutgoingPipeMessages()
	}

	// Run the pprof profiler if enabled.
	if len(cfg.Profile) > 0 {
		if done(ctx) {
			return ctx.Err()
		}

		profileRedirect := http.RedirectHandler("/debug/pprof", http.StatusSeeOther)
		http.Handle("/", profileRedirect)
		for _, listenAddr := range cfg.Profile {
			listenAddr := listenAddr // copy for closure
			go func() {
				log.Infof("Starting profile server on %s", listenAddr)
				err := http.ListenAndServe(listenAddr, nil)
				if err != nil {
					fatalf("Unable to run profiler: %v", err)
				}
			}()
		}
	}

	// Write mem profile if requested.
	if cfg.MemProfile != "" {
		if done(ctx) {
			return ctx.Err()
		}

		f, err := os.Create(cfg.MemProfile)
		if err != nil {
			log.Errorf("Unable to create cpu profile: %v", err)
			return err
		}
		timer := time.NewTimer(time.Minute * 5) // 5 minutes
		go func() {
			<-timer.C
			pprof.WriteHeapProfile(f)
			f.Close()
		}()
	}

	if done(ctx) {
		return ctx.Err()
	}

	// Create the loader which is used to load and unload the wallet.  If
	// --noinitialload is not set, this function is responsible for loading the
	// wallet.  Otherwise, loading is deferred so it can be performed over RPC.
	dbDir := networkDir(cfg.AppDataDir.Value, activeNet.Params)
	stakeOptions := &ldr.StakeOptions{
		VotingEnabled:       cfg.EnableVoting,
		AddressReuse:        cfg.ReuseAddresses,
		VotingAddress:       cfg.TBOpts.VotingAddress.Address,
		PoolAddress:         cfg.PoolAddress.Address,
		PoolFees:            cfg.PoolFees,
		StakePoolColdExtKey: cfg.StakePoolColdExtKey,
		TicketFee:           cfg.TicketFee.ToCoin(),
	}
	loader := ldr.NewLoader(activeNet.Params, dbDir, stakeOptions,
		cfg.AddrIdxScanLen, cfg.AllowHighFees, cfg.RelayFee.ToCoin())

	// Stop any services started by the loader after the shutdown procedure is
	// initialized and this function returns.
	defer func() {
		err := loader.UnloadWallet()
		if err != nil && err != ldr.ErrWalletNotLoaded {
			log.Errorf("Failed to close wallet: %v", err)
		} else if err == nil {
			log.Infof("Closed wallet")
		}
	}()

	// Open the wallet when --noinitialload was not set.
	passphrase := []byte{}
	if !cfg.NoInitialLoad {
		walletPass := []byte(cfg.WalletPass)
		defer zero.Bytes(walletPass)

		if cfg.PromptPublicPass {
			walletPass, _ = passPrompt(ctx, "Enter public wallet passphrase", false)
		}

		if done(ctx) {
			return ctx.Err()
		}

		// Load the wallet.  It must have been created already or this will
		// return an appropriate error.
		w, err := loader.OpenExistingWallet(walletPass)
		if err != nil {
			log.Errorf("Failed to open wallet: %v", err)
			return err
		}

		if done(ctx) {
			return ctx.Err()
		}

		// TODO(jrick): I think that this prompt should be removed
		// entirely instead of enabling it when --noinitialload is
		// unset.  It can be replaced with an RPC request (either
		// providing the private passphrase as a parameter, or require
		// unlocking the wallet first) to trigger a full accounts
		// rescan.
		//
		// Until then, since --noinitialload users are expecting to use
		// the wallet only over RPC, disable this feature for them.
		if cfg.Pass != "" {
			passphrase = []byte(cfg.Pass)
			err = w.Unlock(passphrase, nil)
			if err != nil {
				log.Errorf("Incorrect passphrase in pass config setting.")
				return err
			}
			w.SetInitiallyUnlocked(true)
		} else {
			passphrase = startPromptPass(ctx, w)
		}
	}

	if done(ctx) {
		return ctx.Err()
	}

	// Create and start the RPC servers to serve wallet client connections.  If
	// any of the servers can not be started, it will be nil.  If none of them
	// can be started, this errors since at least one server must run for the
	// wallet to be useful.
	//
	// Servers will be associated with a loaded wallet if it has already been
	// loaded, or after it is loaded later on.
	gRPCServer, jsonRPCServer, err := startRPCServers(loader)
	if err != nil {
		log.Errorf("Unable to create RPC servers: %v", err)
		return err
	}
	if gRPCServer != nil {
		// Start wallet and voting gRPC services after a wallet is loaded.
		loader.RunAfterLoad(func(w *wallet.Wallet) {
			rpcserver.StartWalletService(gRPCServer, w)
			rpcserver.StartVotingService(gRPCServer, w)
		})
		defer func() {
			log.Warn("Stopping gRPC server...")
			gRPCServer.Stop()
			log.Info("gRPC server shutdown")
		}()
	}
	if jsonRPCServer != nil {
		go func() {
			for range jsonRPCServer.RequestProcessShutdown() {
				requestShutdown()
			}
		}()
		defer func() {
			log.Warn("Stopping JSON-RPC server...")
			jsonRPCServer.Stop()
			log.Info("JSON-RPC server shutdown")
		}()
	}

	// Stop the ticket buyer (if running) on shutdown.  This returns an error
	// that can be ignored when the ticket buyer was never started.
	defer loader.StopTicketPurchase()

	// When not running with --noinitialload, it is the main package's
	// responsibility to connect the loaded wallet to the dcrd RPC server for
	// wallet synchronization.  This function blocks until cancelled.
	if !cfg.NoInitialLoad {
		if done(ctx) {
			return ctx.Err()
		}

		rpcClientConnectLoop(ctx, passphrase, jsonRPCServer, loader)
	}

	// Wait until shutdown is signaled before returning and running deferred
	// shutdown tasks.
	<-ctx.Done()
	return ctx.Err()
}

func passPrompt(ctx context.Context, prefix string, confirm bool) (passphrase []byte, err error) {
	os.Stdout.Sync()
	c := make(chan struct{}, 1)
	go func() {
		passphrase, err = prompt.PassPrompt(bufio.NewReader(os.Stdin), prefix, confirm)
		c <- struct{}{}
	}()
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-c:
		return passphrase, err
	}
}

// startPromptPass prompts the user for a password to unlock their wallet in
// the event that it was restored from seed or --promptpass flag is set.
func startPromptPass(ctx context.Context, w *wallet.Wallet) []byte {
	promptPass := cfg.PromptPass

	// Watching only wallets never require a password.
	if w.Manager.WatchingOnly() {
		return nil
	}

	// The wallet is totally desynced, so we need to resync accounts.
	// Prompt for the password. Then, set the flag it wallet so it
	// knows which address functions to call when resyncing.
	needSync, err := w.NeedsAccountsSync()
	if err != nil {
		log.Errorf("Error determining whether an accounts sync is necessary: %v", err)
	}
	if err == nil && needSync {
		fmt.Println("*** ATTENTION ***")
		fmt.Println("Since this is your first time running we need to sync accounts. Please enter")
		fmt.Println("the private wallet passphrase. This will complete syncing of the wallet")
		fmt.Println("accounts and then leave your wallet unlocked. You may relock wallet after by")
		fmt.Println("calling 'walletlock' through the RPC.")
		fmt.Println("*****************")
		promptPass = true
	}
	if cfg.EnableTicketBuyer {
		promptPass = true
	}

	if !promptPass {
		return nil
	}

	// We need to rescan accounts for the initial sync. Unlock the
	// wallet after prompting for the passphrase. The special case
	// of a --createtemp simnet wallet is handled by first
	// attempting to automatically open it with the default
	// passphrase. The wallet should also request to be unlocked
	// if stake mining is currently on, so users with this flag
	// are prompted here as well.
	for {
		if w.ChainParams() == &chaincfg.SimNetParams {
			err := w.Unlock(wallet.SimulationPassphrase, nil)
			if err == nil {
				// Unlock success with the default password.
				w.SetInitiallyUnlocked(true)
				return wallet.SimulationPassphrase
			}
		}

		passphrase, err := passPrompt(ctx, "Enter private passphrase", false)
		if err != nil {
			return nil
		}

		err = w.Unlock(passphrase, nil)
		if err != nil {
			fmt.Println("Incorrect password entered. Please " +
				"try again.")
			continue
		}
		w.SetInitiallyUnlocked(true)
		return passphrase
	}
}

// rpcClientConnectLoop loops forever, attempting to create a connection to the
// consensus RPC server.  If this connection succeeds, the RPC client is used as
// the loaded wallet's network backend and used to keep the wallet synchronized
// to the network.  If/when the RPC connection is lost, the wallet is
// disassociated from the client and a new connection is attempmted.
//
// The JSON-RPC server is optional.  If set, the connected RPC client will be
// associated with the server for RPC passthrough and to enable additional
// methods.
//
// This function panics if the wallet has not already been loaded.
func rpcClientConnectLoop(ctx context.Context, passphrase []byte, jsonRPCServer *legacyrpc.Server, loader *ldr.Loader) {
	w, ok := loader.LoadedWallet()
	if !ok {
		panic("rpcClientConnectLoop: called without loaded wallet")
	}

	certs := readCAFile()

	for {
		chainClient, err := startChainRPC(ctx, certs)
		if err != nil {
			return
		}

		n := chain.BackendFromRPCClient(chainClient.Client)
		w.SetNetworkBackend(n)
		if jsonRPCServer != nil {
			jsonRPCServer.SetChainServer(chainClient)
		}
		loader.SetChainClient(chainClient.Client)

		if cfg.EnableTicketBuyer {
			err = loader.StartTicketPurchase(passphrase, &cfg.tbCfg)
			if err != nil {
				log.Errorf("Unable to start ticket buyer: %v", err)
			}
		}

		// Run wallet synchronization until it is cancelled or errors.  If the
		// context was cancelled, return immediately instead of trying to
		// reconnect.
		syncer := chain.NewRPCSyncer(w, chainClient)
		err = syncer.Run(ctx, true)
		if err == context.Canceled {
			return
		}
		if err != nil {
			syncLog.Errorf("Wallet synchronization stopped: %v", err)
		}

		// Disassociate the RPC client from all subsystems until reconnection
		// occurs.
		w.SetNetworkBackend(nil)
		if jsonRPCServer != nil {
			jsonRPCServer.SetChainServer(nil)
		}
		loader.SetChainClient(nil)
		loader.StopTicketPurchase()
	}
}

func readCAFile() []byte {
	// Read certificate file if TLS is not disabled.
	var certs []byte
	if !cfg.DisableClientTLS {
		var err error
		certs, err = ioutil.ReadFile(cfg.CAFile.Value)
		if err != nil {
			log.Warnf("Cannot open CA file: %v", err)
			// If there's an error reading the CA file, continue
			// with nil certs and without the client connection.
			certs = nil
		}
	} else {
		log.Info("Chain server RPC TLS is disabled")
	}

	return certs
}

// startChainRPC opens a RPC client connection to a dtcd server for blockchain
// services.  This function uses the RPC options from the global config and
// there is no recovery in case the server is not available or if there is an
// authentication error.  Instead, all requests to the client will simply error.
func startChainRPC(ctx context.Context, certs []byte) (*chain.RPCClient, error) {
	log.Infof("Attempting RPC client connection to %v", cfg.RPCConnect)
	rpcc, err := chain.NewRPCClient(activeNet.Params, cfg.RPCConnect,
		cfg.DcrdUsername, cfg.DcrdPassword, certs, cfg.DisableClientTLS)
	if err != nil {
		return nil, err
	}
	err = rpcc.Start(ctx, true)
	return rpcc, err
}
<commit_msg>legacyrpc: rename GetChainServer to requireChainClient & update sig. (#1006)<commit_after>// Copyright (c) 2013-2015 The btcsuite developers
// Copyright (c) 2015-2017 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime"
	"runtime/pprof"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrwallet/apperrors"
	"github.com/decred/dcrwallet/chain"
	"github.com/decred/dcrwallet/internal/prompt"
	"github.com/decred/dcrwallet/internal/zero"
	ldr "github.com/decred/dcrwallet/loader"
	"github.com/decred/dcrwallet/rpc/legacyrpc"
	"github.com/decred/dcrwallet/rpc/rpcserver"
	"github.com/decred/dcrwallet/version"
	"github.com/decred/dcrwallet/wallet"
)

var (
	cfg *config
)

func main() {
	// Create a context that is cancelled when a shutdown request is received
	// through an interrupt signal or an RPC request.
	ctx := withShutdownCancel(context.Background())
	go shutdownListener()

	// Run the wallet until permanent failure or shutdown is requested.
	if err := run(ctx); err != nil && err != context.Canceled {
		os.Exit(1)
	}
}

// done returns whether the context's Done channel was closed due to
// cancellation or exceeded deadline.
func done(ctx context.Context) bool {
	select {
	case <-ctx.Done():
		return true
	default:
		return false
	}
}

// run is the main startup and teardown logic performed by the main package.  It
// is responsible for parsing the config, starting RPC servers, loading and
// syncing the wallet (if necessary), and stopping all started services when the
// context is cancelled.
func run(ctx context.Context) error {
	// Load configuration and parse command line.  This function also
	// initializes logging and configures it accordingly.
	tcfg, _, err := loadConfig(ctx)
	if err != nil {
		return err
	}
	cfg = tcfg
	defer func() {
		if logRotator != nil {
			logRotator.Close()
		}
	}()

	// Show version at startup.
	log.Infof("Version %s (Go version %s)", version.String(), runtime.Version())

	// Read IPC messages from the read end of a pipe created and passed by the
	// parent process, if any.  When this pipe is closed, shutdown is
	// initialized.
	if cfg.PipeRx != nil {
		go serviceControlPipeRx(uintptr(*cfg.PipeRx))
	}
	if cfg.PipeTx != nil {
		go serviceControlPipeTx(uintptr(*cfg.PipeTx))
	} else {
		go drainOutgoingPipeMessages()
	}

	// Run the pprof profiler if enabled.
	if len(cfg.Profile) > 0 {
		if done(ctx) {
			return ctx.Err()
		}

		profileRedirect := http.RedirectHandler("/debug/pprof", http.StatusSeeOther)
		http.Handle("/", profileRedirect)
		for _, listenAddr := range cfg.Profile {
			listenAddr := listenAddr // copy for closure
			go func() {
				log.Infof("Starting profile server on %s", listenAddr)
				err := http.ListenAndServe(listenAddr, nil)
				if err != nil {
					fatalf("Unable to run profiler: %v", err)
				}
			}()
		}
	}

	// Write mem profile if requested.
	if cfg.MemProfile != "" {
		if done(ctx) {
			return ctx.Err()
		}

		f, err := os.Create(cfg.MemProfile)
		if err != nil {
			log.Errorf("Unable to create cpu profile: %v", err)
			return err
		}
		timer := time.NewTimer(time.Minute * 5) // 5 minutes
		go func() {
			<-timer.C
			pprof.WriteHeapProfile(f)
			f.Close()
		}()
	}

	if done(ctx) {
		return ctx.Err()
	}

	// Create the loader which is used to load and unload the wallet.  If
	// --noinitialload is not set, this function is responsible for loading the
	// wallet.  Otherwise, loading is deferred so it can be performed over RPC.
	dbDir := networkDir(cfg.AppDataDir.Value, activeNet.Params)
	stakeOptions := &ldr.StakeOptions{
		VotingEnabled:       cfg.EnableVoting,
		AddressReuse:        cfg.ReuseAddresses,
		VotingAddress:       cfg.TBOpts.VotingAddress.Address,
		PoolAddress:         cfg.PoolAddress.Address,
		PoolFees:            cfg.PoolFees,
		StakePoolColdExtKey: cfg.StakePoolColdExtKey,
		TicketFee:           cfg.TicketFee.ToCoin(),
	}
	loader := ldr.NewLoader(activeNet.Params, dbDir, stakeOptions,
		cfg.AddrIdxScanLen, cfg.AllowHighFees, cfg.RelayFee.ToCoin())

	// Stop any services started by the loader after the shutdown procedure is
	// initialized and this function returns.
	defer func() {
		err := loader.UnloadWallet()
		if err != nil && err != ldr.ErrWalletNotLoaded {
			log.Errorf("Failed to close wallet: %v", err)
		} else if err == nil {
			log.Infof("Closed wallet")
		}
	}()

	// Open the wallet when --noinitialload was not set.
	passphrase := []byte{}
	if !cfg.NoInitialLoad {
		walletPass := []byte(cfg.WalletPass)
		defer zero.Bytes(walletPass)

		if cfg.PromptPublicPass {
			walletPass, _ = passPrompt(ctx, "Enter public wallet passphrase", false)
		}

		if done(ctx) {
			return ctx.Err()
		}

		// Load the wallet.  It must have been created already or this will
		// return an appropriate error.
		w, err := loader.OpenExistingWallet(walletPass)
		if err != nil {
			log.Errorf("Failed to open wallet: %v", err)
			if apperrors.IsError(err, apperrors.ErrWrongPassphrase) {
				// walletpass not provided, advice using --walletpass or --promptpublicpass
				if len(walletPass) == 0 {
					log.Info("Retry with --walletpass or --promptpublicpass.")
				}
			}

			return err
		}

		if done(ctx) {
			return ctx.Err()
		}

		// TODO(jrick): I think that this prompt should be removed
		// entirely instead of enabling it when --noinitialload is
		// unset.  It can be replaced with an RPC request (either
		// providing the private passphrase as a parameter, or require
		// unlocking the wallet first) to trigger a full accounts
		// rescan.
		//
		// Until then, since --noinitialload users are expecting to use
		// the wallet only over RPC, disable this feature for them.
		if cfg.Pass != "" {
			passphrase = []byte(cfg.Pass)
			err = w.Unlock(passphrase, nil)
			if err != nil {
				log.Errorf("Incorrect passphrase in pass config setting.")
				return err
			}
			w.SetInitiallyUnlocked(true)
		} else {
			passphrase = startPromptPass(ctx, w)
		}
	}

	if done(ctx) {
		return ctx.Err()
	}

	// Create and start the RPC servers to serve wallet client connections.  If
	// any of the servers can not be started, it will be nil.  If none of them
	// can be started, this errors since at least one server must run for the
	// wallet to be useful.
	//
	// Servers will be associated with a loaded wallet if it has already been
	// loaded, or after it is loaded later on.
	gRPCServer, jsonRPCServer, err := startRPCServers(loader)
	if err != nil {
		log.Errorf("Unable to create RPC servers: %v", err)
		return err
	}
	if gRPCServer != nil {
		// Start wallet and voting gRPC services after a wallet is loaded.
		loader.RunAfterLoad(func(w *wallet.Wallet) {
			rpcserver.StartWalletService(gRPCServer, w)
			rpcserver.StartVotingService(gRPCServer, w)
		})
		defer func() {
			log.Warn("Stopping gRPC server...")
			gRPCServer.Stop()
			log.Info("gRPC server shutdown")
		}()
	}
	if jsonRPCServer != nil {
		go func() {
			for range jsonRPCServer.RequestProcessShutdown() {
				requestShutdown()
			}
		}()
		defer func() {
			log.Warn("Stopping JSON-RPC server...")
			jsonRPCServer.Stop()
			log.Info("JSON-RPC server shutdown")
		}()
	}

	// Stop the ticket buyer (if running) on shutdown.  This returns an error
	// that can be ignored when the ticket buyer was never started.
	defer loader.StopTicketPurchase()

	// When not running with --noinitialload, it is the main package's
	// responsibility to connect the loaded wallet to the dcrd RPC server for
	// wallet synchronization.  This function blocks until cancelled.
	if !cfg.NoInitialLoad {
		if done(ctx) {
			return ctx.Err()
		}

		rpcClientConnectLoop(ctx, passphrase, jsonRPCServer, loader)
	}

	// Wait until shutdown is signaled before returning and running deferred
	// shutdown tasks.
	<-ctx.Done()
	return ctx.Err()
}

func passPrompt(ctx context.Context, prefix string, confirm bool) (passphrase []byte, err error) {
	os.Stdout.Sync()
	c := make(chan struct{}, 1)
	go func() {
		passphrase, err = prompt.PassPrompt(bufio.NewReader(os.Stdin), prefix, confirm)
		c <- struct{}{}
	}()
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-c:
		return passphrase, err
	}
}

// startPromptPass prompts the user for a password to unlock their wallet in
// the event that it was restored from seed or --promptpass flag is set.
func startPromptPass(ctx context.Context, w *wallet.Wallet) []byte {
	promptPass := cfg.PromptPass

	// Watching only wallets never require a password.
	if w.Manager.WatchingOnly() {
		return nil
	}

	// The wallet is totally desynced, so we need to resync accounts.
	// Prompt for the password. Then, set the flag it wallet so it
	// knows which address functions to call when resyncing.
	needSync, err := w.NeedsAccountsSync()
	if err != nil {
		log.Errorf("Error determining whether an accounts sync is necessary: %v", err)
	}
	if err == nil && needSync {
		fmt.Println("*** ATTENTION ***")
		fmt.Println("Since this is your first time running we need to sync accounts. Please enter")
		fmt.Println("the private wallet passphrase. This will complete syncing of the wallet")
		fmt.Println("accounts and then leave your wallet unlocked. You may relock wallet after by")
		fmt.Println("calling 'walletlock' through the RPC.")
		fmt.Println("*****************")
		promptPass = true
	}
	if cfg.EnableTicketBuyer {
		promptPass = true
	}

	if !promptPass {
		return nil
	}

	// We need to rescan accounts for the initial sync. Unlock the
	// wallet after prompting for the passphrase. The special case
	// of a --createtemp simnet wallet is handled by first
	// attempting to automatically open it with the default
	// passphrase. The wallet should also request to be unlocked
	// if stake mining is currently on, so users with this flag
	// are prompted here as well.
	for {
		if w.ChainParams() == &chaincfg.SimNetParams {
			err := w.Unlock(wallet.SimulationPassphrase, nil)
			if err == nil {
				// Unlock success with the default password.
				w.SetInitiallyUnlocked(true)
				return wallet.SimulationPassphrase
			}
		}

		passphrase, err := passPrompt(ctx, "Enter private passphrase", false)
		if err != nil {
			return nil
		}

		err = w.Unlock(passphrase, nil)
		if err != nil {
			fmt.Println("Incorrect password entered. Please " +
				"try again.")
			continue
		}
		w.SetInitiallyUnlocked(true)
		return passphrase
	}
}

// rpcClientConnectLoop loops forever, attempting to create a connection to the
// consensus RPC server.  If this connection succeeds, the RPC client is used as
// the loaded wallet's network backend and used to keep the wallet synchronized
// to the network.  If/when the RPC connection is lost, the wallet is
// disassociated from the client and a new connection is attempmted.
//
// The JSON-RPC server is optional.  If set, the connected RPC client will be
// associated with the server for RPC passthrough and to enable additional
// methods.
//
// This function panics if the wallet has not already been loaded.
func rpcClientConnectLoop(ctx context.Context, passphrase []byte, jsonRPCServer *legacyrpc.Server, loader *ldr.Loader) {
	w, ok := loader.LoadedWallet()
	if !ok {
		panic("rpcClientConnectLoop: called without loaded wallet")
	}

	certs := readCAFile()

	for {
		chainClient, err := startChainRPC(ctx, certs)
		if err != nil {
			return
		}

		n := chain.BackendFromRPCClient(chainClient.Client)
		w.SetNetworkBackend(n)
		if jsonRPCServer != nil {
			jsonRPCServer.SetChainServer(chainClient)
		}
		loader.SetChainClient(chainClient.Client)

		if cfg.EnableTicketBuyer {
			err = loader.StartTicketPurchase(passphrase, &cfg.tbCfg)
			if err != nil {
				log.Errorf("Unable to start ticket buyer: %v", err)
			}
		}

		// Run wallet synchronization until it is cancelled or errors.  If the
		// context was cancelled, return immediately instead of trying to
		// reconnect.
		syncer := chain.NewRPCSyncer(w, chainClient)
		err = syncer.Run(ctx, true)
		if err == context.Canceled {
			return
		}
		if err != nil {
			syncLog.Errorf("Wallet synchronization stopped: %v", err)
		}

		// Disassociate the RPC client from all subsystems until reconnection
		// occurs.
		w.SetNetworkBackend(nil)
		if jsonRPCServer != nil {
			jsonRPCServer.SetChainServer(nil)
		}
		loader.SetChainClient(nil)
		loader.StopTicketPurchase()
	}
}

func readCAFile() []byte {
	// Read certificate file if TLS is not disabled.
	var certs []byte
	if !cfg.DisableClientTLS {
		var err error
		certs, err = ioutil.ReadFile(cfg.CAFile.Value)
		if err != nil {
			log.Warnf("Cannot open CA file: %v", err)
			// If there's an error reading the CA file, continue
			// with nil certs and without the client connection.
			certs = nil
		}
	} else {
		log.Info("Chain server RPC TLS is disabled")
	}

	return certs
}

// startChainRPC opens a RPC client connection to a dtcd server for blockchain
// services.  This function uses the RPC options from the global config and
// there is no recovery in case the server is not available or if there is an
// authentication error.  Instead, all requests to the client will simply error.
func startChainRPC(ctx context.Context, certs []byte) (*chain.RPCClient, error) {
	log.Infof("Attempting RPC client connection to %v", cfg.RPCConnect)
	rpcc, err := chain.NewRPCClient(activeNet.Params, cfg.RPCConnect,
		cfg.DcrdUsername, cfg.DcrdPassword, certs, cfg.DisableClientTLS)
	if err != nil {
		return nil, err
	}
	err = rpcc.Start(ctx, true)
	return rpcc, err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"strings"
	"time"
	"mime"
	"io"

	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/net/publicsuffix"
)

type export struct {
	URL       string `json:"url"`
	Job       string `json:"job"`
	Finished  bool   `json:"finished"`
	Processed int    `json:"processed"`
}

func prompt() (username, password string) {
	// Get the username
	fmt.Print("Username: ")
	reader := bufio.NewReader(os.Stdin)
	usernameinput, err := reader.ReadString('\n')
	if err != nil {
		fmt.Println("Unexpected error reading username.")
		fmt.Println(err)
		os.Exit(1)
	}

	// Get the password
	fmt.Print("Password: ")
	passwordinput, err := terminal.ReadPassword(0)
	if err != nil {
		fmt.Println("Unexpected error reading password.")
		fmt.Println(err)
		os.Exit(1)
	}
	fmt.Println()

	return strings.TrimSpace(usernameinput), string(passwordinput)
}

func login(client *http.Client) {

	for {
		username, password := prompt()

		resp, err := client.PostForm("https://ir.library.carleton.ca/login",
			url.Values{"userid": {username}, "password": {password}})
		if err != nil {
			fmt.Println("Unexpected error logging in.")
			fmt.Println(err)
			os.Exit(1)
		}
		resp.Body.Close()

		if resp.Request.URL.Path == "/" {
			break
		} else {
			fmt.Println("Error logging in, retry...")
		}
	}

}

func exportData(irtype string, client *http.Client) {

	path := fmt.Sprintf("/%v/export", irtype)
	var curExport = &export{}

	for !curExport.Finished {

		url := fmt.Sprintf("https://ir.library.carleton.ca%v?job=%v", path, curExport.Job)

		resp, err := client.Get(url)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		if resp.StatusCode != 200 || resp.Request.URL.Path != path {
			buf := new(bytes.Buffer)
			buf.ReadFrom(resp.Body)
			fmt.Println(buf.String())
			resp.Body.Close()
			os.Exit(1)
		}

		err = json.NewDecoder(resp.Body).Decode(curExport)
		resp.Body.Close()

		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		fmt.Printf("%+v\n", curExport)
		time.Sleep(1 * time.Second)
	}

	resp, err := client.Get(curExport.URL)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	_, params, err := mime.ParseMediaType(resp.Header.Get("Content-Disposition"))
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	filename := fmt.Sprintf("%v-%v", time.Now().Format("2006-01-02"), params["filename"])

	fmt.Printf("Saving to filename %v\n", filename)

	//open a file for writing
    file, err := os.Create(filename)
    if err != nil {
		fmt.Println(err)
		os.Exit(1)
    }
    // Use io.Copy to just dump the response body to the file. This supports huge files
    _, err = io.Copy(file, resp.Body)
    if err != nil {
		fmt.Println(err)
		os.Exit(1)
    }
    file.Close()
    fmt.Println("Success!")

}

func main() {

	jar, err := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
	if err != nil {
		fmt.Println("Unexpected error creating cookie storage.")
		fmt.Println(err)
		os.Exit(1)
	}

	client := &http.Client{
		Jar: jar,
	}

	login(client)

	fmt.Println("Login complete!")

	exportData("col", client)

}
<commit_msg>Removed print statements<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"strings"
	"time"

	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/net/publicsuffix"
)

type export struct {
	URL       string `json:"url"`
	Job       string `json:"job"`
	Finished  bool   `json:"finished"`
	Processed int    `json:"processed"`
}

func prompt() (username, password string) {
	// Get the username
	fmt.Print("Username: ")
	reader := bufio.NewReader(os.Stdin)
	usernameinput, err := reader.ReadString('\n')
	if err != nil {
		fmt.Println("Unexpected error reading username.")
		fmt.Println(err)
		os.Exit(1)
	}

	// Get the password
	fmt.Print("Password: ")
	passwordinput, err := terminal.ReadPassword(0)
	if err != nil {
		fmt.Println("Unexpected error reading password.")
		fmt.Println(err)
		os.Exit(1)
	}
	fmt.Println()

	return strings.TrimSpace(usernameinput), string(passwordinput)
}

func login(client *http.Client) {

	for {
		username, password := prompt()

		resp, err := client.PostForm("https://ir.library.carleton.ca/login",
			url.Values{"userid": {username}, "password": {password}})
		if err != nil {
			fmt.Println("Unexpected error logging in.")
			fmt.Println(err)
			os.Exit(1)
		}
		resp.Body.Close()

		if resp.Request.URL.Path == "/" {
			break
		} else {
			fmt.Println("Error logging in, retry...")
		}
	}

}

func exportData(irtype string, client *http.Client) {

	path := fmt.Sprintf("/%v/export", irtype)
	var curExport = &export{}

	for !curExport.Finished {

		url := fmt.Sprintf("https://ir.library.carleton.ca%v?job=%v", path, curExport.Job)

		resp, err := client.Get(url)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		if resp.StatusCode != 200 || resp.Request.URL.Path != path {
			buf := new(bytes.Buffer)
			buf.ReadFrom(resp.Body)
			fmt.Println(buf.String())
			resp.Body.Close()
			os.Exit(1)
		}

		err = json.NewDecoder(resp.Body).Decode(curExport)
		resp.Body.Close()

		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		fmt.Printf("%+v\n", curExport)
		time.Sleep(1 * time.Second)
	}

	resp, err := client.Get(curExport.URL)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	_, params, err := mime.ParseMediaType(resp.Header.Get("Content-Disposition"))
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	filename := fmt.Sprintf("%v-%v", time.Now().Format("2006-01-02"), params["filename"])

	fmt.Printf("Saving to filename %v\n", filename)

	file, err := os.Create(filename)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	_, err = io.Copy(file, resp.Body)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	file.Close()

}

func main() {

	jar, err := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
	if err != nil {
		fmt.Println("Unexpected error creating cookie storage.")
		fmt.Println(err)
		os.Exit(1)
	}

	client := &http.Client{
		Jar: jar,
	}

	login(client)

	fmt.Println("Login complete!")

	exportData("col", client)

}
<|endoftext|>


<commit_before>package app

import (
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/pierrec/lz4"
)

// D is the handler for decompression
func D(files []string) (msg string, err error) {
	if len(files) == 0 {
		return "", errors.New("decompress: no file specified")
	}
	// go through each file, check for compression type, and call appropriately
	for _, file := range files {

	}
	return msg, err
}
<commit_msg>add d command and support for lz4<commit_after>package app

import (
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/pierrec/lz4"
)

// D is the handler for decompression
func D(files []string) (msg string, err error) {
	if len(files) == 0 {
		return "", errors.New("decompress: no file specified")
	}
	// Process each file
	for _, file := range files {
		err = DToFile(file)
	}
	return msg, err
}

// DToFile decompresses each source to a file.
func DToFile(fname string) error {
	// open the file
	src, err := os.Open(fname)
	if err != nil {
		return err
	}
	// check its format
	format := getFileFormat(src)
	if format == Unsupported {
		return fmt.Errorf("%s's format is unsupported", fname)
	}
	dFname, err := dOutFile(fname)
	// open output file
	dst, err := os.OpenFile(dFname, os.O_CREATE|os.O_WRONLY, 0755)
	if err != nil {
		return err
	}
	// extract
	err = dlz4(dst, src)
	return nil
}

func dlz4(dst io.Writer, src io.Reader) error {
	lzr := lz4.NewReader(src)
	_, err := io.Copy(dst, lzr)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>