<commit_before>package main

import "bytes"

const (
	templateSQL    = "sql"
	contentTypeSQL = "text/sql"
)

type TemplateSQL struct {
	schema     Schema
	linePrefix []byte
}

func (template TemplateSQL) Generate(chunk *Chunk) ([]byte, error) {
	var buffer bytes.Buffer
	buffer.Write(template.getLinePrefix(chunk))
	for _, column := range template.schema.columns {
		val, err := column.Value(chunk)
		if err != nil {
			return nil, err
		}
		chunk.values[column.Column().name] = val
		buffer.WriteString(" '")
		buffer.WriteString(val)
		buffer.WriteString("',")
	}
	buffer.Truncate(buffer.Len() - 1)
	buffer.WriteString(" );")
	buffer.WriteByte('\n')
	return buffer.Bytes(), nil
}

func (template TemplateSQL) ContentType() string {
	return contentTypeSQL
}

func (template TemplateSQL) getLinePrefix(chunk *Chunk) []byte {
	if template.linePrefix != nil {
		return template.linePrefix
	}

	var buffer bytes.Buffer
	buffer.WriteString("INSERT INTO ")
	buffer.WriteString(template.schema.document)
	buffer.WriteString(" ( ")
	for _, column := range template.schema.columns {
		buffer.WriteString(column.Column().name)
		buffer.WriteString(", ")
	}
	buffer.Truncate(buffer.Len() - 2)
	buffer.WriteString(" ) VALUES (")
	return buffer.Bytes()
}
<commit_msg>fix: lineprefix caching for template: sql<commit_after>package main

import "bytes"

const (
	templateSQL    = "sql"
	contentTypeSQL = "text/sql"
)

type TemplateSQL struct {
	schema     Schema
	linePrefix []byte
}

func (template TemplateSQL) Generate(chunk *Chunk) ([]byte, error) {
	var buffer bytes.Buffer
	buffer.Write(template.getLinePrefix(chunk))
	for _, column := range template.schema.columns {
		val, err := column.Value(chunk)
		if err != nil {
			return nil, err
		}
		chunk.values[column.Column().name] = val
		buffer.WriteString(" '")
		buffer.WriteString(val)
		buffer.WriteString("',")
	}
	buffer.Truncate(buffer.Len() - 1)
	buffer.WriteString(" );")
	buffer.WriteByte('\n')
	return buffer.Bytes(), nil
}

func (template TemplateSQL) ContentType() string {
	return contentTypeSQL
}

func (template TemplateSQL) getLinePrefix(chunk *Chunk) []byte {
	if template.linePrefix == nil {
		var buffer bytes.Buffer
		buffer.WriteString("INSERT INTO ")
		buffer.WriteString(template.schema.document)
		buffer.WriteString(" ( ")
		for _, column := range template.schema.columns {
			buffer.WriteString(column.Column().name)
			buffer.WriteString(", ")
		}
		buffer.Truncate(buffer.Len() - 2)
		buffer.WriteString(" ) VALUES (")
		template.linePrefix = buffer.Bytes()
	}

	return template.linePrefix
}
<|endoftext|>


<commit_before>// This is a simple implementation of a console-based Tetris clone. It only works where termbox-go works
// (linux/mac should be fine).
package tetris

import (
	"github.com/nsf/termbox-go"
	"math"
	"math/rand"
	"time"
)

type Direction int

const (
	Up Direction = iota + 1
	Down
	Left
	Right
)

// A Game tracks the entire game state of tetris, including the Board, the upcoming piece, the game speed
// (dropDelayMillis), the score, and various other internal data.
type Game struct {
	board           *Board
	nextPiece       *Piece
	pieces          []Piece
	paused          bool
	over            bool
	dropDelayMillis int
	ticker          *time.Ticker
	score           int
}

// Initialize a new game, ready to be started with Start().
func NewGame() *Game {
	game := new(Game)
	game.pieces = tetrisPieces()
	game.board = newBoard()
	game.board.currentPiece = game.GeneratePiece()
	game.board.currentPosition = Vector{initialX, 0}
	game.nextPiece = game.GeneratePiece()
	game.paused = false
	game.over = false
	game.score = 0
	game.startTicker()
	return game
}

// Start up the ticker with the appropriate interval based on the current score.
func (game *Game) startTicker() {
	// Set the speed as a function of score. Starts at 800ms, decreases to 200ms by 100ms each 500 points.
	game.dropDelayMillis = 800 - game.score/5
	if game.dropDelayMillis < 200 {
		game.dropDelayMillis = 200
	}
	game.ticker = time.NewTicker(time.Duration(game.dropDelayMillis) * time.Millisecond)
}

// Stop the game ticker. This stops automatic advancement of the piece.
func (game *Game) stopTicker() {
	game.ticker.Stop()
}

// A game event, generated by user input or by the game ticker.
type GameEvent int

const (
	MoveLeft GameEvent = iota
	MoveRight
	MoveDown
	Rotate
	QuickDrop
	Pause
	Quit
	// An event that doesn't cause a change to game state but causes a full redraw; e.g., a window resize.
	Redraw
)

// Start running the game. It will continue indefinitely until the user exits.
func (game *Game) Start() {

	drawStaticBoardParts()
	game.DrawDynamic(false)

	eventQueue := make(chan GameEvent, 100)
	go func() {
		for {
			eventQueue <- waitForUserEvent()
		}
	}()
gameOver:
	for {
		var event GameEvent
		select {
		case event = <-eventQueue:
		case <-game.ticker.C:
			event = MoveDown
		}
		// If the game is paused, ignore all commands except for Pause, Quit, and Redraw. On Redraw, redraw
		// the pause screen.
		if game.paused {
			switch event {
			case Pause:
				game.PauseToggle()
			case Quit:
				return
			case Redraw:
				drawStaticBoardParts()
				game.DrawPauseScreen()
			}
		} else {
			switch event {
			case MoveLeft:
				game.Move(Left)
			case MoveRight:
				game.Move(Right)
			case MoveDown:
				game.Move(Down)
			case QuickDrop:
				game.QuickDrop()
			case Rotate:
				game.Rotate()
			case Pause:
				game.PauseToggle()
			case Quit:
				return
			case Redraw:
				drawStaticBoardParts()
			}
			// Update screen only if game is not now paused.
			if !game.paused {
				game.DrawDynamic(false)
			}
		}
		if game.over {
			break gameOver
		}
	}
	game.DrawGameOver()
	for waitForUserEvent() != Quit {
	}
}

// A blocking function that waits on a ticker and then emits a MoveDown event.
func waitForTick(ticker *time.Ticker) GameEvent {
	<-ticker.C
	return MoveDown
}

// A blocking function that waits for user input and then emits the appropriate GameEvent.
func waitForUserEvent() GameEvent {
	switch event := termbox.PollEvent(); event.Type {
	// Movement: arrow keys or vim controls (h, j, k, l)
	// Pause: 'p'
	// Exit: 'q' or ctrl-c.
	case termbox.EventKey:
		if event.Ch == 0 { // A special key combo was pressed
			switch event.Key {
			case termbox.KeyCtrlC:
				return Quit
			case termbox.KeyArrowLeft:
				return MoveLeft
			case termbox.KeyArrowUp:
				return Rotate
			case termbox.KeyArrowRight:
				return MoveRight
			case termbox.KeyArrowDown:
				return MoveDown
			case termbox.KeySpace:
				return QuickDrop
			}
		} else {
			switch event.Ch {
			case 'p':
				return Pause
			case 'q':
				return Quit
			case 'h':
				return MoveLeft
			case 'k':
				return Rotate
			case 'l':
				return MoveRight
			case 'j':
				return MoveDown
			}
		}
	case termbox.EventResize:
		return Redraw
	case termbox.EventError:
		panic(event.Err)
	}
	return Redraw // Should never be reached
}

// Randomly choose a new game piece from among the the available pieces.
func (game *Game) GeneratePiece() *Piece {
	return &game.pieces[rand.Intn(len(game.pieces))]
}

// Anchor the current piece to the board, clears lines, increments the score, and generate a new piece. Sets
// the 'game over' state if the new piece overlaps existing pieces.
func (game *Game) anchor() {
	game.board.mergeCurrentPiece()

	// Clear any completed rows (with animation) and increment the score if necessary.
	rowsCleared := game.board.clearedRows()

	if len(rowsCleared) > 0 {
		// Animate the cleared rows disappearing
		game.stopTicker()
		flickerCells := make(map[Vector]termbox.Attribute)
		for _, y := range rowsCleared {
			for x := 0; x < width; x++ {
				point := Vector{x, y}
				flickerCells[point] = game.board.cells[point]
			}
		}
		for i := 0; i < 5; i++ {
			for point, color := range flickerCells {
				if i%2 == 0 {
					color = backgroundColor
				}
				setBoardCell((point.x*2)+2, headerHeight+point.y+2, color)
			}
			termbox.Flush()
			time.Sleep(80 * time.Millisecond)
		}

		// Get rid of the rows
		game.board.clearRows()

		// Scoring -- 1 row -> 100, 2 rows -> 200, ... 4 rows -> 800
		points := 100 * math.Pow(2, float64(len(rowsCleared)-1))
		game.score += int(points)

		game.startTicker()
	}

	// Bring in the next piece.
	game.board.currentPiece = game.nextPiece
	game.board.currentPosition = Vector{initialX, 0}
	game.nextPiece = game.GeneratePiece()
	game.nextPiece.currentRotation = 0

	if game.board.currentPieceInCollision() {
		game.over = true
	}
}

// Attempt to move.
func (game *Game) Move(where Direction) {
	translation := Vector{0, 0}
	switch where {
	case Down:
		translation = Vector{0, 1}
	case Left:
		translation = Vector{-1, 0}
	case Right:
		translation = Vector{1, 0}
	}
	// Attempt to make the move.
	moved := game.board.moveIfPossible(translation)

	// Perform anchoring if we tried to move down but we were unsuccessful.
	if where == Down && !moved {
		game.anchor()
	}
}

// Drop the piece all the way and anchor it.
func (game *Game) QuickDrop() {
	// Move down as far as possible
	for game.board.moveIfPossible(Vector{0, 1}) {
	}
	game.DrawDynamic(false)
	game.anchor()
}

// Rotates the current game piece, if possible.
func (game *Game) Rotate() {
	game.board.currentPiece.rotate()
	if game.board.currentPieceInCollision() {
		game.board.currentPiece.unrotate()
	}
}

// Draw the dynamic parts of the game interface (the board, the next piece preview pane, and the score).  The
// static parts should be drawn with the drawStaticBoardParts() function, if needed.  If clearOnly is true,
// the board and preview pane will be cleared rather than redrawn.
func (game *Game) DrawDynamic(clearOnly bool) {

	// Print the board contents. Each block will correspond to a side-by-side pair of cells in the termbox, so
	// that the visible blocks will be roughly square.  If clearOnly is true, draw background color.
	for x := 0; x < width; x++ {
		for y := 0; y < height; y++ {
			if clearOnly {
				setBoardCell((x*2)+2, headerHeight+y+2, backgroundColor)
			} else {
				color := game.board.CellColor(Vector{x, y})
				setBoardCell((x*2)+2, headerHeight+y+2, color)
			}
		}
	}

	// Print the preview piece. Need to clear the box first.  Draw next piece only if clearOnly is false
	previewPieceOffset := Vector{(width * 2) + 8, headerHeight + 3}
	for x := 0; x < 6; x++ {
		for y := 0; y < 4; y++ {
			cursor := previewPieceOffset.plus(Vector{x, y})
			setCell(cursor.x, cursor.y, ' ', termbox.ColorDefault)
		}
	}
	if !clearOnly {
		for _, point := range game.nextPiece.rotations[0] {
			cursor := previewPieceOffset.plus(Vector{point.x * 2, point.y})
			setBoardCell(cursor.x, cursor.y, game.nextPiece.color)
		}
	}

	// Draw the current score.  If clearOnly, do the same.
	score := game.score
	cursor := Vector{(width * 2) + 18, headerHeight + previewHeight + 7}
	for {
		digit := score % 10
		score /= 10
		drawDigitAsAscii(cursor.x, cursor.y, digit)
		cursor = cursor.plus(Vector{-4, 0})
		if score == 0 {
			break
		}
	}

	// Flush termbox's internal state to the screen.
	termbox.Flush()
}

// Pause or unpause the game, depending on game.paused.
func (game *Game) PauseToggle() {
	if game.paused {
		drawStaticBoardParts()
		game.DrawDynamic(false)
		game.startTicker()
	} else {
		game.stopTicker()
		game.DrawPauseScreen()
	}
	game.paused = !game.paused
}

// Draw the pause screen, hiding the game board and next piece.
func (game *Game) DrawPauseScreen() {
	// Clear the board and preview screen
	game.DrawDynamic(true)

	// Draw PAUSED overlay
	for y := (totalHeight/2 - 1); y <= (totalHeight/2)+1; y++ {
		for x := 1; x < totalWidth+3; x++ {
			termbox.SetCell(x, y, ' ', termbox.ColorDefault, termbox.ColorBlue)
		}
	}
	for i, ch := range "PAUSED" {
		termbox.SetCell(totalWidth/2-2+i, totalHeight/2, ch, termbox.ColorWhite, termbox.ColorBlue)
	}

	// Flush termbox to screen
	termbox.Flush()
}

// Draw the "GAME OVER" overlay on top of the game interface.
func (game *Game) DrawGameOver() {
	for y := (totalHeight/2 - 1); y <= (totalHeight/2)+1; y++ {
		for x := 1; x < totalWidth+3; x++ {
			termbox.SetCell(x, y, ' ', termbox.ColorDefault, termbox.ColorBlue)
		}
	}
	for i, ch := range "GAME OVER" {
		termbox.SetCell(totalWidth/2-4+i, totalHeight/2, ch, termbox.ColorWhite, termbox.ColorBlue)
	}
	termbox.Flush()
}
<commit_msg>Fix a bug with initial piece rotation.<commit_after>// This is a simple implementation of a console-based Tetris clone. It only works where termbox-go works
// (linux/mac should be fine).
package tetris

import (
	"github.com/nsf/termbox-go"
	"math"
	"math/rand"
	"time"
)

type Direction int

const (
	Up Direction = iota + 1
	Down
	Left
	Right
)

// A Game tracks the entire game state of tetris, including the Board, the upcoming piece, the game speed
// (dropDelayMillis), the score, and various other internal data.
type Game struct {
	board           *Board
	nextPiece       *Piece
	pieces          []Piece
	paused          bool
	over            bool
	dropDelayMillis int
	ticker          *time.Ticker
	score           int
}

// Initialize a new game, ready to be started with Start().
func NewGame() *Game {
	game := new(Game)
	game.pieces = tetrisPieces()
	game.board = newBoard()
	game.board.currentPiece = game.GeneratePiece()
	game.board.currentPosition = Vector{initialX, 0}
	game.nextPiece = game.GeneratePiece()
	game.paused = false
	game.over = false
	game.score = 0
	game.startTicker()
	return game
}

// Start up the ticker with the appropriate interval based on the current score.
func (game *Game) startTicker() {
	// Set the speed as a function of score. Starts at 800ms, decreases to 200ms by 100ms each 500 points.
	game.dropDelayMillis = 800 - game.score/5
	if game.dropDelayMillis < 200 {
		game.dropDelayMillis = 200
	}
	game.ticker = time.NewTicker(time.Duration(game.dropDelayMillis) * time.Millisecond)
}

// Stop the game ticker. This stops automatic advancement of the piece.
func (game *Game) stopTicker() {
	game.ticker.Stop()
}

// A game event, generated by user input or by the game ticker.
type GameEvent int

const (
	MoveLeft GameEvent = iota
	MoveRight
	MoveDown
	Rotate
	QuickDrop
	Pause
	Quit
	// An event that doesn't cause a change to game state but causes a full redraw; e.g., a window resize.
	Redraw
)

// Start running the game. It will continue indefinitely until the user exits.
func (game *Game) Start() {

	drawStaticBoardParts()
	game.DrawDynamic(false)

	eventQueue := make(chan GameEvent, 100)
	go func() {
		for {
			eventQueue <- waitForUserEvent()
		}
	}()
gameOver:
	for {
		var event GameEvent
		select {
		case event = <-eventQueue:
		case <-game.ticker.C:
			event = MoveDown
		}
		// If the game is paused, ignore all commands except for Pause, Quit, and Redraw. On Redraw, redraw
		// the pause screen.
		if game.paused {
			switch event {
			case Pause:
				game.PauseToggle()
			case Quit:
				return
			case Redraw:
				drawStaticBoardParts()
				game.DrawPauseScreen()
			}
		} else {
			switch event {
			case MoveLeft:
				game.Move(Left)
			case MoveRight:
				game.Move(Right)
			case MoveDown:
				game.Move(Down)
			case QuickDrop:
				game.QuickDrop()
			case Rotate:
				game.Rotate()
			case Pause:
				game.PauseToggle()
			case Quit:
				return
			case Redraw:
				drawStaticBoardParts()
			}
			// Update screen only if game is not now paused.
			if !game.paused {
				game.DrawDynamic(false)
			}
		}
		if game.over {
			break gameOver
		}
	}
	game.DrawGameOver()
	for waitForUserEvent() != Quit {
	}
}

// A blocking function that waits on a ticker and then emits a MoveDown event.
func waitForTick(ticker *time.Ticker) GameEvent {
	<-ticker.C
	return MoveDown
}

// A blocking function that waits for user input and then emits the appropriate GameEvent.
func waitForUserEvent() GameEvent {
	switch event := termbox.PollEvent(); event.Type {
	// Movement: arrow keys or vim controls (h, j, k, l)
	// Pause: 'p'
	// Exit: 'q' or ctrl-c.
	case termbox.EventKey:
		if event.Ch == 0 { // A special key combo was pressed
			switch event.Key {
			case termbox.KeyCtrlC:
				return Quit
			case termbox.KeyArrowLeft:
				return MoveLeft
			case termbox.KeyArrowUp:
				return Rotate
			case termbox.KeyArrowRight:
				return MoveRight
			case termbox.KeyArrowDown:
				return MoveDown
			case termbox.KeySpace:
				return QuickDrop
			}
		} else {
			switch event.Ch {
			case 'p':
				return Pause
			case 'q':
				return Quit
			case 'h':
				return MoveLeft
			case 'k':
				return Rotate
			case 'l':
				return MoveRight
			case 'j':
				return MoveDown
			}
		}
	case termbox.EventResize:
		return Redraw
	case termbox.EventError:
		panic(event.Err)
	}
	return Redraw // Should never be reached
}

// Randomly choose a new game piece from among the the available pieces.
func (game *Game) GeneratePiece() *Piece {
	return &game.pieces[rand.Intn(len(game.pieces))]
}

// Anchor the current piece to the board, clears lines, increments the score, and generate a new piece. Sets
// the 'game over' state if the new piece overlaps existing pieces.
func (game *Game) anchor() {
	game.board.mergeCurrentPiece()

	// Clear any completed rows (with animation) and increment the score if necessary.
	rowsCleared := game.board.clearedRows()

	if len(rowsCleared) > 0 {
		// Animate the cleared rows disappearing
		game.stopTicker()
		flickerCells := make(map[Vector]termbox.Attribute)
		for _, y := range rowsCleared {
			for x := 0; x < width; x++ {
				point := Vector{x, y}
				flickerCells[point] = game.board.cells[point]
			}
		}
		for i := 0; i < 5; i++ {
			for point, color := range flickerCells {
				if i%2 == 0 {
					color = backgroundColor
				}
				setBoardCell((point.x*2)+2, headerHeight+point.y+2, color)
			}
			termbox.Flush()
			time.Sleep(80 * time.Millisecond)
		}

		// Get rid of the rows
		game.board.clearRows()

		// Scoring -- 1 row -> 100, 2 rows -> 200, ... 4 rows -> 800
		points := 100 * math.Pow(2, float64(len(rowsCleared)-1))
		game.score += int(points)

		game.startTicker()
	}

	// Bring in the next piece.
	game.board.currentPiece = game.nextPiece
	game.board.currentPosition = Vector{initialX, 0}
	game.board.currentPiece.currentRotation = 0
	game.nextPiece = game.GeneratePiece()

	if game.board.currentPieceInCollision() {
		game.over = true
	}
}

// Attempt to move.
func (game *Game) Move(where Direction) {
	translation := Vector{0, 0}
	switch where {
	case Down:
		translation = Vector{0, 1}
	case Left:
		translation = Vector{-1, 0}
	case Right:
		translation = Vector{1, 0}
	}
	// Attempt to make the move.
	moved := game.board.moveIfPossible(translation)

	// Perform anchoring if we tried to move down but we were unsuccessful.
	if where == Down && !moved {
		game.anchor()
	}
}

// Drop the piece all the way and anchor it.
func (game *Game) QuickDrop() {
	// Move down as far as possible
	for game.board.moveIfPossible(Vector{0, 1}) {
	}
	game.DrawDynamic(false)
	game.anchor()
}

// Rotates the current game piece, if possible.
func (game *Game) Rotate() {
	game.board.currentPiece.rotate()
	if game.board.currentPieceInCollision() {
		game.board.currentPiece.unrotate()
	}
}

// Draw the dynamic parts of the game interface (the board, the next piece preview pane, and the score).  The
// static parts should be drawn with the drawStaticBoardParts() function, if needed.  If clearOnly is true,
// the board and preview pane will be cleared rather than redrawn.
func (game *Game) DrawDynamic(clearOnly bool) {

	// Print the board contents. Each block will correspond to a side-by-side pair of cells in the termbox, so
	// that the visible blocks will be roughly square.  If clearOnly is true, draw background color.
	for x := 0; x < width; x++ {
		for y := 0; y < height; y++ {
			if clearOnly {
				setBoardCell((x*2)+2, headerHeight+y+2, backgroundColor)
			} else {
				color := game.board.CellColor(Vector{x, y})
				setBoardCell((x*2)+2, headerHeight+y+2, color)
			}
		}
	}

	// Print the preview piece. Need to clear the box first.  Draw next piece only if clearOnly is false
	previewPieceOffset := Vector{(width * 2) + 8, headerHeight + 3}
	for x := 0; x < 6; x++ {
		for y := 0; y < 4; y++ {
			cursor := previewPieceOffset.plus(Vector{x, y})
			setCell(cursor.x, cursor.y, ' ', termbox.ColorDefault)
		}
	}
	if !clearOnly {
		for _, point := range game.nextPiece.rotations[0] {
			cursor := previewPieceOffset.plus(Vector{point.x * 2, point.y})
			setBoardCell(cursor.x, cursor.y, game.nextPiece.color)
		}
	}

	// Draw the current score.  If clearOnly, do the same.
	score := game.score
	cursor := Vector{(width * 2) + 18, headerHeight + previewHeight + 7}
	for {
		digit := score % 10
		score /= 10
		drawDigitAsAscii(cursor.x, cursor.y, digit)
		cursor = cursor.plus(Vector{-4, 0})
		if score == 0 {
			break
		}
	}

	// Flush termbox's internal state to the screen.
	termbox.Flush()
}

// Pause or unpause the game, depending on game.paused.
func (game *Game) PauseToggle() {
	if game.paused {
		drawStaticBoardParts()
		game.DrawDynamic(false)
		game.startTicker()
	} else {
		game.stopTicker()
		game.DrawPauseScreen()
	}
	game.paused = !game.paused
}

// Draw the pause screen, hiding the game board and next piece.
func (game *Game) DrawPauseScreen() {
	// Clear the board and preview screen
	game.DrawDynamic(true)

	// Draw PAUSED overlay
	for y := (totalHeight/2 - 1); y <= (totalHeight/2)+1; y++ {
		for x := 1; x < totalWidth+3; x++ {
			termbox.SetCell(x, y, ' ', termbox.ColorDefault, termbox.ColorBlue)
		}
	}
	for i, ch := range "PAUSED" {
		termbox.SetCell(totalWidth/2-2+i, totalHeight/2, ch, termbox.ColorWhite, termbox.ColorBlue)
	}

	// Flush termbox to screen
	termbox.Flush()
}

// Draw the "GAME OVER" overlay on top of the game interface.
func (game *Game) DrawGameOver() {
	for y := (totalHeight/2 - 1); y <= (totalHeight/2)+1; y++ {
		for x := 1; x < totalWidth+3; x++ {
			termbox.SetCell(x, y, ' ', termbox.ColorDefault, termbox.ColorBlue)
		}
	}
	for i, ch := range "GAME OVER" {
		termbox.SetCell(totalWidth/2-4+i, totalHeight/2, ch, termbox.ColorWhite, termbox.ColorBlue)
	}
	termbox.Flush()
}
<|endoftext|>


<commit_before>package nsqthumbnailer

import (
	"bytes"
	"fmt"
	"image"
	"log"
	"math"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/disintegration/imaging"
	"gopkg.in/amz.v1/aws"
	"gopkg.in/amz.v1/s3"
)

var (
	awsAuth aws.Auth
)

func init() {
	awsAuth = newAwsAuth()
}

// This function used internally to convert any image type to NRGBA if needed.
// copied from `imaging`
func toNRGBA(img image.Image) *image.NRGBA {
	srcBounds := img.Bounds()
	if srcBounds.Min.X == 0 && srcBounds.Min.Y == 0 {
		if src0, ok := img.(*image.NRGBA); ok {
			return src0
		}
	}
	return imaging.Clone(img)
}

func newAwsAuth() aws.Auth {
	// Authenticate and Create an aws S3 service
	auth, err := aws.EnvAuth()
	if err != nil {
		panic(err.Error())
	}
	return auth
}

type imageOpenSaverError struct {
	url *url.URL
}

func (e imageOpenSaverError) Error() string {
	return fmt.Sprintf("imageOpenSaverError with URL:%v", e.url)
}

// ImageOpenSaver interface that can Open and Close images from a given backend:fs,  s3, ...
type ImageOpenSaver interface {
	Open() (image.Image, error)
	Save(img image.Image) error
}

// filesystem implementation of the ImageOpenSaver interface
type fsImageOpenSaver struct {
	URL *url.URL
}

func (s fsImageOpenSaver) Open() (image.Image, error) {
	return imaging.Open(s.URL.Path)
}

func (s fsImageOpenSaver) Save(img image.Image) error {
	return imaging.Save(img, s.URL.Path)
}

// s3 implementation of the s3ImageOpenSaver interface
type s3ImageOpenSaver struct {
	URL *url.URL
}

func (s s3ImageOpenSaver) Open() (image.Image, error) {
	conn := s3.New(awsAuth, aws.USEast)
	bucket := conn.Bucket(s.URL.Host)
	reader, err := bucket.GetReader(s.URL.Path)
	if err != nil {
		return nil, err
	}
	defer reader.Close()
	return imaging.Decode(reader)
}

func (s s3ImageOpenSaver) Save(img image.Image) error {
	var buffer bytes.Buffer
	formats := map[string]imaging.Format{
		".jpg":  imaging.JPEG,
		".jpeg": imaging.JPEG,
		".png":  imaging.PNG,
		".tif":  imaging.TIFF,
		".tiff": imaging.TIFF,
		".bmp":  imaging.BMP,
		".gif":  imaging.GIF,
	}
	ext := strings.ToLower(filepath.Ext(s.URL.Path))
	f, ok := formats[ext]
	if !ok {
		return imaging.ErrUnsupportedFormat
	}
	err := imaging.Encode(&buffer, img, f)
	if err != nil {
		log.Println("An error occured while encoding ", s.URL)
		return err
	}
	conn := s3.New(awsAuth, aws.USEast)
	bucket := conn.Bucket(s.URL.Host)

	err = bucket.Put(s.URL.Path, buffer.Bytes(), fmt.Sprintf("image/%s", imaging.JPEG), s3.PublicRead)
	if err != nil {
		log.Println("An error occured while putting on S3", s.URL)
		return err
	}
	return nil
}

// NewImageOpenSaver return the relevant implementation of ImageOpenSaver based on
// the url.Scheme
func NewImageOpenSaver(url *url.URL) (ImageOpenSaver, error) {
	switch url.Scheme {
	case "file":
		return &fsImageOpenSaver{url}, nil
	case "s3":
		return &s3ImageOpenSaver{url}, nil
	default:
		return nil, imageOpenSaverError{url}
	}
}

type rectangle struct {
	Min [2]int `json:"min"`
	Max [2]int `json:"max"`
}

func (r *rectangle) String() string {
	return fmt.Sprintf("min: %v, max: %v", r.Min, r.Max)
}

func (r *rectangle) newImageRect() image.Rectangle {
	return image.Rect(r.Min[0], r.Min[1], r.Max[0], r.Max[1])
}

type ThumbnailOpt struct {
	Rect   *rectangle `json:"rect,omitempty"`
	Width  int        `json:"width"`
	Height int        `json:"height"`
}

type ThumbnailerMessage struct {
	SrcImage  string         `json:"srcImage"`
	DstFolder string         `json:"dstFolder"`
	Opts      []ThumbnailOpt `json:"opts"`
}

func (tm *ThumbnailerMessage) thumbURL(baseName string, opt ThumbnailOpt) *url.URL {
	fURL, err := url.Parse(tm.DstFolder)
	if err != nil {
		log.Fatalln("An error occured while parsing the DstFolder", err)
	}

	if opt.Rect != nil {
		fURL.Path = filepath.Join(
			fURL.Path,
			fmt.Sprintf("%s_c%d-%d-%d-%d_s%dx%d.jpg", baseName, opt.Rect.Min[0], opt.Rect.Min[1], opt.Rect.Max[0], opt.Rect.Max[1], opt.Width, opt.Height))
	} else if opt.Width == 0 && opt.Height == 0 {
		fURL.Path = filepath.Join(fURL.Path, baseName)
	} else {
		fURL.Path = filepath.Join(fURL.Path, fmt.Sprintf("%s_s%dx%d.jpg", baseName, opt.Width, opt.Height))
	}
	return fURL
}

// Resize the src image to the biggest thumb sizes in tm.opts.
func (tm *ThumbnailerMessage) maxThumbnail(src image.Image) image.Image {
	maxW, maxH := 0, 0
	srcW := src.Bounds().Max.X
	srcH := src.Bounds().Max.Y
	for _, opt := range tm.Opts {
		dstW, dstH := opt.Width, opt.Height
		// if new width or height is 0 then preserve aspect ratio, minimum 1px
		if dstW == 0 {
			tmpW := float64(dstH) * float64(srcW) / float64(srcH)
			dstW = int(math.Max(1.0, math.Floor(tmpW+0.5)))
		}
		if dstH == 0 {
			tmpH := float64(dstW) * float64(srcH) / float64(srcW)
			dstH = int(math.Max(1.0, math.Floor(tmpH+0.5)))
		}
		if dstW > maxW {
			maxW = dstW
		}
		if dstH > maxH {
			maxH = dstH
		}
	}
	fmt.Println("thumbnail max: ", maxW, maxH, "for :", tm.Opts)
	return imaging.Resize(src, maxW, maxH, imaging.CatmullRom)
}

func (tm *ThumbnailerMessage) generateThumbnail(errorChan chan error, srcURL *url.URL, img image.Image, opt ThumbnailOpt) {
	timerStart := time.Now()
	var thumbImg *image.NRGBA
	if opt.Rect != nil {
		img = imaging.Crop(img, opt.Rect.newImageRect())
	}

	if opt.Width == 0 && opt.Height == 0 {
		thumbImg = toNRGBA(img)
	} else {
		thumbImg = imaging.Resize(img, opt.Width, opt.Height, imaging.CatmullRom)
	}

	thumbURL := tm.thumbURL(filepath.Base(srcURL.Path), opt)
	timerThumbDone := time.Now()
	log.Println("thumb :", thumbURL, " generated in : ", timerThumbDone.Sub(timerStart))

	timerSaveStart := time.Now()
	thumb, err := NewImageOpenSaver(thumbURL)
	if err != nil {
		log.Println("An error occured while creating an instance of ImageOpenSaver", err)
		errorChan <- err
		return
	}
	err = thumb.Save(thumbImg)
	if err != nil {
		log.Println("An error occured while saving the thumb", err)
		errorChan <- err
		return
	}
	errorChan <- nil
	timerEnd := time.Now()
	log.Println("thumb :", thumbURL, " saved in : ", timerEnd.Sub(timerSaveStart))
	return
}

func (tm *ThumbnailerMessage) GenerateThumbnails() error {
	sURL, err := url.Parse(tm.SrcImage)
	if err != nil {
		log.Println("An error occured while parsing the SrcImage", err)
		return err
	}
	src, err := NewImageOpenSaver(sURL)
	if err != nil {
		log.Println("An error occured while creating an instance of ImageOpenSaver", err)
		return err
	}
	img, err := src.Open()
	if err != nil {
		log.Println("An error occured while opening SrcImage", err)
		return err
	}
	// From now on we will deal with an NRGBA image
	img = toNRGBA(img)
	fmt.Println("image Bounds: ", img.Bounds())

	if len(tm.Opts) > 1 {
		// The resized image will be used to generate all the thumbs
		img = tm.maxThumbnail(img)
	}

	errorChan := make(chan error, 1)
	for _, opt := range tm.Opts {
		go tm.generateThumbnail(errorChan, sURL, img, opt)
	}

	for i := 0; i < len(tm.Opts); i++ {
		select {
		case err := <-errorChan:
			if err != nil {
				return err
			}
		}

	}
	return nil
}
<commit_msg>carry over the ext to the thumbs use mime to set the content type in S3<commit_after>package nsqthumbnailer

import (
	"bytes"
	"fmt"
	"image"
	"log"
	"math"
	"mime"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/disintegration/imaging"
	"gopkg.in/amz.v1/aws"
	"gopkg.in/amz.v1/s3"
)

var (
	awsAuth aws.Auth
)

func init() {
	awsAuth = newAwsAuth()
}

// This function used internally to convert any image type to NRGBA if needed.
// copied from `imaging`
func toNRGBA(img image.Image) *image.NRGBA {
	srcBounds := img.Bounds()
	if srcBounds.Min.X == 0 && srcBounds.Min.Y == 0 {
		if src0, ok := img.(*image.NRGBA); ok {
			return src0
		}
	}
	return imaging.Clone(img)
}

func newAwsAuth() aws.Auth {
	// Authenticate and Create an aws S3 service
	auth, err := aws.EnvAuth()
	if err != nil {
		panic(err.Error())
	}
	return auth
}

type imageOpenSaverError struct {
	url *url.URL
}

func (e imageOpenSaverError) Error() string {
	return fmt.Sprintf("imageOpenSaverError with URL:%v", e.url)
}

// ImageOpenSaver interface that can Open and Close images from a given backend:fs,  s3, ...
type ImageOpenSaver interface {
	Open() (image.Image, error)
	Save(img image.Image) error
}

// filesystem implementation of the ImageOpenSaver interface
type fsImageOpenSaver struct {
	URL *url.URL
}

func (s fsImageOpenSaver) Open() (image.Image, error) {
	return imaging.Open(s.URL.Path)
}

func (s fsImageOpenSaver) Save(img image.Image) error {
	return imaging.Save(img, s.URL.Path)
}

// s3 implementation of the s3ImageOpenSaver interface
type s3ImageOpenSaver struct {
	URL *url.URL
}

func (s s3ImageOpenSaver) Open() (image.Image, error) {
	conn := s3.New(awsAuth, aws.USEast)
	bucket := conn.Bucket(s.URL.Host)
	reader, err := bucket.GetReader(s.URL.Path)
	if err != nil {
		return nil, err
	}
	defer reader.Close()
	return imaging.Decode(reader)
}

func (s s3ImageOpenSaver) Save(img image.Image) error {
	var buffer bytes.Buffer
	formats := map[string]imaging.Format{
		".jpg":  imaging.JPEG,
		".jpeg": imaging.JPEG,
		".png":  imaging.PNG,
		".tif":  imaging.TIFF,
		".tiff": imaging.TIFF,
		".bmp":  imaging.BMP,
		".gif":  imaging.GIF,
	}
	ext := strings.ToLower(filepath.Ext(s.URL.Path))
	f, ok := formats[ext]
	if !ok {
		return imaging.ErrUnsupportedFormat
	}
	err := imaging.Encode(&buffer, img, f)
	if err != nil {
		log.Println("An error occured while encoding ", s.URL)
		return err
	}
	conn := s3.New(awsAuth, aws.USEast)
	bucket := conn.Bucket(s.URL.Host)

	err = bucket.Put(s.URL.Path, buffer.Bytes(), mime.TypeByExtension(ext), s3.PublicRead)
	if err != nil {
		log.Println("An error occured while putting on S3", s.URL)
		return err
	}
	return nil
}

// NewImageOpenSaver return the relevant implementation of ImageOpenSaver based on
// the url.Scheme
func NewImageOpenSaver(url *url.URL) (ImageOpenSaver, error) {
	switch url.Scheme {
	case "file":
		return &fsImageOpenSaver{url}, nil
	case "s3":
		return &s3ImageOpenSaver{url}, nil
	default:
		return nil, imageOpenSaverError{url}
	}
}

type rectangle struct {
	Min [2]int `json:"min"`
	Max [2]int `json:"max"`
}

func (r *rectangle) String() string {
	return fmt.Sprintf("min: %v, max: %v", r.Min, r.Max)
}

func (r *rectangle) newImageRect() image.Rectangle {
	return image.Rect(r.Min[0], r.Min[1], r.Max[0], r.Max[1])
}

type ThumbnailOpt struct {
	Rect   *rectangle `json:"rect,omitempty"`
	Width  int        `json:"width"`
	Height int        `json:"height"`
}

type ThumbnailerMessage struct {
	SrcImage  string         `json:"srcImage"`
	DstFolder string         `json:"dstFolder"`
	Opts      []ThumbnailOpt `json:"opts"`
}

func (tm *ThumbnailerMessage) thumbURL(baseName string, opt ThumbnailOpt) *url.URL {
	fURL, err := url.Parse(tm.DstFolder)
	if err != nil {
		log.Fatalln("An error occured while parsing the DstFolder", err)
	}
	// TODO (yml): I am pretty sure that we do not really want to always do this.
	ext := strings.ToLower(filepath.Ext(tm.SrcImage))
	if opt.Rect != nil {
		fURL.Path = filepath.Join(
			fURL.Path,
			fmt.Sprintf("%s_c%d-%d-%d-%d_s%dx%d%s", baseName, opt.Rect.Min[0], opt.Rect.Min[1], opt.Rect.Max[0], opt.Rect.Max[1], opt.Width, opt.Height, ext))
	} else if opt.Width == 0 && opt.Height == 0 {
		fURL.Path = filepath.Join(fURL.Path, baseName)
	} else {
		fURL.Path = filepath.Join(fURL.Path, fmt.Sprintf("%s_s%dx%d%s", baseName, opt.Width, opt.Height, ext))
	}
	return fURL
}

// Resize the src image to the biggest thumb sizes in tm.opts.
func (tm *ThumbnailerMessage) maxThumbnail(src image.Image) image.Image {
	maxW, maxH := 0, 0
	srcW := src.Bounds().Max.X
	srcH := src.Bounds().Max.Y
	for _, opt := range tm.Opts {
		dstW, dstH := opt.Width, opt.Height
		// if new width or height is 0 then preserve aspect ratio, minimum 1px
		if dstW == 0 {
			tmpW := float64(dstH) * float64(srcW) / float64(srcH)
			dstW = int(math.Max(1.0, math.Floor(tmpW+0.5)))
		}
		if dstH == 0 {
			tmpH := float64(dstW) * float64(srcH) / float64(srcW)
			dstH = int(math.Max(1.0, math.Floor(tmpH+0.5)))
		}
		if dstW > maxW {
			maxW = dstW
		}
		if dstH > maxH {
			maxH = dstH
		}
	}
	fmt.Println("thumbnail max: ", maxW, maxH, "for :", tm.Opts)
	return imaging.Resize(src, maxW, maxH, imaging.CatmullRom)
}

func (tm *ThumbnailerMessage) generateThumbnail(errorChan chan error, srcURL *url.URL, img image.Image, opt ThumbnailOpt) {
	timerStart := time.Now()
	var thumbImg *image.NRGBA
	if opt.Rect != nil {
		img = imaging.Crop(img, opt.Rect.newImageRect())
	}

	if opt.Width == 0 && opt.Height == 0 {
		thumbImg = toNRGBA(img)
	} else {
		thumbImg = imaging.Resize(img, opt.Width, opt.Height, imaging.CatmullRom)
	}

	thumbURL := tm.thumbURL(filepath.Base(srcURL.Path), opt)
	timerThumbDone := time.Now()
	log.Println("thumb :", thumbURL, " generated in : ", timerThumbDone.Sub(timerStart))

	timerSaveStart := time.Now()
	thumb, err := NewImageOpenSaver(thumbURL)
	if err != nil {
		log.Println("An error occured while creating an instance of ImageOpenSaver", err)
		errorChan <- err
		return
	}
	err = thumb.Save(thumbImg)
	if err != nil {
		log.Println("An error occured while saving the thumb", err)
		errorChan <- err
		return
	}
	errorChan <- nil
	timerEnd := time.Now()
	log.Println("thumb :", thumbURL, " saved in : ", timerEnd.Sub(timerSaveStart))
	return
}

func (tm *ThumbnailerMessage) GenerateThumbnails() error {
	sURL, err := url.Parse(tm.SrcImage)
	if err != nil {
		log.Println("An error occured while parsing the SrcImage", err)
		return err
	}
	src, err := NewImageOpenSaver(sURL)
	if err != nil {
		log.Println("An error occured while creating an instance of ImageOpenSaver", err)
		return err
	}
	img, err := src.Open()
	if err != nil {
		log.Println("An error occured while opening SrcImage", err)
		return err
	}
	// From now on we will deal with an NRGBA image
	img = toNRGBA(img)
	fmt.Println("image Bounds: ", img.Bounds())

	if len(tm.Opts) > 1 {
		// The resized image will be used to generate all the thumbs
		img = tm.maxThumbnail(img)
	}

	errorChan := make(chan error, 1)
	for _, opt := range tm.Opts {
		go tm.generateThumbnail(errorChan, sURL, img, opt)
	}

	for i := 0; i < len(tm.Opts); i++ {
		select {
		case err := <-errorChan:
			if err != nil {
				return err
			}
		}

	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"os"
	"fmt"
	"path"
	"sync"
	"regexp"
	"strings"
	"net/http"
	"encoding/json"

	"golang.org/x/net/html"

	"github.com/parnurzeal/gorequest"
)

// Configuration object
type Configuration struct {
	Directory string `json:"directory"`
}

// Track class
type Track struct {
	Title 	string 	`json:"title"`
	FileURL File 	`json:"file"`
	Album string
}

// File - Track URL helper class
type File struct{
	URL string `json:"mp3-128"`
}

const CONCURRENCY = 4
var throttle = make(chan int, CONCURRENCY)

func _contains(s []string, e string) bool {
    for _, a := range s {
        if a == e {
            return true
        }
    }
    return false
}

func downloadTrack(path string, track Track, wg *sync.WaitGroup, throttle chan int){
	defer wg.Done()

	fmt.Printf("[DEBUG] Downloading %s (%s)\n", track.Title, track.Album)

	var fileName = fmt.Sprintf("%s/%s/%s.mp3" , path, track.Album, track.Title)
	output, err := os.Create(fileName)
	if err != nil {
		fmt.Println("[ERROR] Failed creating", fileName, "-", err)
		return
	}
	defer output.Close()

	var url = "https:" + track.FileURL.URL
	response, err := http.Get(url)
	if err != nil {
		fmt.Println("[ERROR] Failed downloading", url, "-", err)
		return
	}
	defer response.Body.Close()

	_, err = io.Copy(output, response.Body)
	if err != nil {
		fmt.Println("[ERROR] Failed downloading", url, "-", err)
		return
	}

	fmt.Printf("[DEBUG] Successfully Downloaded %s (%s)\n", track.Title, track.Album)

	<-throttle
}

func fetchAlbumTracks(band string, album string) (tracks []Track) {
	var url = "https://" + band + ".bandcamp.com" + album
	_, body, errs := gorequest.New().Get(url).End()

	if errs != nil {
		fmt.Println("[ERROR] Failed to crawl \"" + url + "\"")
		return
	}

	pattern, _ := regexp.Compile(`trackinfo.+}]`)
	result := pattern.FindString(body)

	json.Unmarshal([]byte(result[strings.Index(result, "[{"):]), &tracks)

	return
}

func fetchAlbums(band string) (albums []string) {
	var url = "https://" + band + ".bandcamp.com/music"
	resp, err := http.Get(url)

	if err != nil {
		fmt.Println("[ERROR] Failed to crawl \"" + url + "\"")
		return
	}

	b := resp.Body
	defer b.Close()

	z := html.NewTokenizer(b)

	for {
		tt := z.Next()

		switch {
		case tt == html.ErrorToken:
			return
		case tt == html.StartTagToken:
			t := z.Token()

			isAnchor := t.Data == "a"
			if !isAnchor {
				continue
			}

			for _, a := range t.Attr {
				if a.Key == "href" {
					href := a.Val
					if strings.Index(href, "/album") == 0 && !_contains(albums, href){
						albums = append(albums, href)
					}
				}
			}
		}
	}
}

func getMusicPathRoot() string {
	root := path.Dir(os.Args[0]) + "/data"

	file, err := os.Open(path.Dir(os.Args[0]) + "/.trabandcamprc")
	if err == nil {
		decoder := json.NewDecoder(file)
		configuration := Configuration{}
		err = decoder.Decode(&configuration)
		if err != nil {
			fmt.Println("[ERROR] Cannot parse configuration file.", err)
			os.Exit(1)
		}

		root = configuration.Directory
	}

	return root
}

func checkBandExistence(band string) bool{
	var url = "https://" + band + ".bandcamp.com/music"
	resp, err := http.Get(url)

	if resp.StatusCode != 200 || err != nil {
		return false
	}

	return true
}

func main(){
	fmt.Println("                          _                     _")
	fmt.Println("___________              | |                   | |")
	fmt.Println("\\__    ___/___________   | |__   __ _ _ __   __| | ___ __ _ _ __ ___  _ __")
	fmt.Println("  |    |  \\_  __ \\__  \\  | '_ \\ / _` | '_ \\ / _` |/ __/ _` | '_ ` _ \\| '_ \\")
	fmt.Println("  |    |   |  | \\// __ \\_| |_) | (_| | | | | (_| | (_| (_| | | | | | | |_) |")
	fmt.Println("  |____|   |__|  (____  /|_.__/ \\__,_|_| |_|\\__,_|\\___\\__,_|_| |_| |_| .__/")
	fmt.Println("                                                                     | |")
	fmt.Println("                                                                     |_| v.0.0.2")

	if len(os.Args) != 2 {
		fmt.Println("[ERROR] Missing Band Name")
		os.Exit(1)
	}

	band := os.Args[1]
	if !checkBandExistence(band) {
		fmt.Printf("[ERROR] Band `%s` doesn't exist\n", band)
		os.Exit(1)
	}

	var musicPath = getMusicPathRoot()
	musicPath = musicPath + "/" + band
	os.MkdirAll(musicPath, 0777)

	fmt.Println("[INFO] Analyzing " + band)

	var albums = fetchAlbums(band)
	fmt.Printf("[INFO] Found %d Albums\n", len(albums))
	fmt.Printf("[DEBUG] %q\n", albums)

	var tracks []Track
	for _, album := range albums{
		var albumPath = musicPath + "/" + album[7:]
		os.MkdirAll(albumPath, 0777)
		tmpTracks := fetchAlbumTracks(band, album)

		for index := range tmpTracks{
			tmpTracks[index].Album = album[7:]
		}

		tracks = append(tracks, tmpTracks...)
	}
	fmt.Printf("[INFO] Found %d Tracks\n", len(tracks))

	var wg sync.WaitGroup
	for _, track := range tracks{
		throttle <- 1
		wg.Add(1)

		go downloadTrack(musicPath, track, &wg, throttle)
	}

	wg.Wait()
}
<commit_msg>Return the throttle chan in case of error<commit_after>package main

import (
	"io"
	"os"
	"fmt"
	"path"
	"sync"
	"regexp"
	"strings"
	"net/http"
	"encoding/json"

	"golang.org/x/net/html"

	"github.com/parnurzeal/gorequest"
)

// Configuration object
type Configuration struct {
	Directory string `json:"directory"`
}

// Track class
type Track struct {
	Title 	string 	`json:"title"`
	FileURL File 	`json:"file"`
	Album string
}

// File - Track URL helper class
type File struct{
	URL string `json:"mp3-128"`
}

const CONCURRENCY = 4
var throttle = make(chan int, CONCURRENCY)

func _contains(s []string, e string) bool {
    for _, a := range s {
        if a == e {
            return true
        }
    }
    return false
}

func downloadTrack(path string, track Track, wg *sync.WaitGroup, throttle chan int){
	defer wg.Done()

	fmt.Printf("[DEBUG] Downloading %s (%s)\n", track.Title, track.Album)

	var fileName = fmt.Sprintf("%s/%s/%s.mp3" , path, track.Album, track.Title)
	output, err := os.Create(fileName)
	if err != nil {
		fmt.Println("[ERROR] Failed creating", fileName, "-", err)
		<-throttle
		return
	}
	defer output.Close()

	var url = "https:" + track.FileURL.URL
	response, err := http.Get(url)
	if err != nil {
		fmt.Println("[ERROR] Failed downloading", url, "-", err)
		<-throttle
		return
	}
	defer response.Body.Close()

	_, err = io.Copy(output, response.Body)
	if err != nil {
		fmt.Println("[ERROR] Failed downloading", url, "-", err)
		<-throttle
		return
	}

	fmt.Printf("[DEBUG] Successfully Downloaded %s (%s)\n", track.Title, track.Album)

	<-throttle
}

func fetchAlbumTracks(band string, album string) (tracks []Track) {
	var url = "https://" + band + ".bandcamp.com" + album
	_, body, errs := gorequest.New().Get(url).End()

	if errs != nil {
		fmt.Println("[ERROR] Failed to crawl \"" + url + "\"")
		return
	}

	pattern, _ := regexp.Compile(`trackinfo.+}]`)
	result := pattern.FindString(body)

	json.Unmarshal([]byte(result[strings.Index(result, "[{"):]), &tracks)

	return
}

func fetchAlbums(band string) (albums []string) {
	var url = "https://" + band + ".bandcamp.com/music"
	resp, err := http.Get(url)

	if err != nil {
		fmt.Println("[ERROR] Failed to crawl \"" + url + "\"")
		return
	}

	b := resp.Body
	defer b.Close()

	z := html.NewTokenizer(b)

	for {
		tt := z.Next()

		switch {
		case tt == html.ErrorToken:
			return
		case tt == html.StartTagToken:
			t := z.Token()

			isAnchor := t.Data == "a"
			if !isAnchor {
				continue
			}

			for _, a := range t.Attr {
				if a.Key == "href" {
					href := a.Val
					if strings.Index(href, "/album") == 0 && !_contains(albums, href){
						albums = append(albums, href)
					}
				}
			}
		}
	}
}

func getMusicPathRoot() string {
	root := path.Dir(os.Args[0]) + "/data"

	file, err := os.Open(path.Dir(os.Args[0]) + "/.trabandcamprc")
	if err == nil {
		decoder := json.NewDecoder(file)
		configuration := Configuration{}
		err = decoder.Decode(&configuration)
		if err != nil {
			fmt.Println("[ERROR] Cannot parse configuration file.", err)
			os.Exit(1)
		}

		root = configuration.Directory
	}

	return root
}

func checkBandExistence(band string) bool{
	var url = "https://" + band + ".bandcamp.com/music"
	resp, err := http.Get(url)

	if resp.StatusCode != 200 || err != nil {
		return false
	}

	return true
}

func main(){
	fmt.Println("                          _                     _")
	fmt.Println("___________              | |                   | |")
	fmt.Println("\\__    ___/___________   | |__   __ _ _ __   __| | ___ __ _ _ __ ___  _ __")
	fmt.Println("  |    |  \\_  __ \\__  \\  | '_ \\ / _` | '_ \\ / _` |/ __/ _` | '_ ` _ \\| '_ \\")
	fmt.Println("  |    |   |  | \\// __ \\_| |_) | (_| | | | | (_| | (_| (_| | | | | | | |_) |")
	fmt.Println("  |____|   |__|  (____  /|_.__/ \\__,_|_| |_|\\__,_|\\___\\__,_|_| |_| |_| .__/")
	fmt.Println("                                                                     | |")
	fmt.Println("                                                                     |_| v.0.0.2")

	if len(os.Args) != 2 {
		fmt.Println("[ERROR] Missing Band Name")
		os.Exit(1)
	}

	band := os.Args[1]
	if !checkBandExistence(band) {
		fmt.Printf("[ERROR] Band `%s` doesn't exist\n", band)
		os.Exit(1)
	}

	var musicPath = getMusicPathRoot()
	musicPath = musicPath + "/" + band
	os.MkdirAll(musicPath, 0777)

	fmt.Println("[INFO] Analyzing " + band)

	var albums = fetchAlbums(band)
	fmt.Printf("[INFO] Found %d Albums\n", len(albums))
	fmt.Printf("[DEBUG] %q\n", albums)

	var tracks []Track
	for _, album := range albums{
		var albumPath = musicPath + "/" + album[7:]
		os.MkdirAll(albumPath, 0777)
		tmpTracks := fetchAlbumTracks(band, album)

		for index := range tmpTracks{
			tmpTracks[index].Album = album[7:]
		}

		tracks = append(tracks, tmpTracks...)
	}
	fmt.Printf("[INFO] Found %d Tracks\n", len(tracks))

	var wg sync.WaitGroup
	for _, track := range tracks{
		throttle <- 1
		wg.Add(1)

		go downloadTrack(musicPath, track, &wg, throttle)
	}

	wg.Wait()
}
<|endoftext|>


<commit_before>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"errors"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const NameKey = "name"

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

var udpConn *net.UDPConn

// used to initialize udpConn inside sendSample
var udpInitOnce sync.Once
var udpInitFunc = func() { initUDPConn(localVeneurAddress) }
var udpConnUninitializedErr = errors.New("UDP connection is not yet initialized")

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	if t.End.IsZero() {
		t.End = time.Now()
	}
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Tags:           t.Tags,
		Service:        Service,
	}

	span.Tags[NameKey] = name
	span.Tags[ResourceKey] = t.Resource

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Tags[NameKey] = name

	err := sendSample(span)
	if err != nil {
		logrus.WithError(err).Error("Error submitting sample")
	}
	return err
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
			opts = append(opts, NameTag(name))
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) error {
	if Disabled() {
		return nil
	}

	udpInitOnce.Do(udpInitFunc)
	if udpConn == nil {
		return udpConnUninitializedErr
	}

	data, err := proto.Marshal(sample)
	if err != nil {
		return err
	}

	_, err = udpConn.Write(data)
	if err != nil {
		return err
	}

	return nil
}

// initUDPConn will initialize the global UDP connection.
// It will panic if it encounters an error.
// It should only be called via the corresponding sync.Once
func initUDPConn(address string) {
	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		panic(err)
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		panic(err)
	}
	udpConn = conn
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<commit_msg>Return error if connection initialization fails<commit_after>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const NameKey = "name"

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

var udpConn *net.UDPConn

// used to initialize udpConn inside sendSample
var udpInitOnce sync.Once
var udpInitFunc = func() { initUDPConn(localVeneurAddress) }
var udpConnUninitializedErr = errors.New("UDP connection is not yet initialized")

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	if t.End.IsZero() {
		t.End = time.Now()
	}
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Tags:           t.Tags,
		Service:        Service,
	}

	span.Tags[NameKey] = name
	span.Tags[ResourceKey] = t.Resource

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Tags[NameKey] = name

	err := sendSample(span)
	if err != nil {
		logrus.WithError(err).Error("Error submitting sample")
	}
	return err
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
			opts = append(opts, NameTag(name))
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err, ok := r.(error)
			if !ok {
				err = fmt.Errorf("encountered panic while sending sample %#v", err)
			}
		}
	}()

	if Disabled() {
		return nil
	}

	udpInitOnce.Do(udpInitFunc)

	// at this point, the connection should already be initialized

	if udpConn == nil {
		return udpConnUninitializedErr
	}

	data, err := proto.Marshal(sample)
	if err != nil {
		return
	}

	_, err = udpConn.Write(data)
	if err != nil {
		return
	}

	return nil
}

// initUDPConn will initialize the global UDP connection.
// It will panic if it encounters an error.
// It should only be called via the corresponding sync.Once
func initUDPConn(address string) {
	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		panic(err)
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		panic(err)
	}
	udpConn = conn
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<|endoftext|>


<commit_before>package travel

import (
	"fmt"
	"os"

	"github.com/akerl/speculate/creds"
	"github.com/akerl/speculate/executors"
	"github.com/akerl/timber/log"

	"github.com/akerl/voyager/cartogram"
	"github.com/akerl/voyager/profiles"
	"github.com/akerl/voyager/prompt"
)

var logger = log.NewLogger("voyager")

type hop struct {
	Profile string
	Account string
	Region  string
	Role    string
	Mfa     bool
	Policy  string
}

// Itinerary describes a travel request
type Itinerary struct {
	Args        []string
	RoleName    []string
	ProfileName []string
	SessionName string
	Policy      string
	Lifetime    int64
	MfaCode     string
	MfaSerial   string
	MfaPrompt   executors.MfaPrompt
	Prompt      prompt.Func
	Store       profiles.Store
	pack        *cartogram.Pack
}

// Travel loads creds from a full set of parameters
func (i *Itinerary) Travel() (creds.Creds, error) {
	return i.getCreds()
}

func (i *Itinerary) getStore() profiles.Store {
	if i.Store == nil {
		logger.InfoMsg("Using default profile store")
		i.Store = profiles.NewDefaultStore()
	}
	return i.Store
}

func (i *Itinerary) getPrompt() prompt.Func {
	if i.Prompt == nil {
		logger.InfoMsg("Using default prompt")
		i.Prompt = prompt.WithDefault
	}
	return i.Prompt
}

func (i *Itinerary) getPack() (cartogram.Pack, error) {
	if i.pack != nil {
		return *i.pack, nil
	}
	i.pack = &cartogram.Pack{}
	err := i.pack.Load()
	return *i.pack, err
}

func (i *Itinerary) getAccount() (cartogram.Account, error) {
	pack, err := i.getPack()
	if err != nil {
		return cartogram.Account{}, err
	}
	return pack.FindWithPrompt(i.Args, i.getPrompt())
}

func (i *Itinerary) getCreds() (creds.Creds, error) {
	var c creds.Creds
	var err error

	path, err := i.getPath()
	if err != nil {
		return c, err
	}

	err = clearEnvironment()
	if err != nil {
		return c, err
	}

	profileHop, stack := path[0], path[1:]
	store := i.getStore()
	err = profiles.SetProfile(profileHop.Profile, store)
	if err != nil {
		return c, err
	}
	err = os.Setenv("AWS_DEFAULT_REGION", profileHop.Region)

	stack[len(stack)-1].Policy = i.Policy
	for _, thisHop := range stack {
		c, err = i.executeHop(thisHop, c)
		if err != nil {
			break
		}
	}
	return c, err
}

func clearEnvironment() error {
	for varName := range creds.Translations["envvar"] {
		logger.InfoMsg(fmt.Sprintf("Unsetting env var: %s", varName))
		err := os.Unsetenv(varName)
		if err != nil {
			return err
		}
	}
	return nil
}

//revive:disable-next-line:cyclomatic
func (i *Itinerary) executeHop(thisHop hop, c creds.Creds) (creds.Creds, error) {
	var newCreds creds.Creds
	logger.InfoMsg(fmt.Sprintf("Executing hop: %+v", thisHop))
	a := executors.Assumption{}
	if err := a.SetAccountID(thisHop.Account); err != nil {
		return newCreds, err
	}
	if err := a.SetRoleName(thisHop.Role); err != nil {
		return newCreds, err
	}
	if err := a.SetSessionName(i.SessionName); err != nil {
		return newCreds, err
	}
	if err := a.SetLifetime(i.Lifetime); err != nil {
		return newCreds, err
	}
	if err := a.SetPolicy(thisHop.Policy); err != nil {
		return newCreds, err
	}
	if thisHop.Mfa {
		if err := a.SetMfa(true); err != nil {
			return newCreds, err
		}
		if err := a.SetMfaSerial(i.MfaSerial); err != nil {
			return newCreds, err
		}
		if err := a.SetMfaCode(i.MfaCode); err != nil {
			return newCreds, err
		}
		if err := a.SetMfaPrompt(i.MfaPrompt); err != nil {
			return newCreds, err
		}
	}
	newCreds, err := a.ExecuteWithCreds(c)
	newCreds.Region = thisHop.Region
	return newCreds, err
}

func (i *Itinerary) getPath() ([]hop, error) { //revive:disable-line:cyclomatic
	var paths [][]hop
	mapProfiles := make(map[string]bool)
	mapRoles := make(map[string]bool)

	account, err := i.getAccount()
	if err != nil {
		return []hop{}, err
	}

	for _, r := range account.Roles {
		p, err := i.tracePath(account, r)
		if err != nil {
			return []hop{}, err
		}
		for _, item := range p {
			paths = append(paths, item)
			mapRoles[item[len(item)-1].Role] = true
		}
	}

	allRoles := keys(mapRoles)
	role, err := i.getPrompt().Filtered(i.RoleName, allRoles, "Desired target role:")
	if err != nil {
		return []hop{}, err
	}
	hopsWithMatchingRoles := [][]hop{}
	for _, item := range paths {
		if item[len(item)-1].Role == role {
			hopsWithMatchingRoles = append(hopsWithMatchingRoles, item)
			mapProfiles[item[0].Profile] = true
		}
	}

	allProfiles := keys(mapProfiles)
	profile, err := i.getPrompt().Filtered(i.ProfileName, allProfiles, "Desired target profile:")
	if err != nil {
		return []hop{}, err
	}
	hopsWithMatchingProfiles := [][]hop{}
	for _, item := range hopsWithMatchingRoles {
		if item[0].Profile == profile {
			hopsWithMatchingProfiles = append(hopsWithMatchingProfiles, item)
		}
	}

	if len(hopsWithMatchingProfiles) > 1 {
		logger.InfoMsg("Multiple valid paths detected. Selecting the first option")
	}
	return hopsWithMatchingProfiles[0], nil
}

func (i *Itinerary) tracePath(acc cartogram.Account, role cartogram.Role) ([][]hop, error) {
	var srcHops [][]hop

	for _, item := range role.Sources {
		pathMatch := cartogram.AccountRegex.FindStringSubmatch(item.Path)
		if len(pathMatch) == 3 {
			err := i.testSourcePath(pathMatch[1], pathMatch[2], &srcHops)
			if err != nil {
				return [][]hop{}, err
			}
		} else {
			srcHops = append(srcHops, []hop{{Profile: item.Path}})
		}
	}

	myHop := hop{
		Role:    role.Name,
		Account: acc.Account,
		Region:  acc.Region,
		Mfa:     role.Mfa,
	}

	for i := range srcHops {
		srcHops[i] = append(srcHops[i], myHop)
	}
	return srcHops, nil
}

func (i *Itinerary) testSourcePath(srcAccID, srcRoleName string, srcHops *[][]hop) error {
	ok, srcAcc := i.pack.Lookup(srcAccID)
	if !ok {
		logger.DebugMsg(fmt.Sprintf("Found dead end due to missing account: %s", srcAccID))
		return nil
	}
	ok, srcRole := srcAcc.Roles.Lookup(srcRoleName)
	if !ok {
		logger.DebugMsg(fmt.Sprintf(
			"Found dead end due to missing role: %s/%s", srcAccID, srcRoleName,
		))
		return nil
	}
	newPaths, err := i.tracePath(srcAcc, srcRole)
	if err != nil {
		return err
	}
	for _, np := range newPaths {
		*srcHops = append(*srcHops, np)
	}
	return nil
}

func keys(input map[string]bool) []string {
	list := []string{}
	for k := range input {
		list = append(list, k)
	}
	return list
}
<commit_msg>set profile region (#11)<commit_after>package travel

import (
	"fmt"
	"os"

	"github.com/akerl/speculate/creds"
	"github.com/akerl/speculate/executors"
	"github.com/akerl/timber/log"

	"github.com/akerl/voyager/cartogram"
	"github.com/akerl/voyager/profiles"
	"github.com/akerl/voyager/prompt"
)

var logger = log.NewLogger("voyager")

type hop struct {
	Profile string
	Account string
	Region  string
	Role    string
	Mfa     bool
	Policy  string
}

// Itinerary describes a travel request
type Itinerary struct {
	Args        []string
	RoleName    []string
	ProfileName []string
	SessionName string
	Policy      string
	Lifetime    int64
	MfaCode     string
	MfaSerial   string
	MfaPrompt   executors.MfaPrompt
	Prompt      prompt.Func
	Store       profiles.Store
	pack        *cartogram.Pack
}

// Travel loads creds from a full set of parameters
func (i *Itinerary) Travel() (creds.Creds, error) {
	return i.getCreds()
}

func (i *Itinerary) getStore() profiles.Store {
	if i.Store == nil {
		logger.InfoMsg("Using default profile store")
		i.Store = profiles.NewDefaultStore()
	}
	return i.Store
}

func (i *Itinerary) getPrompt() prompt.Func {
	if i.Prompt == nil {
		logger.InfoMsg("Using default prompt")
		i.Prompt = prompt.WithDefault
	}
	return i.Prompt
}

func (i *Itinerary) getPack() (cartogram.Pack, error) {
	if i.pack != nil {
		return *i.pack, nil
	}
	i.pack = &cartogram.Pack{}
	err := i.pack.Load()
	return *i.pack, err
}

func (i *Itinerary) getAccount() (cartogram.Account, error) {
	pack, err := i.getPack()
	if err != nil {
		return cartogram.Account{}, err
	}
	return pack.FindWithPrompt(i.Args, i.getPrompt())
}

func (i *Itinerary) getCreds() (creds.Creds, error) {
	var c creds.Creds
	var err error

	path, err := i.getPath()
	if err != nil {
		return c, err
	}

	err = clearEnvironment()
	if err != nil {
		return c, err
	}

	profileHop, stack := path[0], path[1:]
	logger.InfoMsg(fmt.Sprintf("Setting origin hop: %+v", profileHop))
	store := i.getStore()
	err = profiles.SetProfile(profileHop.Profile, store)
	if err != nil {
		return c, err
	}
	logger.InfoMsg(fmt.Sprintf("Setting AWS_DEFAULT_REGION to %s", profileHop.Region))
	err = os.Setenv("AWS_DEFAULT_REGION", profileHop.Region)
	if err != nil {
		return c, err
	}

	stack[len(stack)-1].Policy = i.Policy
	for _, thisHop := range stack {
		c, err = i.executeHop(thisHop, c)
		if err != nil {
			break
		}
	}
	return c, err
}

func clearEnvironment() error {
	for varName := range creds.Translations["envvar"] {
		logger.InfoMsg(fmt.Sprintf("Unsetting env var: %s", varName))
		err := os.Unsetenv(varName)
		if err != nil {
			return err
		}
	}
	return nil
}

//revive:disable-next-line:cyclomatic
func (i *Itinerary) executeHop(thisHop hop, c creds.Creds) (creds.Creds, error) {
	var newCreds creds.Creds
	logger.InfoMsg(fmt.Sprintf("Executing hop: %+v", thisHop))
	a := executors.Assumption{}
	if err := a.SetAccountID(thisHop.Account); err != nil {
		return newCreds, err
	}
	if err := a.SetRoleName(thisHop.Role); err != nil {
		return newCreds, err
	}
	if err := a.SetSessionName(i.SessionName); err != nil {
		return newCreds, err
	}
	if err := a.SetLifetime(i.Lifetime); err != nil {
		return newCreds, err
	}
	if err := a.SetPolicy(thisHop.Policy); err != nil {
		return newCreds, err
	}
	if thisHop.Mfa {
		if err := a.SetMfa(true); err != nil {
			return newCreds, err
		}
		if err := a.SetMfaSerial(i.MfaSerial); err != nil {
			return newCreds, err
		}
		if err := a.SetMfaCode(i.MfaCode); err != nil {
			return newCreds, err
		}
		if err := a.SetMfaPrompt(i.MfaPrompt); err != nil {
			return newCreds, err
		}
	}
	newCreds, err := a.ExecuteWithCreds(c)
	newCreds.Region = thisHop.Region
	return newCreds, err
}

func (i *Itinerary) getPath() ([]hop, error) { //revive:disable-line:cyclomatic
	var paths [][]hop
	mapProfiles := make(map[string]bool)
	mapRoles := make(map[string]bool)

	account, err := i.getAccount()
	if err != nil {
		return []hop{}, err
	}

	for _, r := range account.Roles {
		p, err := i.tracePath(account, r)
		if err != nil {
			return []hop{}, err
		}
		for _, item := range p {
			paths = append(paths, item)
			mapRoles[item[len(item)-1].Role] = true
		}
	}

	allRoles := keys(mapRoles)
	role, err := i.getPrompt().Filtered(i.RoleName, allRoles, "Desired target role:")
	if err != nil {
		return []hop{}, err
	}
	hopsWithMatchingRoles := [][]hop{}
	for _, item := range paths {
		if item[len(item)-1].Role == role {
			hopsWithMatchingRoles = append(hopsWithMatchingRoles, item)
			mapProfiles[item[0].Profile] = true
		}
	}

	allProfiles := keys(mapProfiles)
	profile, err := i.getPrompt().Filtered(i.ProfileName, allProfiles, "Desired target profile:")
	if err != nil {
		return []hop{}, err
	}
	hopsWithMatchingProfiles := [][]hop{}
	for _, item := range hopsWithMatchingRoles {
		if item[0].Profile == profile {
			hopsWithMatchingProfiles = append(hopsWithMatchingProfiles, item)
		}
	}

	if len(hopsWithMatchingProfiles) > 1 {
		logger.InfoMsg("Multiple valid paths detected. Selecting the first option")
	}
	return hopsWithMatchingProfiles[0], nil
}

func (i *Itinerary) tracePath(acc cartogram.Account, role cartogram.Role) ([][]hop, error) {
	var srcHops [][]hop

	for _, item := range role.Sources {
		pathMatch := cartogram.AccountRegex.FindStringSubmatch(item.Path)
		if len(pathMatch) == 3 {
			err := i.testSourcePath(pathMatch[1], pathMatch[2], &srcHops)
			if err != nil {
				return [][]hop{}, err
			}
		} else {
			srcHops = append(srcHops, []hop{{
				Profile: item.Path,
				Region:  acc.Region,
			}})
		}
	}

	myHop := hop{
		Role:    role.Name,
		Account: acc.Account,
		Region:  acc.Region,
		Mfa:     role.Mfa,
	}

	for i := range srcHops {
		srcHops[i] = append(srcHops[i], myHop)
	}
	return srcHops, nil
}

func (i *Itinerary) testSourcePath(srcAccID, srcRoleName string, srcHops *[][]hop) error {
	ok, srcAcc := i.pack.Lookup(srcAccID)
	if !ok {
		logger.DebugMsg(fmt.Sprintf("Found dead end due to missing account: %s", srcAccID))
		return nil
	}
	ok, srcRole := srcAcc.Roles.Lookup(srcRoleName)
	if !ok {
		logger.DebugMsg(fmt.Sprintf(
			"Found dead end due to missing role: %s/%s", srcAccID, srcRoleName,
		))
		return nil
	}
	newPaths, err := i.tracePath(srcAcc, srcRole)
	if err != nil {
		return err
	}
	for _, np := range newPaths {
		*srcHops = append(*srcHops, np)
	}
	return nil
}

func keys(input map[string]bool) []string {
	list := []string{}
	for k := range input {
		list = append(list, k)
	}
	return list
}
<|endoftext|>


<commit_before>package tspec

import (
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"strings"
	"sync"

	"github.com/go-openapi/spec"
	"github.com/pkg/errors"
)

// Parser ...
type Parser struct {
	fset       *token.FileSet
	fileMap    map[string]*ast.File
	dirPkgMap  map[string]*ast.Package
	pkgObjsMap map[*ast.Package]map[string]*ast.Object
	typeMap    map[string]*spec.Schema
}

// NewParser ...
func NewParser() (parser *Parser) {
	parser = new(Parser)
	parser.fset = token.NewFileSet()
	parser.fileMap = make(map[string]*ast.File)
	parser.dirPkgMap = make(map[string]*ast.Package)
	parser.pkgObjsMap = make(map[*ast.Package]map[string]*ast.Object)
	parser.typeMap = make(map[string]*spec.Schema)
	return
}

// ParseFile ...
func (t *Parser) ParseFile(filePath string) (f *ast.File, err error) {
	if tmpF, ok := t.fileMap[filePath]; ok {
		f = tmpF
		return
	}
	f, err = parser.ParseFile(t.fset, filePath, nil, parser.ParseComments)
	if err != nil {
		err = errors.WithStack(err)
		return
	}

	t.fileMap[filePath] = f
	return
}

// ParseDir ...
func (t *Parser) ParseDir(dirPath string, pkgName string) (pkg *ast.Package, err error) {
	if tmpPkg, ok := t.dirPkgMap[dirPath]; ok {
		pkg = tmpPkg
		return
	}

	pkgs, err := parser.ParseDir(t.fset, dirPath, nil, parser.ParseComments)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	for k := range pkgs {
		if k == pkgName {
			pkg = pkgs[k]
			break
		}
	}
	if pkg == nil {
		err = errors.Errorf("%s not found in %s", pkgName, dirPath)
		return
	}

	t.dirPkgMap[dirPath] = pkg
	return
}

// Import ...
func (t *Parser) Import(ispec *ast.ImportSpec) (pkg *ast.Package, err error) {
	pkgPath := strings.Trim(ispec.Path.Value, "\"")

	wd, err := os.Getwd()
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	importPkg, err := build.Import(pkgPath, wd, build.ImportComment)
	if err != nil {
		err = errors.WithStack(err)
		return
	}

	pkg, err = t.ParseDir(importPkg.Dir, importPkg.Name)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	return
}

// ParsePkg ...
func (t *Parser) ParsePkg(pkg *ast.Package) (objs map[string]*ast.Object, err error) {
	if tmpObjs, ok := t.pkgObjsMap[pkg]; ok {
		objs = tmpObjs
		return
	}

	objs = make(map[string]*ast.Object)
	for _, f := range pkg.Files {
		for key, obj := range f.Scope.Objects {
			if obj.Kind == ast.Typ {
				objs[key] = obj
			}
		}
	}

	t.pkgObjsMap[pkg] = objs
	return
}

func (t *Parser) parseTypeStr(oPkg *ast.Package, typeStr string) (pkg *ast.Package,
	obj *ast.Object, err error) {
	var objs map[string]*ast.Object
	var ok bool

	if !strings.Contains(typeStr, ".") {
		pkg = oPkg
		objs, err = t.ParsePkg(pkg)
		if err != nil {
			err = errors.WithStack(err)
			return
		}
		if obj, ok = objs[typeStr]; !ok {
			err = errors.Errorf("%s not found in package %s", typeStr, pkg.Name)
			return
		}
		return
	}

	strs := strings.Split(typeStr, ".")
	if len(strs) != 2 {
		err = errors.Errorf("invalid type str %s", typeStr)
		return
	}
	pkgName := strs[0]
	typeTitle := strs[1]

	var p *ast.Package
	for _, file := range oPkg.Files {
		for _, ispec := range file.Imports {
			p, err = t.Import(ispec)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
			if !(pkgName == p.Name || (ispec.Name != nil && pkgName == ispec.Name.Name)) {
				continue
			}
			objs, err = t.ParsePkg(p)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
			if _, ok = objs[typeTitle]; !ok {
				continue
			} else {
				pkg = p
				obj = objs[typeTitle]
				break
			}
		}
	}
	if pkg == nil || obj == nil {
		err = errors.Errorf("%s.%s not found", pkgName, typeTitle)
		return
	}

	return
}

func (t *Parser) parseIdentExpr(oExpr ast.Expr, pkg *ast.Package) (expr ast.Expr, err error) {
	expr = starExprX(oExpr)
	if ident, ok := expr.(*ast.Ident); ok {
		if ident.Obj != nil {
			ts, e := objDeclTypeSpec(ident.Obj)
			if e != nil {
				err = errors.WithStack(e)
				return
			}
			expr = starExprX(ts.Type)
		} else {
			// try to find obj in pkg
			objs, e := t.ParsePkg(pkg)
			if e != nil {
				err = errors.WithStack(e)
				return
			}
			if obj, ok := objs[ident.Name]; ok {
				ts, e := objDeclTypeSpec(obj)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				expr = starExprX(ts.Type)
			}
		}
	}
	return
}

func (t *Parser) parseTypeRef(pkg *ast.Package, expr ast.Expr, typeTitle, typeID string) (
	schema *spec.Schema, err error) {
	ident, isIdent := starExprX(expr).(*ast.Ident)
	typeExpr, err := t.parseIdentExpr(expr, pkg)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	switch typ := typeExpr.(type) {
	case *ast.StructType:
		if isIdent {
			typeID = pkg.Name + "." + ident.Name
			typeTitle = ident.Name
		}
		schema = spec.RefProperty("#/definitions/" + typeID)
		_, err = t.parseType(pkg, typ, typeTitle, typeID)
		if err != nil {
			err = errors.WithStack(err)
			return
		}
		return
	case *ast.SelectorExpr:
		typeStr, e := selectorExprTypeStr(typ)
		if e != nil {
			err = errors.WithStack(err)
			return
		}
		if typeStr != "time.Time" {
			typeID := typeStr
			schema = spec.RefProperty("#/definitions/" + typeID)
			_, err = t.Parse(pkg, typeStr)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
			return
		}
	}
	return t.parseType(pkg, typeExpr, typeTitle, typeID)
}

var parseTypeLock sync.Mutex

func (t *Parser) parseType(pkg *ast.Package, expr ast.Expr, typeTitle, typeID string) (schema *spec.Schema,
	err error) {
	parseTypeLock.Lock()
	if tmpSchema, ok := t.typeMap[typeID]; ok {
		schema = tmpSchema
		parseTypeLock.Unlock()
		return
	}
	if typeID != "" {
		t.typeMap[typeID] = nil
	}
	parseTypeLock.Unlock()

	// parse ident expr
	expr, err = t.parseIdentExpr(expr, pkg)
	if err != nil {
		err = errors.WithStack(err)
		return
	}

	schema = new(spec.Schema)
	schema.WithID(typeID)
	schema.WithTitle(typeTitle)
	switch typ := expr.(type) {
	case *ast.StructType:
		schema.Typed("object", "")
		if typ.Fields.List == nil {
			break
		}
		for _, field := range typ.Fields.List {
			if len(field.Names) != 0 {
				fieldName := field.Names[0].Name
				var fTypeID, fTypeTitle string
				if _, isAnonymousStruct := starExprX(field.Type).(*ast.StructType); isAnonymousStruct {
					fTypeTitle = typeTitle + "_" + fieldName
					fTypeID = pkg.Name + "." + fTypeTitle
				}
				prop, e := t.parseTypeRef(pkg, field.Type, fTypeTitle, fTypeID)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				schema.SetProperty(fieldName, *prop)
			} else {
				// inherited struct
				var fieldTypeTitle, fieldTypeID string
				ident, isIdent := starExprX(field.Type).(*ast.Ident)
				fieldExpr, e := t.parseIdentExpr(field.Type, pkg)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				switch fieldTyp := fieldExpr.(type) {
				case *ast.StructType:
					if isIdent {
						fieldTypeID = pkg.Name + "." + ident.Name
						fieldTypeTitle = ident.Name
					}
				case *ast.SelectorExpr:
					typeStr, e := selectorExprTypeStr(fieldTyp)
					if e != nil {
						err = errors.WithStack(err)
						return
					}
					if typeStr != "time.Time" {
						fieldTypeID = typeStr
						fieldTypeTitle = typeStr
					}
				}
				inheritedSchema, e := t.parseType(pkg, field.Type, fieldTypeTitle,
					fieldTypeID)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				for propKey, propSchema := range inheritedSchema.Properties {
					if _, ok := schema.Properties[propKey]; !ok {
						schema.SetProperty(propKey, propSchema)
					}
				}
			}
		}
	case *ast.ArrayType:
		itemsSchema, e := t.parseTypeRef(pkg, typ.Elt, "", "")
		if e != nil {
			err = errors.WithStack(e)
			return
		}
		schema = spec.ArrayProperty(itemsSchema)
	case *ast.MapType:
		valueSchema, e := t.parseTypeRef(pkg, typ.Value, "", "")
		if e != nil {
			err = errors.WithStack(e)
			return
		}
		schema = spec.MapProperty(valueSchema)
	case *ast.SelectorExpr:
		typeStr, e := selectorExprTypeStr(typ)
		if e != nil {
			err = errors.WithStack(err)
			return
		}
		if typeStr == "time.Time" {
			typeType, typeFormat, e := parseBasicType("time")
			if e != nil {
				err = errors.WithStack(e)
				return
			}
			schema.Typed(typeType, typeFormat)
		} else {
			schema, err = t.Parse(pkg, typeStr)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
		}
	case *ast.Ident: // basic type only
		typeType, typeFormat, e := parseBasicType(typ.Name)
		if e != nil {
			err = errors.WithStack(e)
			return
		}
		schema.Typed(typeType, typeFormat)
	case *ast.InterfaceType:
	default:
		err = errors.Errorf("invalid expr type %T", typ)
		return
	}

	if typeID != "" {
		t.typeMap[typeID] = schema
	}
	return
}

// Parse ...
func (t *Parser) Parse(oPkg *ast.Package, typeStr string) (
	schema *spec.Schema, err error) {
	pkg, obj, err := t.parseTypeStr(oPkg, typeStr)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	ts, err := objDeclTypeSpec(obj)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	schema, err = t.parseType(pkg, ts.Type, obj.Name, pkg.Name+"."+obj.Name)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	return
}

// Definitions ...
func (t *Parser) Definitions() (defs spec.Definitions) {
	defs = make(spec.Definitions)
	for k, v := range t.typeMap {
		defs[k] = *v
	}
	return
}

func starExprX(expr ast.Expr) ast.Expr {
	if star, ok := expr.(*ast.StarExpr); ok {
		return star.X
	}
	return expr
}

func objDeclTypeSpec(obj *ast.Object) (ts *ast.TypeSpec, err error) {
	ts, ok := obj.Decl.(*ast.TypeSpec)
	if !ok {
		err = errors.Errorf("invalid object decl, want *ast.TypeSpec, got %T", ts)
		return
	}
	return
}

func selectorExprTypeStr(expr *ast.SelectorExpr) (typeStr string, err error) {
	xIdent, ok := expr.X.(*ast.Ident)
	if !ok {
		err = errors.Errorf("invalid selector expr %#v", expr)
		return
	}
	typeStr = xIdent.Name + "." + expr.Sel.Name
	return
}

var basicTypes = map[string]string{
	"bool": "boolean:",
	"uint": "integer:uint", "uint8": "integer:uint8", "uint16": "integer:uint16",
	"uint32": "integer:uint32", "uint64": "integer:uint64",
	"int": "integer:int", "int8": "integer:int8", "int16": "integer:int16",
	"int32": "integer:int32", "int64": "integer:int64",
	"uintptr": "integer:int64",
	"float32": "number:float32", "float64": "number:float64",
	"string":    "string",
	"complex64": "number:float", "complex128": "number:double",
	"byte": "string:byte", "rune": "string:byte", "time": "string:date-time",
}

func parseBasicType(typeTitle string) (typ, format string, err error) {
	typeStr, ok := basicTypes[typeTitle]
	if !ok {
		err = errors.Errorf("invalid ident %s", typeTitle)
		return
	}
	exprs := strings.Split(typeStr, ":")
	typ = exprs[0]
	if len(exprs) > 1 {
		format = exprs[1]
	}
	return
}
<commit_msg>[tspec] Fix 'parseTypeStr'<commit_after>package tspec

import (
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"strings"
	"sync"

	"github.com/go-openapi/spec"
	"github.com/pkg/errors"
)

// Parser ...
type Parser struct {
	fset       *token.FileSet
	fileMap    map[string]*ast.File
	dirPkgMap  map[string]*ast.Package
	pkgObjsMap map[*ast.Package]map[string]*ast.Object
	typeMap    map[string]*spec.Schema
}

// NewParser ...
func NewParser() (parser *Parser) {
	parser = new(Parser)
	parser.fset = token.NewFileSet()
	parser.fileMap = make(map[string]*ast.File)
	parser.dirPkgMap = make(map[string]*ast.Package)
	parser.pkgObjsMap = make(map[*ast.Package]map[string]*ast.Object)
	parser.typeMap = make(map[string]*spec.Schema)
	return
}

// ParseFile ...
func (t *Parser) ParseFile(filePath string) (f *ast.File, err error) {
	if tmpF, ok := t.fileMap[filePath]; ok {
		f = tmpF
		return
	}
	f, err = parser.ParseFile(t.fset, filePath, nil, parser.ParseComments)
	if err != nil {
		err = errors.WithStack(err)
		return
	}

	t.fileMap[filePath] = f
	return
}

// ParseDir ...
func (t *Parser) ParseDir(dirPath string, pkgName string) (pkg *ast.Package, err error) {
	if tmpPkg, ok := t.dirPkgMap[dirPath]; ok {
		pkg = tmpPkg
		return
	}

	pkgs, err := parser.ParseDir(t.fset, dirPath, nil, parser.ParseComments)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	for k := range pkgs {
		if k == pkgName {
			pkg = pkgs[k]
			break
		}
	}
	if pkg == nil {
		err = errors.Errorf("%s not found in %s", pkgName, dirPath)
		return
	}

	t.dirPkgMap[dirPath] = pkg
	return
}

// Import ...
func (t *Parser) Import(ispec *ast.ImportSpec) (pkg *ast.Package, err error) {
	pkgPath := strings.Trim(ispec.Path.Value, "\"")

	wd, err := os.Getwd()
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	importPkg, err := build.Import(pkgPath, wd, build.ImportComment)
	if err != nil {
		err = errors.WithStack(err)
		return
	}

	pkg, err = t.ParseDir(importPkg.Dir, importPkg.Name)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	return
}

// ParsePkg ...
func (t *Parser) ParsePkg(pkg *ast.Package) (objs map[string]*ast.Object, err error) {
	if tmpObjs, ok := t.pkgObjsMap[pkg]; ok {
		objs = tmpObjs
		return
	}

	objs = make(map[string]*ast.Object)
	for _, f := range pkg.Files {
		for key, obj := range f.Scope.Objects {
			if obj.Kind == ast.Typ {
				objs[key] = obj
			}
		}
	}

	t.pkgObjsMap[pkg] = objs
	return
}

func (t *Parser) parseTypeStr(oPkg *ast.Package, typeStr string) (pkg *ast.Package,
	obj *ast.Object, err error) {
	var objs map[string]*ast.Object
	var pkgName, typeTitle string
	var ok bool

	strs := strings.Split(typeStr, ".")
	l := len(strs)
	if l == 0 || l > 2 {
		err = errors.Errorf("invalid type str %s", typeStr)
		return
	}
	if l == 1 {
		pkgName = oPkg.Name
		typeTitle = strs[0]
	} else {
		pkgName = strs[0]
		typeTitle = strs[1]
	}
	if pkgName == oPkg.Name {
		pkg = oPkg
		objs, err = t.ParsePkg(pkg)
		if err != nil {
			err = errors.WithStack(err)
			return
		}
		if obj, ok = objs[typeTitle]; !ok {
			err = errors.Errorf("%s not found in package %s", typeTitle, pkg.Name)
			return
		}
		return
	}
	var p *ast.Package
	for _, file := range oPkg.Files {
		for _, ispec := range file.Imports {
			p, err = t.Import(ispec)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
			if !(pkgName == p.Name || (ispec.Name != nil && pkgName == ispec.Name.Name)) {
				continue
			}
			objs, err = t.ParsePkg(p)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
			if _, ok = objs[typeTitle]; !ok {
				continue
			} else {
				pkg = p
				obj = objs[typeTitle]
				break
			}
		}
	}
	if pkg == nil || obj == nil {
		err = errors.Errorf("%s.%s not found", pkgName, typeTitle)
		return
	}

	return
}

func (t *Parser) parseIdentExpr(oExpr ast.Expr, pkg *ast.Package) (expr ast.Expr, err error) {
	expr = starExprX(oExpr)
	if ident, ok := expr.(*ast.Ident); ok {
		if ident.Obj != nil {
			ts, e := objDeclTypeSpec(ident.Obj)
			if e != nil {
				err = errors.WithStack(e)
				return
			}
			expr = starExprX(ts.Type)
		} else {
			// try to find obj in pkg
			objs, e := t.ParsePkg(pkg)
			if e != nil {
				err = errors.WithStack(e)
				return
			}
			if obj, ok := objs[ident.Name]; ok {
				ts, e := objDeclTypeSpec(obj)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				expr = starExprX(ts.Type)
			}
		}
	}
	return
}

func (t *Parser) parseTypeRef(pkg *ast.Package, expr ast.Expr, typeTitle, typeID string) (
	schema *spec.Schema, err error) {
	ident, isIdent := starExprX(expr).(*ast.Ident)
	typeExpr, err := t.parseIdentExpr(expr, pkg)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	switch typ := typeExpr.(type) {
	case *ast.StructType:
		if isIdent {
			typeID = pkg.Name + "." + ident.Name
			typeTitle = ident.Name
		}
		schema = spec.RefProperty("#/definitions/" + typeID)
		_, err = t.parseType(pkg, typ, typeTitle, typeID)
		if err != nil {
			err = errors.WithStack(err)
			return
		}
		return
	case *ast.SelectorExpr:
		typeStr, e := selectorExprTypeStr(typ)
		if e != nil {
			err = errors.WithStack(err)
			return
		}
		if typeStr != "time.Time" {
			typeID := typeStr
			schema = spec.RefProperty("#/definitions/" + typeID)
			_, err = t.Parse(pkg, typeStr)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
			return
		}
	}
	return t.parseType(pkg, typeExpr, typeTitle, typeID)
}

var parseTypeLock sync.Mutex

func (t *Parser) parseType(pkg *ast.Package, expr ast.Expr, typeTitle, typeID string) (schema *spec.Schema,
	err error) {
	parseTypeLock.Lock()
	if tmpSchema, ok := t.typeMap[typeID]; ok {
		schema = tmpSchema
		parseTypeLock.Unlock()
		return
	}
	if typeID != "" {
		t.typeMap[typeID] = nil
	}
	parseTypeLock.Unlock()

	// parse ident expr
	expr, err = t.parseIdentExpr(expr, pkg)
	if err != nil {
		err = errors.WithStack(err)
		return
	}

	schema = new(spec.Schema)
	schema.WithID(typeID)
	schema.WithTitle(typeTitle)
	switch typ := expr.(type) {
	case *ast.StructType:
		schema.Typed("object", "")
		if typ.Fields.List == nil {
			break
		}
		for _, field := range typ.Fields.List {
			if len(field.Names) != 0 {
				fieldName := field.Names[0].Name
				var fTypeID, fTypeTitle string
				if _, isAnonymousStruct := starExprX(field.Type).(*ast.StructType); isAnonymousStruct {
					fTypeTitle = typeTitle + "_" + fieldName
					fTypeID = pkg.Name + "." + fTypeTitle
				}
				prop, e := t.parseTypeRef(pkg, field.Type, fTypeTitle, fTypeID)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				schema.SetProperty(fieldName, *prop)
			} else {
				// inherited struct
				var fieldTypeTitle, fieldTypeID string
				ident, isIdent := starExprX(field.Type).(*ast.Ident)
				fieldExpr, e := t.parseIdentExpr(field.Type, pkg)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				switch fieldTyp := fieldExpr.(type) {
				case *ast.StructType:
					if isIdent {
						fieldTypeID = pkg.Name + "." + ident.Name
						fieldTypeTitle = ident.Name
					}
				case *ast.SelectorExpr:
					typeStr, e := selectorExprTypeStr(fieldTyp)
					if e != nil {
						err = errors.WithStack(err)
						return
					}
					if typeStr != "time.Time" {
						fieldTypeID = typeStr
						fieldTypeTitle = typeStr
					}
				}
				inheritedSchema, e := t.parseType(pkg, field.Type, fieldTypeTitle,
					fieldTypeID)
				if e != nil {
					err = errors.WithStack(e)
					return
				}
				for propKey, propSchema := range inheritedSchema.Properties {
					if _, ok := schema.Properties[propKey]; !ok {
						schema.SetProperty(propKey, propSchema)
					}
				}
			}
		}
	case *ast.ArrayType:
		itemsSchema, e := t.parseTypeRef(pkg, typ.Elt, "", "")
		if e != nil {
			err = errors.WithStack(e)
			return
		}
		schema = spec.ArrayProperty(itemsSchema)
	case *ast.MapType:
		valueSchema, e := t.parseTypeRef(pkg, typ.Value, "", "")
		if e != nil {
			err = errors.WithStack(e)
			return
		}
		schema = spec.MapProperty(valueSchema)
	case *ast.SelectorExpr:
		typeStr, e := selectorExprTypeStr(typ)
		if e != nil {
			err = errors.WithStack(err)
			return
		}
		if typeStr == "time.Time" {
			typeType, typeFormat, e := parseBasicType("time")
			if e != nil {
				err = errors.WithStack(e)
				return
			}
			schema.Typed(typeType, typeFormat)
		} else {
			schema, err = t.Parse(pkg, typeStr)
			if err != nil {
				err = errors.WithStack(err)
				return
			}
		}
	case *ast.Ident: // basic type only
		typeType, typeFormat, e := parseBasicType(typ.Name)
		if e != nil {
			err = errors.WithStack(e)
			return
		}
		schema.Typed(typeType, typeFormat)
	case *ast.InterfaceType:
	default:
		err = errors.Errorf("invalid expr type %T", typ)
		return
	}

	if typeID != "" {
		t.typeMap[typeID] = schema
	}
	return
}

// Parse ...
func (t *Parser) Parse(oPkg *ast.Package, typeStr string) (
	schema *spec.Schema, err error) {
	pkg, obj, err := t.parseTypeStr(oPkg, typeStr)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	ts, err := objDeclTypeSpec(obj)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	schema, err = t.parseType(pkg, ts.Type, obj.Name, pkg.Name+"."+obj.Name)
	if err != nil {
		err = errors.WithStack(err)
		return
	}
	return
}

// Definitions ...
func (t *Parser) Definitions() (defs spec.Definitions) {
	defs = make(spec.Definitions)
	for k, v := range t.typeMap {
		defs[k] = *v
	}
	return
}

func starExprX(expr ast.Expr) ast.Expr {
	if star, ok := expr.(*ast.StarExpr); ok {
		return star.X
	}
	return expr
}

func objDeclTypeSpec(obj *ast.Object) (ts *ast.TypeSpec, err error) {
	ts, ok := obj.Decl.(*ast.TypeSpec)
	if !ok {
		err = errors.Errorf("invalid object decl, want *ast.TypeSpec, got %T", ts)
		return
	}
	return
}

func selectorExprTypeStr(expr *ast.SelectorExpr) (typeStr string, err error) {
	xIdent, ok := expr.X.(*ast.Ident)
	if !ok {
		err = errors.Errorf("invalid selector expr %#v", expr)
		return
	}
	typeStr = xIdent.Name + "." + expr.Sel.Name
	return
}

var basicTypes = map[string]string{
	"bool": "boolean:",
	"uint": "integer:uint", "uint8": "integer:uint8", "uint16": "integer:uint16",
	"uint32": "integer:uint32", "uint64": "integer:uint64",
	"int": "integer:int", "int8": "integer:int8", "int16": "integer:int16",
	"int32": "integer:int32", "int64": "integer:int64",
	"uintptr": "integer:int64",
	"float32": "number:float32", "float64": "number:float64",
	"string":    "string",
	"complex64": "number:float", "complex128": "number:double",
	"byte": "string:byte", "rune": "string:byte", "time": "string:date-time",
}

func parseBasicType(typeTitle string) (typ, format string, err error) {
	typeStr, ok := basicTypes[typeTitle]
	if !ok {
		err = errors.Errorf("invalid ident %s", typeTitle)
		return
	}
	exprs := strings.Split(typeStr, ":")
	typ = exprs[0]
	if len(exprs) > 1 {
		format = exprs[1]
	}
	return
}
<|endoftext|>