<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strconv"

	"flag"
	"github.com/syndtr/goleveldb/leveldb"
)

var (
	// go的正则在匹配ReadAll出来的数据时有点奇怪，必须用[\w\W]{0,}?来代替\n

	// 最大同时开启线程数
	maxThread int
	// 用于获取版块名称
	getForumName = regexp.MustCompile(`<h1 class="xs2">[\w\W]{0,}?<a[^>]*>([^<]*)</a>`)
	// 用于获取版块帖子分页数量
	getForumPageNumber = regexp.MustCompile(`<a href="[^"]+" class="last">\.\.\. ([0-9]+)</a>`)
	// 用于获取版块介绍
	getForumIntroduction = regexp.MustCompile(`<div id="forum_rules_[0-9]*"[^>]*>([\w\W]{0,}?)(?:</div>[\w\W]{0,}?){3}<div class="drag">`)
	// 用于获取帖子信息
	getPostInfo = regexp.MustCompile(`id="normalthread_[0-9]+"[\w\W]{0,}?class="icn"[\w\W]{0,}?href="([^"]*)"[\w\W]{0,}?\n[\w\W]{0,}?>([^<]*)[\w\W]{0,}?class="xst" >([^<]*)[\w\W]{0,}?[\w\W]{0,}?>([^<]*)[\w\W]{0,}?(?:)?([^<]*)`)
	// 用于获取帖子内容
	getPostBody = regexp.MustCompile(`<div class="pcb">([\w\W]{0,}?)<div id="comment_[0-9]*" class="cm">`)
)

func main() {
	flag.IntVar(&maxThread, "m", 10, "最大线程数量")
	flag.Parse()

	if len(flag.Args()) == 0 {
		fmt.Println("请在命令后添加论坛版块FID参数")
		fmt.Println("例如：")
		fmt.Println(os.Args[0] + " 45 -m=10")
		fmt.Println("其中m为可选参数，用于设置最大线程数量")
		fmt.Println("FID可从版块URL中寻找")
		fmt.Println("数据会保存在db/FID路径中")
		return
	}

	printInfo("最大线程数量", maxThread)

	// 论坛版块fid
	fid := flag.Arg(0)
	printInfo("版块FID", fid)

	resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=1")
	if err != nil {
		printError("http.Get", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		printError("ReadAll", err)
		os.Exit(1)
	}

	if bytes.Contains(body, []byte(`<p>抱歉，指定的版块不存在</p>`)) {
		printError("GetForum", "版块不存在，请检查PID是否正确")
		os.Exit(1)
	}

	// 获取版块名称
	n := getForumName.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumName", "获取版块名称出错")
		os.Exit(1)
	}
	forumName := string(n[1])
	printInfo("版块名称", forumName)

	// 获取版块帖子分页数量
	n = getForumPageNumber.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumPageNumber", "获取版块分页数量出错")
		os.Exit(1)
	}
	maxPagesNum, err := strconv.Atoi(string(n[1]))
	if err != nil {
		printError("Atoi", err)
		os.Exit(1)
	}
	printInfo("本版块全部分页数量", maxPagesNum)

	// 获取版块介绍
	var forumIntroduction string
	n = getForumIntroduction.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumIntroduction", "获取版块介绍出错，版块介绍为空")
	} else {
		forumIntroduction = string(n[1])
	}

	// 创建数据库
	db, err := leveldb.OpenFile("db/"+fid, nil)
	if err != nil {
		printError("OpenDB", err)
		os.Exit(1)
	}
	defer db.Close()

	// 保存版块信息到数据库
	buf, err := encode(&forumInfo{
		Name:         forumName,
		PageNumber:   maxPagesNum,
		Introduction: forumIntroduction,
	})
	if err != nil {
		printError("EncodeForumInfo", err)
	}
	err = db.Put([]byte("info"), buf, nil)
	if err != nil {
		printError("PutForumInfo", err)
		os.Exit(1)
	}

	var done = make(chan int)
	// 用于统计还有多少页未完成
	var pageAmount = maxPagesNum

	getPage := func(page int) {
		postList, err := getPagesList(fid, page)
		if err != nil {
			printError("getPagesList"+string(page), err)
		} else {
			// 获取每个帖子的内容
			for i, v := range postList {
				// for用于重试
				for {
					resp, err := http.Get("http://www.mcbbs.net/" + v.Url)
					if err != nil {
						printError("GetPost", err)
					} else {
						defer resp.Body.Close()

						// 检查服务器是否返回成功
						if resp.StatusCode != 200 {
							// 服务器错误
							printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
							if resp.StatusCode == 404 {
								// 如果为404，则没有重试的必要，跳出重试
								printError("GetPost.Retry", "帖子《"+v.Title+"》不存在")
								break
							}
						} else {
							body, err := ioutil.ReadAll(resp.Body)
							if err != nil {
								printError("GetPost.ReadAll", err)
							} else {
								n := getPostBody.FindSubmatch(body)
								// 检查是否获取body成功
								if len(n) == 0 {
									printError("GetPost.FindSubmatch", "未找到页面内文章部分")
									// 获取失败，一般是帖子被屏蔽或者论坛数据库错误
									break
								} else {
									postBody := string(n[1])
									// 存入Body
									postList[i].Body = postBody
									printInfo("GetPost", v.Title, "[OK]")
									// 跳出循环重试
									break
								}
							}
						}
					}
					// 获取失败，重试
					printError("GetPost.Retry", "获取帖子《"+v.Title+"》失败，正在重试")
				}
			}
		}
		// 编码
		byt, err := encode(postList)
		if err != nil {
			printError("Encode", err)
			return
		}
		// 存入数据库
		err = db.Put([]byte("page_"+strconv.Itoa(page)), byt, nil)
		if err != nil {
			printError("db.Put", err)
			return
		}

		pageAmount--

		printInfo("OK", "线程", page, "执行完毕")
		printInfo("OK", "版块分页", page, "中的所有帖子已经储存到本地，还有", pageAmount, "页正在下载中")
		done <- page
	}

	// 如果页面数量小于线程数量
	// 直接开启maxPagesNum个线程
	if maxPagesNum < maxThread {
		for i := 0; i < maxPagesNum; i++ {
			go getPage(i + 1)
		}
	} else {
		// 开启设置数量的线程
		for i := 0; i < maxThread; i++ {
			go getPage(i + 1)
		}

		// 用于统计已开启线程数
		var i = maxThread

		// 单个线程执行完毕，补充线程
		// 使线程数量始终保持设置的最大数量
		for {
			// 等待线程完成
			<-done
			// 因为前面已经开启了maxThread个线程
			// 但是没有返回maxThread个done
			// 所以i最终会等于  maxThread + maxPagesNum
			// 所以需要检查已完成页数是否超出总页数
			if i <= maxPagesNum {
				i++
				// 开启新线程
				go getPage(i)
			}
		}
	}

	printInfo("OK", "FID为", fid, "的版块中的所有帖子已储存到本地")
}

// 获取单页面的帖子列表
func getPagesList(fid string, pageNum int) (pageList []postInfo, err error) {
	for {
		resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=" + strconv.Itoa(pageNum))
		if err != nil {
			printError("GetPageList", err)
		} else {
			defer resp.Body.Close()

			// 检查服务器是否返回成功
			if resp.StatusCode != 200 {
				// 服务器错误
				printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
			} else {
				body, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					printError("GetPageList.ReadAll", err)
				} else {
					/*
						// getPage所匹配的信息
						// 获取到的参数已用“m[x]说明”标出

						[m[2]分类] m[3]标题

						 ...234
						New

						m[4]作者
						m[6]发帖时间2/span>

						111189

						谢普
						5 天前

					*/
					m := getPostInfo.FindAllStringSubmatch(string(body), -1)

					// 检查是否获取成功
					if len(m) == 0 {
						printError("GetPageList.FindAllStringSubmatch", "未找到分页内帖子")
						// 失败，跳出
					} else {
						// 处理单个帖子信息
						for _, v := range m {
							date := v[5]
							// 如果发帖时间1为空，那么使用发帖时间2
							if date == "" {
								date = v[6]
							}
							// 储存帖子信息
							postInf := postInfo{
								Url:      v[1],
								Category: v[2],
								Title:    v[3],
								Author:   v[4],
								Date:     date,
							}
							pageList = append(pageList, postInf)
						}
					}
					// 跳出重试
					break
				}
			}
		}
		// 获取失败，重试
		printError("getPageList", "获取版块分页", pageNum, "失败，正在重试")
	}

	return pageList, nil
}

func printInfo(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[INFO]", s + ":"}, v...)...)
}

func printError(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[ERROR]", s + ":"}, v...)...)
}

func encode(data interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(data)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decode(data []byte, to interface{}) error {
	buf := bytes.NewBuffer(data)
	dec := gob.NewDecoder(buf)
	return dec.Decode(to)
}

type forumInfo struct {
	Name         string
	PageNumber   int
	Introduction string
}

type postInfo struct {
	Category string
	Url      string
	Title    string
	Author   string
	Date     string
	Body     string
}
<commit_msg>fixed: 请求一个不存在分页<commit_after>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strconv"

	"flag"
	"github.com/syndtr/goleveldb/leveldb"
)

var (
	// go的正则在匹配ReadAll出来的数据时有点奇怪，必须用[\w\W]{0,}?来代替\n

	// 最大同时开启线程数
	maxThread int
	// 用于获取版块名称
	getForumName = regexp.MustCompile(`<h1 class="xs2">[\w\W]{0,}?<a[^>]*>([^<]*)</a>`)
	// 用于获取版块帖子分页数量
	getForumPageNumber = regexp.MustCompile(`<a href="[^"]+" class="last">\.\.\. ([0-9]+)</a>`)
	// 用于获取版块介绍
	getForumIntroduction = regexp.MustCompile(`<div id="forum_rules_[0-9]*"[^>]*>([\w\W]{0,}?)(?:</div>[\w\W]{0,}?){3}<div class="drag">`)
	// 用于获取帖子信息
	getPostInfo = regexp.MustCompile(`id="normalthread_[0-9]+"[\w\W]{0,}?class="icn"[\w\W]{0,}?href="([^"]*)"[\w\W]{0,}?\n[\w\W]{0,}?>([^<]*)[\w\W]{0,}?class="xst" >([^<]*)[\w\W]{0,}?[\w\W]{0,}?>([^<]*)[\w\W]{0,}?(?:)?([^<]*)`)
	// 用于获取帖子内容
	getPostBody = regexp.MustCompile(`<div class="pcb">([\w\W]{0,}?)<div id="comment_[0-9]*" class="cm">`)
)

func main() {
	flag.IntVar(&maxThread, "m", 10, "最大线程数量")
	flag.Parse()

	if len(flag.Args()) == 0 {
		fmt.Println("请在命令后添加论坛版块FID参数")
		fmt.Println("例如：")
		fmt.Println(os.Args[0] + " 45 -m=10")
		fmt.Println("其中m为可选参数，用于设置最大线程数量")
		fmt.Println("FID可从版块URL中寻找")
		fmt.Println("数据会保存在db/FID路径中")
		return
	}

	printInfo("最大线程数量", maxThread)

	// 论坛版块fid
	fid := flag.Arg(0)
	printInfo("版块FID", fid)

	resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=1")
	if err != nil {
		printError("http.Get", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		printError("ReadAll", err)
		os.Exit(1)
	}

	if bytes.Contains(body, []byte(`<p>抱歉，指定的版块不存在</p>`)) {
		printError("GetForum", "版块不存在，请检查PID是否正确")
		os.Exit(1)
	}

	// 获取版块名称
	n := getForumName.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumName", "获取版块名称出错")
		os.Exit(1)
	}
	forumName := string(n[1])
	printInfo("版块名称", forumName)

	// 获取版块帖子分页数量
	n = getForumPageNumber.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumPageNumber", "获取版块分页数量出错")
		os.Exit(1)
	}
	maxPagesNum, err := strconv.Atoi(string(n[1]))
	if err != nil {
		printError("Atoi", err)
		os.Exit(1)
	}
	printInfo("本版块全部分页数量", maxPagesNum)

	// 获取版块介绍
	var forumIntroduction string
	n = getForumIntroduction.FindSubmatch(body)
	if len(n) == 0 {
		printError("GetForumIntroduction", "获取版块介绍出错，版块介绍为空")
	} else {
		forumIntroduction = string(n[1])
	}

	// 创建数据库
	db, err := leveldb.OpenFile("db/"+fid, nil)
	if err != nil {
		printError("OpenDB", err)
		os.Exit(1)
	}
	defer db.Close()

	// 保存版块信息到数据库
	buf, err := encode(&forumInfo{
		Name:         forumName,
		PageNumber:   maxPagesNum,
		Introduction: forumIntroduction,
	})
	if err != nil {
		printError("EncodeForumInfo", err)
	}
	err = db.Put([]byte("info"), buf, nil)
	if err != nil {
		printError("PutForumInfo", err)
		os.Exit(1)
	}

	var done = make(chan int)
	// 用于统计还有多少页未完成
	var pageAmount = maxPagesNum

	getPage := func(page int) {
		postList, err := getPagesList(fid, page)
		if err != nil {
			printError("getPagesList"+string(page), err)
		} else {
			// 获取每个帖子的内容
			for i, v := range postList {
				// for用于重试
				for {
					resp, err := http.Get("http://www.mcbbs.net/" + v.Url)
					if err != nil {
						printError("GetPost", err)
					} else {
						defer resp.Body.Close()

						// 检查服务器是否返回成功
						if resp.StatusCode != 200 {
							// 服务器错误
							printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
							if resp.StatusCode == 404 {
								// 如果为404，则没有重试的必要，跳出重试
								printError("GetPost.Retry", "帖子《"+v.Title+"》不存在")
								break
							}
						} else {
							body, err := ioutil.ReadAll(resp.Body)
							if err != nil {
								printError("GetPost.ReadAll", err)
							} else {
								n := getPostBody.FindSubmatch(body)
								// 检查是否获取body成功
								if len(n) == 0 {
									printError("GetPost.FindSubmatch", "未找到页面内文章部分")
									// 获取失败，一般是帖子被屏蔽或者论坛数据库错误
									break
								} else {
									postBody := string(n[1])
									// 存入Body
									postList[i].Body = postBody
									printInfo("GetPost", v.Title, "[OK]")
									// 跳出循环重试
									break
								}
							}
						}
					}
					// 获取失败，重试
					printError("GetPost.Retry", "获取帖子《"+v.Title+"》失败，正在重试")
				}
			}
		}
		// 编码
		byt, err := encode(postList)
		if err != nil {
			printError("Encode", err)
			return
		}
		// 存入数据库
		err = db.Put([]byte("page_"+strconv.Itoa(page)), byt, nil)
		if err != nil {
			printError("db.Put", err)
			return
		}

		pageAmount--

		printInfo("OK", "线程", page, "执行完毕")
		printInfo("OK", "版块分页", page, "中的所有帖子已经储存到本地，还有", pageAmount, "页正在下载中")
		done <- page
	}

	// 如果页面数量小于线程数量
	// 直接开启maxPagesNum个线程
	if maxPagesNum < maxThread {
		for i := 0; i < maxPagesNum; i++ {
			go getPage(i + 1)
		}
	} else {
		// 开启设置数量的线程
		for i := 0; i < maxThread; i++ {
			go getPage(i + 1)
		}

		// 用于统计已开启线程数
		var i = maxThread

		// 单个线程执行完毕，补充线程
		// 使线程数量始终保持设置的最大数量
		for {
			// 等待线程完成
			<-done
			i++
			// 因为前面已经开启了maxThread个线程
			// 但是没有返回maxThread个done
			// 所以i最终会等于  maxThread + maxPagesNum
			// 所以需要检查已完成页数是否超出总页数
			if i <= maxPagesNum {
				// 开启新线程
				go getPage(i)
			}
		}
	}

	printInfo("OK", "FID为", fid, "的版块中的所有帖子已储存到本地")
}

// 获取单页面的帖子列表
func getPagesList(fid string, pageNum int) (pageList []postInfo, err error) {
	for {
		resp, err := http.Get("http://www.mcbbs.net/forum.php?mod=forumdisplay&fid=" + fid + "&orderby=dateline&page=" + strconv.Itoa(pageNum))
		if err != nil {
			printError("GetPageList", err)
		} else {
			defer resp.Body.Close()

			// 检查服务器是否返回成功
			if resp.StatusCode != 200 {
				// 服务器错误
				printError("GetPost.ServerError", "服务器错误，错误码：", resp.StatusCode)
			} else {
				body, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					printError("GetPageList.ReadAll", err)
				} else {
					/*
						// getPage所匹配的信息
						// 获取到的参数已用“m[x]说明”标出

						[m[2]分类] m[3]标题

						 ...234
						New

						m[4]作者
						m[6]发帖时间2/span>

						111189

						谢普
						5 天前

					*/
					m := getPostInfo.FindAllStringSubmatch(string(body), -1)

					// 检查是否获取成功
					if len(m) == 0 {
						printError("GetPageList.FindAllStringSubmatch", "未找到分页内帖子")
						// 失败，跳出
					} else {
						// 处理单个帖子信息
						for _, v := range m {
							date := v[5]
							// 如果发帖时间1为空，那么使用发帖时间2
							if date == "" {
								date = v[6]
							}
							// 储存帖子信息
							postInf := postInfo{
								Url:      v[1],
								Category: v[2],
								Title:    v[3],
								Author:   v[4],
								Date:     date,
							}
							pageList = append(pageList, postInf)
						}
					}
					// 跳出重试
					break
				}
			}
		}
		// 获取失败，重试
		printError("getPageList", "获取版块分页", pageNum, "失败，正在重试")
	}

	return pageList, nil
}

func printInfo(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[INFO]", s + ":"}, v...)...)
}

func printError(s string, v ...interface{}) {
	log.Println(append([]interface{}{"[ERROR]", s + ":"}, v...)...)
}

func encode(data interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(data)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decode(data []byte, to interface{}) error {
	buf := bytes.NewBuffer(data)
	dec := gob.NewDecoder(buf)
	return dec.Decode(to)
}

type forumInfo struct {
	Name         string
	PageNumber   int
	Introduction string
}

type postInfo struct {
	Category string
	Url      string
	Title    string
	Author   string
	Date     string
	Body     string
}
<|endoftext|>


<commit_before>// Copyright 2011 Numrotron Inc.
// Use of this source code is governed by an MIT-style license
// that can be found in the LICENSE file.
//
// Developed at www.stathat.com by Patrick Crosby
// Contact us on twitter with any questions:  twitter.com/stat_hat

// splint is a little Go application to analyze Go source files.  It finds any functions that are
// too long or have too many parameters or results.
package main // import "stathat.com/c/splint"

import (
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path"
)

var statementThreshold = flag.Int("s", 30, "function statement count threshold")
var paramThreshold = flag.Int("p", 5, "parameter list length threshold")
var resultThreshold = flag.Int("r", 5, "result list length threshold")
var ifChainThreshold = flag.Int("c", 2, "if/else chain length threshold")
var outputJSON = flag.Bool("j", false, "output results as json")
var ignoreTestFiles = flag.Bool("i", true, "ignore test files")
var outputSummary = flag.Bool("sum", false, "output summary")

// Parser parses go source files, looking for potentially complex
// code.
type Parser struct {
	filename string
	first    bool
	summary  *Summary
	fileset  *token.FileSet
}

// Offender contains the file, function, position, and count of
// a block of code that splint has recognized as an issue.
type Offender struct {
	Filename string
	Function string
	Count    int
	Position token.Position
}

func (o *Offender) warning(msg string) {
	if *outputJSON {
		return
	}
	fmt.Printf("%s:\tfunction %s %s: %d\n", o.Position, o.Function, msg, o.Count)
}

func (o *Offender) warnNoCount(msg string) {
	if *outputJSON {
		return
	}
	fmt.Printf("%s:\tfunction %s %s\n", o.Position, o.Function, msg)
}

// Summary is a collection of Offenders for all the different
// checks that splint performs.
type Summary struct {
	Statement  []*Offender
	Param      []*Offender
	Result     []*Offender
	EmptyIfs   []*Offender
	IfChains   []*Offender
	BoolParams []*Offender

	// redundant, but using these for easy json output
	NumAboveStatementThreshold int
	NumAboveParamThreshold     int
	NumAboveResultThreshold    int
	NumIfChains                int
	NumEmptyIfs                int
}

func (s *Summary) IsClean() bool {
	return len(s.Statement) == 0 && len(s.Param) == 0 && len(s.Result) == 0 && len(s.EmptyIfs) == 0 && len(s.IfChains) == 0
}

func (s *Summary) addStatement(o *Offender) {
	s.Statement = append(s.Statement, o)
	s.NumAboveStatementThreshold++
	o.warning("too long")
}

func (s *Summary) addParam(o *Offender) {
	s.Param = append(s.Param, o)
	s.NumAboveParamThreshold++
	o.warning("too many params")
}

func (s *Summary) addBoolParam(o *Offender) {
	s.BoolParams = append(s.BoolParams, o)
	o.warning("bool function param")
}

func (s *Summary) addResult(o *Offender) {
	s.Result = append(s.Result, o)
	s.NumAboveResultThreshold++
	o.warning("too many results")
}

func (s *Summary) addEmptyIfBody(o *Offender) {
	s.EmptyIfs = append(s.EmptyIfs, o)
	s.NumEmptyIfs++
	o.warnNoCount("if with empty body")
}

func (s *Summary) addIfChain(o *Offender) {
	s.IfChains = append(s.IfChains, o)
	s.NumIfChains++
	o.warning("long if/else chain")
}

// NewParser creates a splint parser for a file.
func NewParser(filename string, summary *Summary) *Parser {
	return &Parser{filename: filename, first: true, summary: summary}
}

func statementCount(n ast.Node) int {
	total := 0
	counter := func(node ast.Node) bool {
		switch node.(type) {
		case ast.Stmt:
			total++
		}
		return true
	}
	ast.Inspect(n, counter)
	return total
}

func (p *Parser) offender(function string, count int, pos token.Pos) *Offender {
	return &Offender{
		Filename: p.filename,
		Function: function,
		Count:    count,
		Position: p.fileset.Position(pos),
	}
}

func (p *Parser) checkFuncLength(x *ast.FuncDecl) {
	numStatements := statementCount(x)
	if numStatements <= *statementThreshold {
		return
	}

	p.summary.addStatement(p.offender(x.Name.String(), numStatements, x.Pos()))
}

func (p *Parser) checkParamCount(x *ast.FuncDecl) {
	numFields := x.Type.Params.NumFields()
	if numFields <= *paramThreshold {
		return
	}

	p.summary.addParam(p.offender(x.Name.String(), numFields, x.Pos()))
}

func (p *Parser) checkBoolParams(x *ast.FuncDecl) {
	for _, f := range x.Type.Params.List {
		// this is ugly, but:
		if fmt.Sprintf("%s", f.Type) != "bool" {
			continue
		}
		p.summary.addBoolParam(p.offender(x.Name.String(), 0, x.Pos()))
	}
}

func (p *Parser) checkResultCount(x *ast.FuncDecl) {
	numResults := x.Type.Results.NumFields()
	if numResults <= *resultThreshold {
		return
	}

	p.summary.addResult(p.offender(x.Name.String(), numResults, x.Pos()))
}

func (p *Parser) checkEmptyIfs(x *ast.FuncDecl) {
	findIf := func(node ast.Node) bool {
		switch y := node.(type) {
		case *ast.IfStmt:
			if y.Body == nil || len(y.Body.List) == 0 {
				p.summary.addEmptyIfBody(p.offender(x.Name.String(), 0, y.Pos()))
			}
		}
		return true
	}
	ast.Inspect(x, findIf)
}

func chainLength(x *ast.IfStmt) int {
	if x.Else == nil {
		return 0
	}
	if ifst, ok := x.Else.(*ast.IfStmt); ok {
		return 1 + chainLength(ifst)
	}
	return 1
}

func (p *Parser) checkIfChains(x *ast.FuncDecl) {
	findIf := func(node ast.Node) bool {
		switch y := node.(type) {
		case *ast.IfStmt:
			n := chainLength(y)
			if n > *ifChainThreshold {
				p.summary.addIfChain(p.offender(x.Name.String(), n, y.Pos()))
			}
			return false // don't go any deeper
		}
		return true
	}
	ast.Inspect(x, findIf)
}

func (p *Parser) examineFunc(x *ast.FuncDecl) {
	p.checkFuncLength(x)
	p.checkParamCount(x)
	p.checkBoolParams(x)
	p.checkResultCount(x)
	p.checkEmptyIfs(x)
	p.checkIfChains(x)
}

func (p *Parser) examineDecls(tree *ast.File) {
	for _, v := range tree.Decls {
		switch x := v.(type) {
		case *ast.FuncDecl:
			p.examineFunc(x)
		}
	}
}

// Parse parses a file, looking for issues in functions.
func (p *Parser) Parse() {
	p.fileset = token.NewFileSet()
	tree, err := parser.ParseFile(p.fileset, p.filename, nil, 0)
	if err != nil {
		fmt.Printf("error parsing %s: %s\n", p.filename, err)
		return
	}

	p.examineDecls(tree)
}

func isTestFile(filename string) bool {
	base := path.Base(filename)
	match, err := path.Match("*_test.go", base)
	if err != nil {
		fmt.Println("match error:", err)
		return false
	}
	return match
}

func parseFile(filename string, summary *Summary) {
	if *ignoreTestFiles && isTestFile(filename) {
		return
	}
	parser := NewParser(filename, summary)
	parser.Parse()
}

func main() {
	flag.Parse()
	args := flag.Args()
	if len(args) == 0 {
		fmt.Println("Usage: splint [options] <go file>...")
		flag.PrintDefaults()
		os.Exit(1)
	}

	summary := new(Summary)

	for _, v := range args {
		parseFile(v, summary)
	}

	if *outputJSON {
		data, err := json.MarshalIndent(summary, "", "\t")
		if err != nil {
			fmt.Println("json encode error:", err)
		}
		fmt.Println(string(data))

	} else if *outputSummary {
		fmt.Println()
		fmt.Println("Number of functions above statement threshold:", summary.NumAboveStatementThreshold)
		fmt.Println("Number of functions above param threshold:", summary.NumAboveParamThreshold)
		fmt.Println("Number of functions above result threshold:", summary.NumAboveResultThreshold)
		fmt.Println("Number of long if/else chains:", summary.NumIfChains)
		fmt.Println("Number of empty if bodies:", summary.NumEmptyIfs)
		if !summary.IsClean() {
			os.Exit(1)
		}
	}
}
<commit_msg>Add flag to turn off bool param check<commit_after>// Copyright 2011 Numrotron Inc.
// Use of this source code is governed by an MIT-style license
// that can be found in the LICENSE file.
//
// Developed at www.stathat.com by Patrick Crosby
// Contact us on twitter with any questions:  twitter.com/stat_hat

// splint is a little Go application to analyze Go source files.  It finds any functions that are
// too long or have too many parameters or results.
package main // import "stathat.com/c/splint"

import (
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path"
)

var statementThreshold = flag.Int("s", 30, "function statement count threshold")
var paramThreshold = flag.Int("p", 5, "parameter list length threshold")
var resultThreshold = flag.Int("r", 5, "result list length threshold")
var ifChainThreshold = flag.Int("c", 2, "if/else chain length threshold")
var skipBoolParamCheck = flag.Bool("b", false, "don't warn on bool function params")
var outputJSON = flag.Bool("j", false, "output results as json")
var ignoreTestFiles = flag.Bool("i", true, "ignore test files")
var outputSummary = flag.Bool("sum", false, "output summary")

// Parser parses go source files, looking for potentially complex
// code.
type Parser struct {
	filename string
	first    bool
	summary  *Summary
	fileset  *token.FileSet
}

// Offender contains the file, function, position, and count of
// a block of code that splint has recognized as an issue.
type Offender struct {
	Filename string
	Function string
	Count    int
	Position token.Position
}

func (o *Offender) warning(msg string) {
	if *outputJSON {
		return
	}
	fmt.Printf("%s:\tfunction %s %s: %d\n", o.Position, o.Function, msg, o.Count)
}

func (o *Offender) warnNoCount(msg string) {
	if *outputJSON {
		return
	}
	fmt.Printf("%s:\tfunction %s %s\n", o.Position, o.Function, msg)
}

// Summary is a collection of Offenders for all the different
// checks that splint performs.
type Summary struct {
	Statement  []*Offender
	Param      []*Offender
	Result     []*Offender
	EmptyIfs   []*Offender
	IfChains   []*Offender
	BoolParams []*Offender

	// redundant, but using these for easy json output
	NumAboveStatementThreshold int
	NumAboveParamThreshold     int
	NumAboveResultThreshold    int
	NumIfChains                int
	NumEmptyIfs                int
}

func (s *Summary) IsClean() bool {
	return len(s.Statement) == 0 && len(s.Param) == 0 && len(s.Result) == 0 && len(s.EmptyIfs) == 0 && len(s.IfChains) == 0
}

func (s *Summary) addStatement(o *Offender) {
	s.Statement = append(s.Statement, o)
	s.NumAboveStatementThreshold++
	o.warning("too long")
}

func (s *Summary) addParam(o *Offender) {
	s.Param = append(s.Param, o)
	s.NumAboveParamThreshold++
	o.warning("too many params")
}

func (s *Summary) addBoolParam(o *Offender) {
	s.BoolParams = append(s.BoolParams, o)
	o.warning("bool function param")
}

func (s *Summary) addResult(o *Offender) {
	s.Result = append(s.Result, o)
	s.NumAboveResultThreshold++
	o.warning("too many results")
}

func (s *Summary) addEmptyIfBody(o *Offender) {
	s.EmptyIfs = append(s.EmptyIfs, o)
	s.NumEmptyIfs++
	o.warnNoCount("if with empty body")
}

func (s *Summary) addIfChain(o *Offender) {
	s.IfChains = append(s.IfChains, o)
	s.NumIfChains++
	o.warning("long if/else chain")
}

// NewParser creates a splint parser for a file.
func NewParser(filename string, summary *Summary) *Parser {
	return &Parser{filename: filename, first: true, summary: summary}
}

func statementCount(n ast.Node) int {
	total := 0
	counter := func(node ast.Node) bool {
		switch node.(type) {
		case ast.Stmt:
			total++
		}
		return true
	}
	ast.Inspect(n, counter)
	return total
}

func (p *Parser) offender(function string, count int, pos token.Pos) *Offender {
	return &Offender{
		Filename: p.filename,
		Function: function,
		Count:    count,
		Position: p.fileset.Position(pos),
	}
}

func (p *Parser) checkFuncLength(x *ast.FuncDecl) {
	numStatements := statementCount(x)
	if numStatements <= *statementThreshold {
		return
	}

	p.summary.addStatement(p.offender(x.Name.String(), numStatements, x.Pos()))
}

func (p *Parser) checkParamCount(x *ast.FuncDecl) {
	numFields := x.Type.Params.NumFields()
	if numFields <= *paramThreshold {
		return
	}

	p.summary.addParam(p.offender(x.Name.String(), numFields, x.Pos()))
}

func (p *Parser) checkBoolParams(x *ast.FuncDecl) {
	if *skipBoolParamCheck {
		return
	}
	for _, f := range x.Type.Params.List {
		// this is ugly, but:
		if fmt.Sprintf("%s", f.Type) != "bool" {
			continue
		}
		p.summary.addBoolParam(p.offender(x.Name.String(), 0, x.Pos()))
	}
}

func (p *Parser) checkResultCount(x *ast.FuncDecl) {
	numResults := x.Type.Results.NumFields()
	if numResults <= *resultThreshold {
		return
	}

	p.summary.addResult(p.offender(x.Name.String(), numResults, x.Pos()))
}

func (p *Parser) checkEmptyIfs(x *ast.FuncDecl) {
	findIf := func(node ast.Node) bool {
		switch y := node.(type) {
		case *ast.IfStmt:
			if y.Body == nil || len(y.Body.List) == 0 {
				p.summary.addEmptyIfBody(p.offender(x.Name.String(), 0, y.Pos()))
			}
		}
		return true
	}
	ast.Inspect(x, findIf)
}

func chainLength(x *ast.IfStmt) int {
	if x.Else == nil {
		return 0
	}
	if ifst, ok := x.Else.(*ast.IfStmt); ok {
		return 1 + chainLength(ifst)
	}
	return 1
}

func (p *Parser) checkIfChains(x *ast.FuncDecl) {
	findIf := func(node ast.Node) bool {
		switch y := node.(type) {
		case *ast.IfStmt:
			n := chainLength(y)
			if n > *ifChainThreshold {
				p.summary.addIfChain(p.offender(x.Name.String(), n, y.Pos()))
			}
			return false // don't go any deeper
		}
		return true
	}
	ast.Inspect(x, findIf)
}

func (p *Parser) examineFunc(x *ast.FuncDecl) {
	p.checkFuncLength(x)
	p.checkParamCount(x)
	p.checkBoolParams(x)
	p.checkResultCount(x)
	p.checkEmptyIfs(x)
	p.checkIfChains(x)
}

func (p *Parser) examineDecls(tree *ast.File) {
	for _, v := range tree.Decls {
		switch x := v.(type) {
		case *ast.FuncDecl:
			p.examineFunc(x)
		}
	}
}

// Parse parses a file, looking for issues in functions.
func (p *Parser) Parse() {
	p.fileset = token.NewFileSet()
	tree, err := parser.ParseFile(p.fileset, p.filename, nil, 0)
	if err != nil {
		fmt.Printf("error parsing %s: %s\n", p.filename, err)
		return
	}

	p.examineDecls(tree)
}

func isTestFile(filename string) bool {
	base := path.Base(filename)
	match, err := path.Match("*_test.go", base)
	if err != nil {
		fmt.Println("match error:", err)
		return false
	}
	return match
}

func parseFile(filename string, summary *Summary) {
	if *ignoreTestFiles && isTestFile(filename) {
		return
	}
	parser := NewParser(filename, summary)
	parser.Parse()
}

func main() {
	flag.Parse()
	args := flag.Args()
	if len(args) == 0 {
		fmt.Println("Usage: splint [options] <go file>...")
		flag.PrintDefaults()
		os.Exit(1)
	}

	summary := new(Summary)

	for _, v := range args {
		parseFile(v, summary)
	}

	if *outputJSON {
		data, err := json.MarshalIndent(summary, "", "\t")
		if err != nil {
			fmt.Println("json encode error:", err)
		}
		fmt.Println(string(data))

	} else if *outputSummary {
		fmt.Println()
		fmt.Println("Number of functions above statement threshold:", summary.NumAboveStatementThreshold)
		fmt.Println("Number of functions above param threshold:", summary.NumAboveParamThreshold)
		fmt.Println("Number of functions above result threshold:", summary.NumAboveResultThreshold)
		fmt.Println("Number of long if/else chains:", summary.NumIfChains)
		fmt.Println("Number of empty if bodies:", summary.NumEmptyIfs)
		if !summary.IsClean() {
			os.Exit(1)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
)

type Result struct {
	Arguments []string `json:"args"`
	About     string   `json:"about"`
	Stdout    string   `json:"stdout"`
	Stderr    string   `json:"stderr"`
	Success   bool     `json:"success"`
}

func Spruce(where string, flavor string, args ...string) (*Result, error) {
	r := &Result{Arguments: args}

	spruce := "spruce"
	if flavor != "" {
		spruce = "spruce-" + flavor
	}
	cmd := exec.Command(spruce, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Printf("failed to get stdout pipe from command: %s", err)
		return nil, err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Printf("failed to get stderr pipe from command: %s", err)
		return nil, err
	}

	if where != "" {
		err = os.Chdir(where)
		if err != nil {
			log.Printf("failed to chdir to %s: %s", where, err)
			return nil, err
		}
	}
	where, _ = os.Getwd()
	log.Printf("running `spruce %v' in %s", args, where)
	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	output := make(chan string)
	errors := make(chan string)
	drain := func(rd io.Reader, out chan string) {
		var buf []string
		s := bufio.NewScanner(rd)
		for s.Scan() {
			buf = append(buf, s.Text()+"\n")
		}
		out <- strings.Join(buf, "")
		close(out)
	}
	go drain(stdout, output)
	go drain(stderr, errors)

	err = cmd.Wait()
	r.Success = (err != nil)
	r.Stdout = <-output
	r.Stderr = <-errors
	return r, nil
}

type Merge struct {
	Flavor string   `json:"flavor"`
	Prune  []string `json:"prune"`
	YAML   []struct {
		Filename string `json:"filename"`
		Contents string `json:"contents"`
	} `json:"yaml"`

	Debug bool `json:"debug"`
	Trace bool `json:"trace"`
}

func SpruceMerge(m Merge) (*Result, error) {
	dir, err := ioutil.TempDir("", "web")
	if err != nil {
		log.Printf("failed to create temporary working directory: %s", err)
		return nil, err
	}
	// defer rm() the dir

	var args []string
	if m.Debug {
		args = append(args, "--debug")
	}
	if m.Trace {
		args = append(args, "--trace")
	}
	args = append(args, "merge")

	if len(m.Prune) > 0 {
		for _, f := range m.Prune {
			args = append(args, "--prune")
			args = append(args, f)
		}
	}
	for _, y := range m.YAML {
		args = append(args, y.Filename)
		ioutil.WriteFile(
			fmt.Sprintf("%s/%s", dir, y.Filename),
			[]byte(y.Contents),
			os.FileMode(0400),
		)
	}

	result, err := Spruce(dir, m.Flavor, args...)
	if version, err := Spruce("", m.Flavor, "-v"); err == nil {
		result.About = version.Stdout + version.Stderr
	} else {
		log.Printf("failed to determine spruce version information: %s", err)
	}
	return result, err
}
<commit_msg>Remove tempdirs<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
)

type Result struct {
	Arguments []string `json:"args"`
	About     string   `json:"about"`
	Stdout    string   `json:"stdout"`
	Stderr    string   `json:"stderr"`
	Success   bool     `json:"success"`
}

func Spruce(where string, flavor string, args ...string) (*Result, error) {
	r := &Result{Arguments: args}

	spruce := "spruce"
	if flavor != "" {
		spruce = "spruce-" + flavor
	}
	cmd := exec.Command(spruce, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Printf("failed to get stdout pipe from command: %s", err)
		return nil, err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Printf("failed to get stderr pipe from command: %s", err)
		return nil, err
	}

	if where != "" {
		err = os.Chdir(where)
		if err != nil {
			log.Printf("failed to chdir to %s: %s", where, err)
			return nil, err
		}
	}
	where, _ = os.Getwd()
	log.Printf("running `spruce %v' in %s", args, where)
	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	output := make(chan string)
	errors := make(chan string)
	drain := func(rd io.Reader, out chan string) {
		var buf []string
		s := bufio.NewScanner(rd)
		for s.Scan() {
			buf = append(buf, s.Text()+"\n")
		}
		out <- strings.Join(buf, "")
		close(out)
	}
	go drain(stdout, output)
	go drain(stderr, errors)

	err = cmd.Wait()
	r.Success = (err != nil)
	r.Stdout = <-output
	r.Stderr = <-errors
	return r, nil
}

type Merge struct {
	Flavor string   `json:"flavor"`
	Prune  []string `json:"prune"`
	YAML   []struct {
		Filename string `json:"filename"`
		Contents string `json:"contents"`
	} `json:"yaml"`

	Debug bool `json:"debug"`
	Trace bool `json:"trace"`
}

func SpruceMerge(m Merge) (*Result, error) {
	dir, err := ioutil.TempDir("", "web")
	if err != nil {
		log.Printf("failed to create temporary working directory: %s", err)
		return nil, err
	}
	defer os.RemoveAll(dir)

	var args []string
	if m.Debug {
		args = append(args, "--debug")
	}
	if m.Trace {
		args = append(args, "--trace")
	}
	args = append(args, "merge")

	if len(m.Prune) > 0 {
		for _, f := range m.Prune {
			args = append(args, "--prune")
			args = append(args, f)
		}
	}
	for _, y := range m.YAML {
		args = append(args, y.Filename)
		ioutil.WriteFile(
			fmt.Sprintf("%s/%s", dir, y.Filename),
			[]byte(y.Contents),
			os.FileMode(0400),
		)
	}

	result, err := Spruce(dir, m.Flavor, args...)
	if version, err := Spruce("", m.Flavor, "-v"); err == nil {
		result.About = version.Stdout + version.Stderr
	} else {
		log.Printf("failed to determine spruce version information: %s", err)
	}
	return result, err
}
<|endoftext|>


<commit_before>package spdystream

import (
	"code.google.com/p/go.net/spdy"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"
)

type Stream struct {
	streamId  spdy.StreamId
	parent    *Stream
	conn      *Connection
	startChan chan error

	dataLock sync.RWMutex
	dataChan chan []byte
	unread   []byte

	priority   uint8
	headers    http.Header
	headerChan chan http.Header
	finishLock sync.Mutex
	replied    bool
	finished   bool
	closeChan  chan bool
}

// WriteData writes data to stream, sending a dataframe per call
func (s *Stream) WriteData(data []byte, fin bool) error {
	var flags spdy.DataFlags

	if fin {
		flags = spdy.DataFlagFin
		s.finishLock.Lock()
		if s.finished {
			s.finishLock.Unlock()
			return ErrWriteClosedStream
		}
		s.finished = true
		s.finishLock.Unlock()
	}

	dataFrame := &spdy.DataFrame{
		StreamId: s.streamId,
		Flags:    flags,
		Data:     data,
	}

	return s.conn.framer.WriteFrame(dataFrame)
}

// Write writes bytes to a stream, calling write data for each call.
func (s *Stream) Write(data []byte) (n int, err error) {
	err = s.WriteData(data, false)
	if err == nil {
		n = len(data)
	}
	return
}

// Read reads bytes from a stream, a single read will never get more
// than what is sent on a single data frame, but a multiple calls to
// read may get data from the same data frame.
func (s *Stream) Read(p []byte) (n int, err error) {
	if s.unread == nil {
		select {
		case <-s.closeChan:
			return 0, io.EOF
		case read, ok := <-s.dataChan:
			if !ok {
				return 0, io.EOF
			}
			s.unread = read
		}
	}
	n = copy(p, s.unread)
	if n < len(s.unread) {
		s.unread = s.unread[n:]
	} else {
		s.unread = nil
	}
	return
}

// Wait waits for the stream to receive a reply.
func (s *Stream) Wait() error {
	return s.WaitTimeout(time.Duration(0))
}

// WaitTimeout waits for the stream to receive a reply or for timeout.
// When the timeout is reached, ErrTimeout will be returned.
func (s *Stream) WaitTimeout(timeout time.Duration) error {
	var timeoutChan <-chan time.Time
	if timeout > time.Duration(0) {
		timeoutChan = time.After(timeout)
	}

	select {
	case err := <-s.startChan:
		if err != nil {
			return err
		}
		break
	case <-timeoutChan:
		return ErrTimeout
	}
	return nil
}

// Close closes the stream by sending a reset stream frame, indicating
// this side is finished with the stream and the other side should close.
func (s *Stream) Close() error {
	s.conn.removeStream(s)
	s.finishLock.Lock()
	if s.finished {
		s.finishLock.Unlock()
		return nil
	}
	s.finished = true
	s.finishLock.Unlock()
	resetFrame := &spdy.RstStreamFrame{
		StreamId: s.streamId,
		Status:   spdy.Cancel,
	}
	return s.conn.framer.WriteFrame(resetFrame)
}

// CreateSubStream creates a stream using the current as the parent
func (s *Stream) CreateSubStream(headers http.Header, fin bool) (*Stream, error) {
	return s.conn.CreateStream(headers, s, fin)
}

// SetPriority sets the stream priority, does not affect the
// remote priority of this stream after Open has been called.
// Valid values are 0 through 7, 0 being the highest priority
// and 7 the lowest.
func (s *Stream) SetPriority(priority uint8) {
	s.priority = priority
}

// SendHeader sends a header frame across the stream
func (s *Stream) SendHeader(headers http.Header, fin bool) error {
	return s.conn.sendHeaders(headers, s, fin)
}

// SendReply sends a reply on a stream, only valid to be called once
// when handling a new stream
func (s *Stream) SendReply(headers http.Header, fin bool) error {
	if s.replied {
		return nil
	}
	s.replied = true
	return s.conn.sendReply(headers, s, fin)
}

// Refuse sends a reset frame with the status refuse, only
// valid to be called once when handling a new stream.  This
// may be used to indicate that a stream is not allowed
// when http status codes are not being used.
func (s *Stream) Refuse() error {
	if s.replied {
		return nil
	}
	s.replied = true
	return s.conn.sendReset(spdy.RefusedStream, s)
}

// Cancel sends a reset frame with the status canceled. This
// can be used at any time by the creator of the Stream to
// indicate the stream is no longer needed.
func (s *Stream) Cancel() error {
	return s.conn.sendReset(spdy.Cancel, s)
}

// ReceiveHeader receives a header sent on the other side
// of the stream.  This function will block until a header
// is received or stream is closed.
func (s *Stream) ReceiveHeader() (http.Header, error) {
	select {
	case <-s.closeChan:
		break
	case header, ok := <-s.headerChan:
		if !ok {
			return nil, fmt.Errorf("header chan closed")
		}
		return header, nil
	}
	return nil, fmt.Errorf("stream closed")
}

// Parent returns the parent stream
func (s *Stream) Parent() *Stream {
	return s.parent
}

// Headers returns the headers used to create the stream
func (s *Stream) Headers() http.Header {
	return s.headers
}

// String returns the string version of stream using the
// streamId to uniquely identify the stream
func (s *Stream) String() string {
	return fmt.Sprintf("stream:%d", s.streamId)
}

// IsFinished returns whether the stream has finished
// sending data
func (s *Stream) IsFinished() bool {
	return s.finished
}
<commit_msg>Add write lock usage in stream<commit_after>package spdystream

import (
	"code.google.com/p/go.net/spdy"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"
)

type Stream struct {
	streamId  spdy.StreamId
	parent    *Stream
	conn      *Connection
	startChan chan error

	dataLock sync.RWMutex
	dataChan chan []byte
	unread   []byte

	priority   uint8
	headers    http.Header
	headerChan chan http.Header
	finishLock sync.Mutex
	replied    bool
	finished   bool
	closeChan  chan bool
}

// WriteData writes data to stream, sending a dataframe per call
func (s *Stream) WriteData(data []byte, fin bool) error {
	var flags spdy.DataFlags

	if fin {
		flags = spdy.DataFlagFin
		s.finishLock.Lock()
		if s.finished {
			s.finishLock.Unlock()
			return ErrWriteClosedStream
		}
		s.finished = true
		s.finishLock.Unlock()
	}

	dataFrame := &spdy.DataFrame{
		StreamId: s.streamId,
		Flags:    flags,
		Data:     data,
	}

	s.conn.writeLock.Lock()
	defer s.conn.writeLock.Unlock()
	return s.conn.framer.WriteFrame(dataFrame)
}

// Write writes bytes to a stream, calling write data for each call.
func (s *Stream) Write(data []byte) (n int, err error) {
	err = s.WriteData(data, false)
	if err == nil {
		n = len(data)
	}
	return
}

// Read reads bytes from a stream, a single read will never get more
// than what is sent on a single data frame, but a multiple calls to
// read may get data from the same data frame.
func (s *Stream) Read(p []byte) (n int, err error) {
	if s.unread == nil {
		select {
		case <-s.closeChan:
			return 0, io.EOF
		case read, ok := <-s.dataChan:
			if !ok {
				return 0, io.EOF
			}
			s.unread = read
		}
	}
	n = copy(p, s.unread)
	if n < len(s.unread) {
		s.unread = s.unread[n:]
	} else {
		s.unread = nil
	}
	return
}

// Wait waits for the stream to receive a reply.
func (s *Stream) Wait() error {
	return s.WaitTimeout(time.Duration(0))
}

// WaitTimeout waits for the stream to receive a reply or for timeout.
// When the timeout is reached, ErrTimeout will be returned.
func (s *Stream) WaitTimeout(timeout time.Duration) error {
	var timeoutChan <-chan time.Time
	if timeout > time.Duration(0) {
		timeoutChan = time.After(timeout)
	}

	select {
	case err := <-s.startChan:
		if err != nil {
			return err
		}
		break
	case <-timeoutChan:
		return ErrTimeout
	}
	return nil
}

// Close closes the stream by sending a reset stream frame, indicating
// this side is finished with the stream and the other side should close.
func (s *Stream) Close() error {
	s.conn.removeStream(s)
	s.finishLock.Lock()
	if s.finished {
		s.finishLock.Unlock()
		return nil
	}
	s.finished = true
	s.finishLock.Unlock()
	resetFrame := &spdy.RstStreamFrame{
		StreamId: s.streamId,
		Status:   spdy.Cancel,
	}
	s.conn.writeLock.Lock()
	defer s.conn.writeLock.Unlock()
	return s.conn.framer.WriteFrame(resetFrame)
}

// CreateSubStream creates a stream using the current as the parent
func (s *Stream) CreateSubStream(headers http.Header, fin bool) (*Stream, error) {
	return s.conn.CreateStream(headers, s, fin)
}

// SetPriority sets the stream priority, does not affect the
// remote priority of this stream after Open has been called.
// Valid values are 0 through 7, 0 being the highest priority
// and 7 the lowest.
func (s *Stream) SetPriority(priority uint8) {
	s.priority = priority
}

// SendHeader sends a header frame across the stream
func (s *Stream) SendHeader(headers http.Header, fin bool) error {
	return s.conn.sendHeaders(headers, s, fin)
}

// SendReply sends a reply on a stream, only valid to be called once
// when handling a new stream
func (s *Stream) SendReply(headers http.Header, fin bool) error {
	if s.replied {
		return nil
	}
	s.replied = true
	return s.conn.sendReply(headers, s, fin)
}

// Refuse sends a reset frame with the status refuse, only
// valid to be called once when handling a new stream.  This
// may be used to indicate that a stream is not allowed
// when http status codes are not being used.
func (s *Stream) Refuse() error {
	if s.replied {
		return nil
	}
	s.replied = true
	return s.conn.sendReset(spdy.RefusedStream, s)
}

// Cancel sends a reset frame with the status canceled. This
// can be used at any time by the creator of the Stream to
// indicate the stream is no longer needed.
func (s *Stream) Cancel() error {
	return s.conn.sendReset(spdy.Cancel, s)
}

// ReceiveHeader receives a header sent on the other side
// of the stream.  This function will block until a header
// is received or stream is closed.
func (s *Stream) ReceiveHeader() (http.Header, error) {
	select {
	case <-s.closeChan:
		break
	case header, ok := <-s.headerChan:
		if !ok {
			return nil, fmt.Errorf("header chan closed")
		}
		return header, nil
	}
	return nil, fmt.Errorf("stream closed")
}

// Parent returns the parent stream
func (s *Stream) Parent() *Stream {
	return s.parent
}

// Headers returns the headers used to create the stream
func (s *Stream) Headers() http.Header {
	return s.headers
}

// String returns the string version of stream using the
// streamId to uniquely identify the stream
func (s *Stream) String() string {
	return fmt.Sprintf("stream:%d", s.streamId)
}

// IsFinished returns whether the stream has finished
// sending data
func (s *Stream) IsFinished() bool {
	return s.finished
}
<|endoftext|>


<commit_before>package glutton

import (
	"bufio"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"math/rand"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/kung-foo/freki"
)

// Mirai botnet  - https://github.com/CymmetriaResearch/MTPot/blob/master/mirai_conf.json
// Hajime botnet - https://security.rapiditynetworks.com/publications/2016-10-16/hajime.pdf
var miraiCom = map[string][]string{
	"ps":                          []string{"1 pts/21   00:00:00 init"},
	"cat /proc/mounts":            []string{"rootfs / rootfs rw 0 0\r\n/dev/root / ext2 rw,relatime,errors=continue 0 0\r\nproc /proc proc rw,relatime 0 0\r\nsysfs /sys sysfs rw,relatime 0 0\r\nudev /dev tmpfs rw,relatime 0 0\r\ndevpts /dev/pts devpts rw,relatime,mode=600,ptmxmode=000 0 0\r\n/dev/mtdblock1 /home/hik jffs2 rw,relatime 0 0\r\ntmpfs /run tmpfs rw,nosuid,noexec,relatime,size=3231524k,mode=755 0 0\r\n"},
	"(cat .s || cp /bin/echo .s)": []string{"cat: .s: No such file or directory"},
	"nc":   []string{"nc: command not found"},
	"wget": []string{"wget: missing URL"},
	"(dd bs=52 count=1 if=.s || cat .s)": []string{"\x7f\x45\x4c\x46\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x28\x00\x01\x00\x00\x00\xbc\x14\x01\x00\x34\x00\x00\x00"},
	"sh": []string{"$"},
	"echo -e \\x6b\\x61\\x6d\\x69/dev > /dev/.nippon":              []string{""},
	"cat /dev/.nippon":                                             []string{"kami/dev"},
	"rm /dev/.nippon":                                              []string{""},
	"echo -e \\x6b\\x61\\x6d\\x69/run > /run/.nippon":              []string{""},
	"cat /run/.nippon":                                             []string{"kami/run"},
	"rm /run/.nippon":                                              []string{""},
	"cat /bin/sh":                                                  []string{""},
	"/bin/busybox ps":                                              []string{"1 pts/21   00:00:00 init"},
	"/bin/busybox cat /proc/mounts":                                []string{"tmpfs /run tmpfs rw,nosuid,noexec,relatime,size=3231524k,mode=755 0 0"},
	"/bin/busybox echo -e \\x6b\\x61\\x6d\\x69/dev > /dev/.nippon": []string{""},
	"/bin/busybox cat /dev/.nippon":                                []string{"kami/dev"},
	"/bin/busybox rm /dev/.nippon":                                 []string{""},
	"/bin/busybox echo -e \\x6b\\x61\\x6d\\x69/run > /run/.nippon": []string{""},
	"/bin/busybox cat /run/.nippon":                                []string{"kami/run"},
	"/bin/busybox rm /run/.nippon":                                 []string{""},
	"/bin/busybox cat /bin/sh":                                     []string{""},
	"/bin/busybox cat /bin/echo":                                   []string{"/bin/busybox cat /bin/echo\r\n\x7f\x45\x4c\x46\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x28\x00\x01\x00\x00\x00\x6c\xb9\x00\x00\x34\x00\x00\x00"},
	"rm /dev/.human":                                               []string{"rm: can't remove '/.t': No such file or directory\r\nrm: can't remove '/.sh': No such file or directory\r\nrm: can't remove '/.human': No such file or directory\r\ncd /dev"},
}

func writeMsg(conn net.Conn, msg string, g *Glutton) error {
	_, err := conn.Write([]byte(msg))
	g.logger.Infof("[telnet  ] send: %q", msg)
	md := g.processor.Connections.GetByFlow(freki.NewConnKeyFromNetConn(conn))
	g.producer.LogHTTP(conn, md, msg, "write")
	return err
}

func readMsg(conn net.Conn, g *Glutton) (msg string, err error) {
	msg, err = bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		return "", err
	}
	g.logger.Infof("[telnet  ] recv: %q", msg)
	md := g.processor.Connections.GetByFlow(freki.NewConnKeyFromNetConn(conn))
	g.producer.LogHTTP(conn, md, msg, "read")
	return msg, err
}

func getSample(cmd string, g *Glutton) error {
	parts := strings.Split(cmd, " ")
	if len(parts) < 2 {
		return nil
	}
	url := parts[1]
	resp, err := http.Get(url)
	if err != nil {
		g.logger.Errorf("[telnet  ] getSample error: %v", err)
		return err
	}
	defer resp.Body.Close()
	p := new([]byte)
	resp.Body.Read(*p)
	sum := sha256.Sum256(*p)
	// Ignoring errors for if the folder already exists
	os.MkdirAll("samples", os.ModePerm)
	sha256Hash := hex.EncodeToString(sum[:])
	path := filepath.Join("samples", sha256Hash)
	out, err := os.Create(path)
	if err != nil {
		g.logger.Errorf("[telnet  ] getSample error: %v", err)
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, resp.Body)
	if err != nil {
		g.logger.Errorf("[telnet  ] getSample error: %v", err)
		return err
	}
	g.logger.Infof("[telnet  ] getSample succcess: %s", path)
	return nil
}

// HandleTelnet handles telnet communication on a connection
func (g *Glutton) HandleTelnet(conn net.Conn) {
	defer conn.Close()

	// TODO (glaslos): Add device banner

	// telnet window size negotiation response
	writeMsg(conn, "\xff\xfd\x18\xff\xfd\x20\xff\xfd\x23\xff\xfd\x27", g)

	// User name prompt
	writeMsg(conn, "Username: ", g)
	_, err := readMsg(conn, g)
	if err != nil {
		g.logger.Errorf("[telnet  ] %v", err)
		return
	}
	writeMsg(conn, "Password: ", g)
	_, err = readMsg(conn, g)
	if err != nil {
		g.logger.Errorf("[telnet  ] %v", err)
		return
	}

	writeMsg(conn, "welcome\r\n> ", g)
	for {
		msg, err := readMsg(conn, g)
		if err != nil {
			g.logger.Errorf("[telnet  ] %v", err)
			return
		}
		for _, cmd := range strings.Split(msg, ";") {
			if strings.HasPrefix(strings.TrimRight(cmd, ""), "wget http") {
				getSample(cmd, g)
			}
			if strings.TrimRight(cmd, "") == " rm /dev/.t" {
				continue
			}
			if strings.TrimRight(cmd, "\r\n") == " rm /dev/.sh" {
				continue
			}
			if strings.TrimRight(cmd, "\r\n") == "cd /dev/" {
				writeMsg(conn, "ECCHI: applet not found\r\n", g)
				writeMsg(conn, "\r\nBusyBox v1.16.1 (2014-03-04 16:00:18 CST) built-it shell (ash)\r\nEnter 'help' for a list of built-in commands.\r\n", g)
				continue
			}

			if resp := miraiCom[strings.TrimSpace(cmd)]; len(resp) > 0 {
				writeMsg(conn, resp[rand.Intn(len(resp))]+"\r\n", g)
			} else {
				// /bin/busybox YDKBI
				re := regexp.MustCompile(`\/bin\/busybox (?P<applet>[A-Z]+)`)
				match := re.FindStringSubmatch(cmd)
				if len(match) > 1 {
					writeMsg(conn, match[1]+": applet not found\r\n", g)
					writeMsg(conn, "BusyBox v1.16.1 (2014-03-04 16:00:18 CST) built-in shell (ash)\r\nEnter 'help' for a list of built-in commands.\r\n", g)
				}
			}
		}
		writeMsg(conn, "> ", g)
	}
}
<commit_msg>timeout on getSample and fixed pattern matching<commit_after>package glutton

import (
	"bufio"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"math/rand"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/kung-foo/freki"
)

// Mirai botnet  - https://github.com/CymmetriaResearch/MTPot/blob/master/mirai_conf.json
// Hajime botnet - https://security.rapiditynetworks.com/publications/2016-10-16/hajime.pdf
var miraiCom = map[string][]string{
	"ps":                          []string{"1 pts/21   00:00:00 init"},
	"cat /proc/mounts":            []string{"rootfs / rootfs rw 0 0\r\n/dev/root / ext2 rw,relatime,errors=continue 0 0\r\nproc /proc proc rw,relatime 0 0\r\nsysfs /sys sysfs rw,relatime 0 0\r\nudev /dev tmpfs rw,relatime 0 0\r\ndevpts /dev/pts devpts rw,relatime,mode=600,ptmxmode=000 0 0\r\n/dev/mtdblock1 /home/hik jffs2 rw,relatime 0 0\r\ntmpfs /run tmpfs rw,nosuid,noexec,relatime,size=3231524k,mode=755 0 0\r\n"},
	"(cat .s || cp /bin/echo .s)": []string{"cat: .s: No such file or directory"},
	"nc":   []string{"nc: command not found"},
	"wget": []string{"wget: missing URL"},
	"(dd bs=52 count=1 if=.s || cat .s)": []string{"\x7f\x45\x4c\x46\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x28\x00\x01\x00\x00\x00\xbc\x14\x01\x00\x34\x00\x00\x00"},
	"sh": []string{"$"},
	"echo -e \\x6b\\x61\\x6d\\x69/dev > /dev/.nippon":              []string{""},
	"cat /dev/.nippon":                                             []string{"kami/dev"},
	"rm /dev/.nippon":                                              []string{""},
	"echo -e \\x6b\\x61\\x6d\\x69/run > /run/.nippon":              []string{""},
	"cat /run/.nippon":                                             []string{"kami/run"},
	"rm /run/.nippon":                                              []string{""},
	"cat /bin/sh":                                                  []string{""},
	"/bin/busybox ps":                                              []string{"1 pts/21   00:00:00 init"},
	"/bin/busybox cat /proc/mounts":                                []string{"tmpfs /run tmpfs rw,nosuid,noexec,relatime,size=3231524k,mode=755 0 0"},
	"/bin/busybox echo -e \\x6b\\x61\\x6d\\x69/dev > /dev/.nippon": []string{""},
	"/bin/busybox cat /dev/.nippon":                                []string{"kami/dev"},
	"/bin/busybox rm /dev/.nippon":                                 []string{""},
	"/bin/busybox echo -e \\x6b\\x61\\x6d\\x69/run > /run/.nippon": []string{""},
	"/bin/busybox cat /run/.nippon":                                []string{"kami/run"},
	"/bin/busybox rm /run/.nippon":                                 []string{""},
	"/bin/busybox cat /bin/sh":                                     []string{""},
	"/bin/busybox cat /bin/echo":                                   []string{"/bin/busybox cat /bin/echo\r\n\x7f\x45\x4c\x46\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x28\x00\x01\x00\x00\x00\x6c\xb9\x00\x00\x34\x00\x00\x00"},
	"rm /dev/.human":                                               []string{"rm: can't remove '/.t': No such file or directory\r\nrm: can't remove '/.sh': No such file or directory\r\nrm: can't remove '/.human': No such file or directory\r\ncd /dev"},
}

func writeMsg(conn net.Conn, msg string, g *Glutton) error {
	_, err := conn.Write([]byte(msg))
	g.logger.Infof("[telnet  ] send: %q", msg)
	md := g.processor.Connections.GetByFlow(freki.NewConnKeyFromNetConn(conn))
	g.producer.LogHTTP(conn, md, msg, "write")
	return err
}

func readMsg(conn net.Conn, g *Glutton) (msg string, err error) {
	msg, err = bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		return "", err
	}
	g.logger.Infof("[telnet  ] recv: %q", msg)
	md := g.processor.Connections.GetByFlow(freki.NewConnKeyFromNetConn(conn))
	g.producer.LogHTTP(conn, md, msg, "read")
	return msg, err
}

func getSample(cmd string, g *Glutton) error {
	parts := strings.Split(cmd, " ")
	if len(parts) < 2 {
		return nil
	}
	url := parts[1]
	timeout := time.Duration(5 * time.Second)
	client := http.Client{
		Timeout: timeout,
	}
	resp, err := client.Get(url)
	if err != nil {
		g.logger.Errorf("[telnet  ] getSample error: %v", err)
		return err
	}
	defer resp.Body.Close()
	p := new([]byte)
	resp.Body.Read(*p)
	sum := sha256.Sum256(*p)
	// Ignoring errors for if the folder already exists
	os.MkdirAll("samples", os.ModePerm)
	sha256Hash := hex.EncodeToString(sum[:])
	path := filepath.Join("samples", sha256Hash)
	out, err := os.Create(path)
	if err != nil {
		g.logger.Errorf("[telnet  ] getSample error: %v", err)
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, resp.Body)
	if err != nil {
		g.logger.Errorf("[telnet  ] getSample error: %v", err)
		return err
	}
	g.logger.Infof("[telnet  ] getSample succcess: %s", path)
	return nil
}

// HandleTelnet handles telnet communication on a connection
func (g *Glutton) HandleTelnet(conn net.Conn) {
	defer conn.Close()

	// TODO (glaslos): Add device banner

	// telnet window size negotiation response
	writeMsg(conn, "\xff\xfd\x18\xff\xfd\x20\xff\xfd\x23\xff\xfd\x27", g)

	// User name prompt
	writeMsg(conn, "Username: ", g)
	_, err := readMsg(conn, g)
	if err != nil {
		g.logger.Errorf("[telnet  ] %v", err)
		return
	}
	writeMsg(conn, "Password: ", g)
	_, err = readMsg(conn, g)
	if err != nil {
		g.logger.Errorf("[telnet  ] %v", err)
		return
	}

	writeMsg(conn, "welcome\r\n> ", g)
	for {
		msg, err := readMsg(conn, g)
		if err != nil {
			g.logger.Errorf("[telnet  ] %v", err)
			return
		}
		for _, cmd := range strings.Split(msg, ";") {
			if strings.HasPrefix(strings.Trim(cmd, " "), "wget http") {
				go getSample(cmd, g)
			}
			if strings.TrimRight(cmd, "") == " rm /dev/.t" {
				continue
			}
			if strings.TrimRight(cmd, "\r\n") == " rm /dev/.sh" {
				continue
			}
			if strings.TrimRight(cmd, "\r\n") == "cd /dev/" {
				writeMsg(conn, "ECCHI: applet not found\r\n", g)
				writeMsg(conn, "\r\nBusyBox v1.16.1 (2014-03-04 16:00:18 CST) built-it shell (ash)\r\nEnter 'help' for a list of built-in commands.\r\n", g)
				continue
			}

			if resp := miraiCom[strings.TrimSpace(cmd)]; len(resp) > 0 {
				writeMsg(conn, resp[rand.Intn(len(resp))]+"\r\n", g)
			} else {
				// /bin/busybox YDKBI
				re := regexp.MustCompile(`\/bin\/busybox (?P<applet>[A-Z]+)`)
				match := re.FindStringSubmatch(cmd)
				if len(match) > 1 {
					writeMsg(conn, match[1]+": applet not found\r\n", g)
					writeMsg(conn, "BusyBox v1.16.1 (2014-03-04 16:00:18 CST) built-in shell (ash)\r\nEnter 'help' for a list of built-in commands.\r\n", g)
				}
			}
		}
		writeMsg(conn, "> ", g)
	}
}
<|endoftext|>


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

package tfgo

import (
	tf "github.com/tensorflow/tensorflow/tensorflow/go"
	"github.com/tensorflow/tensorflow/tensorflow/go/op"
)

// Tensor is an high level abstraction for the
// tf.Output structure, associating a scope to the Tensor
type Tensor struct {
	Path   *op.Scope
	Output tf.Output
}

// NewTensor creates a *Tensor from a tf.Output
// Place the cloned tensor within the specified scope
func NewTensor(scope *op.Scope, tfout tf.Output) (tensor *Tensor) {
	tensor = new(Tensor)
	tensor.Path = NewScope(scope)
	// Copy the tensor to a new node in the graph
	tensor.Output = op.Identity(tensor.Path.SubScope("Identity"), tfout)
	return tensor
}

// Check checks if the previous operation caused an error
// and thus tensor.Path.Err is not nil.
// If it's not, panics because we're defining the graph in a wrong way
func (tensor *Tensor) Check() {
	err := tensor.Path.Err()
	if err != nil {
		panic(err.Error())
	}
}

// Scope returns the scope associated to the tensor
func (tensor *Tensor) Scope() *op.Scope {
	return tensor.Path
}

// Shape64 returns the shape of the tensor as []int64.
// If firstDimension is true a 4 elements slice is returned.
// Otherwise a 3 elements sice is returned.
func (tensor *Tensor) Shape64(firstDimension bool) []int64 {
	dims64, _ := tensor.Output.Shape().ToSlice()
	if firstDimension {
		return dims64
	}
	return dims64[1:]
}

// Shape32 returns the shape of the tensor as []int32.
// If firstDimension is true a 4 elements slice is returned.
// Otherwise a 3 elements sice is returned.
func (tensor *Tensor) Shape32(firstDimension bool) []int32 {
	dims64 := tensor.Shape64(firstDimension)
	var dims32 = make([]int32, len(dims64))
	for idx, dim := range dims64 {
		dims32[idx] = int32(dim)
	}
	return dims32
}

// Dtype returns the tensor dtype
func (tensor *Tensor) Dtype() tf.DataType {
	return tensor.Output.DataType()
}

// --------
// Methods returning *Tensor
// Structs that embed *Tensor should ""override"" them
// Returning their own type
// --------

// Clone returns a copy of the current tensor in a new scope
// Clone must be used when one want to create a different tensor
// from the output of an operation.
func (tensor *Tensor) Clone() *Tensor {
	defer tensor.Check()
	scope := NewScope(tensor.Path)
	return NewTensor(scope, tensor.Output)
}

// Cast casts the current tensor to the requested dtype
func (tensor *Tensor) Cast(dtype tf.DataType) *Tensor {
	defer tensor.Check()
	tensor.Output = Cast(tensor.Path, tensor.Output, dtype)
	return tensor
}

// Add defines the add operation between the tensor and tfout
// `tfout` dtype is converted to tensor.Dtype() before adding
func (tensor *Tensor) Add(tfout tf.Output) *Tensor {
	defer tensor.Check()
	s := tensor.Path.SubScope("Add")
	tensor.Output = op.Add(s, tensor.Output, Cast(s, tfout, tensor.Dtype()))
	return tensor
}

// Mul defines the multiplication operation between the tensor
// and `tfout`.
// `tfout` dtype is converted to tensor.Dtype() before multiplying
func (tensor *Tensor) Mul(tfout tf.Output) *Tensor {
	defer tensor.Check()
	s := tensor.Path.SubScope("Mul")
	tensor.Output = op.Mul(s, tensor.Output, Cast(s, tfout, tensor.Dtype()))
	return tensor
}

// Pow defines the pow operation x^y, where x are the tensor values
// y dtype is converted to tensor.Dtype() before executing Pow
func (tensor *Tensor) Pow(y tf.Output) *Tensor {
	defer tensor.Check()
	s := tensor.Path.SubScope("Pow")
	tensor.Output = op.Pow(s, tensor.Output, Cast(s, y, tensor.Dtype()))
	return tensor
}

// Square defines the square operation for the tensor values
func (tensor *Tensor) Square() *Tensor {
	defer tensor.Check()
	return tensor.Pow(op.Const(tensor.Path.SubScope("y"), 2.))
}

// Sqrt defines the square root operation for the tensor values
func (tensor *Tensor) Sqrt() *Tensor {
	defer tensor.Check()
	return tensor.Pow(op.Const(tensor.Path.SubScope("y"), 0.5))
}
<commit_msg>added tensor.MatMul<commit_after>/*
Copyright 2017 Paolo Galeone. All right reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package tfgo

import (
	tf "github.com/tensorflow/tensorflow/tensorflow/go"
	"github.com/tensorflow/tensorflow/tensorflow/go/op"
)

// Tensor is an high level abstraction for the
// tf.Output structure, associating a scope to the Tensor
type Tensor struct {
	Path   *op.Scope
	Output tf.Output
}

// NewTensor creates a *Tensor from a tf.Output
// Place the cloned tensor within the specified scope
func NewTensor(scope *op.Scope, tfout tf.Output) (tensor *Tensor) {
	tensor = new(Tensor)
	tensor.Path = NewScope(scope)
	// Copy the tensor to a new node in the graph
	tensor.Output = op.Identity(tensor.Path.SubScope("Identity"), tfout)
	return tensor
}

// Check checks if the previous operation caused an error
// and thus tensor.Path.Err is not nil.
// If it's not, panics because we're defining the graph in a wrong way
func (tensor *Tensor) Check() {
	err := tensor.Path.Err()
	if err != nil {
		panic(err.Error())
	}
}

// Scope returns the scope associated to the tensor
func (tensor *Tensor) Scope() *op.Scope {
	return tensor.Path
}

// Shape64 returns the shape of the tensor as []int64.
// If firstDimension is true a 4 elements slice is returned.
// Otherwise a 3 elements sice is returned.
func (tensor *Tensor) Shape64(firstDimension bool) []int64 {
	dims64, _ := tensor.Output.Shape().ToSlice()
	if firstDimension {
		return dims64
	}
	return dims64[1:]
}

// Shape32 returns the shape of the tensor as []int32.
// If firstDimension is true a 4 elements slice is returned.
// Otherwise a 3 elements sice is returned.
func (tensor *Tensor) Shape32(firstDimension bool) []int32 {
	dims64 := tensor.Shape64(firstDimension)
	var dims32 = make([]int32, len(dims64))
	for idx, dim := range dims64 {
		dims32[idx] = int32(dim)
	}
	return dims32
}

// Dtype returns the tensor dtype
func (tensor *Tensor) Dtype() tf.DataType {
	return tensor.Output.DataType()
}

// --------
// Methods returning *Tensor
// Structs that embed *Tensor should ""override"" them
// Returning their own type
// --------

// Clone returns a copy of the current tensor in a new scope
// Clone must be used when one want to create a different tensor
// from the output of an operation.
func (tensor *Tensor) Clone() *Tensor {
	defer tensor.Check()
	scope := NewScope(tensor.Path)
	return NewTensor(scope, tensor.Output)
}

// Cast casts the current tensor to the requested dtype
func (tensor *Tensor) Cast(dtype tf.DataType) *Tensor {
	defer tensor.Check()
	tensor.Output = Cast(tensor.Path, tensor.Output, dtype)
	return tensor
}

// Add defines the add operation between the tensor and tfout
// `tfout` dtype is converted to tensor.Dtype() before adding
func (tensor *Tensor) Add(tfout tf.Output) *Tensor {
	defer tensor.Check()
	s := tensor.Path.SubScope("Add")
	tensor.Output = op.Add(s, tensor.Output, Cast(s, tfout, tensor.Dtype()))
	return tensor
}

// Mul defines the multiplication operation between the tensor
// and `tfout`. It's the multiplication element-wise with broadcasting support.
// `tfout` dtype is converted to tensor.Dtype() before multiplying
func (tensor *Tensor) Mul(tfout tf.Output) *Tensor {
	defer tensor.Check()
	s := tensor.Path.SubScope("Mul")
	tensor.Output = op.Mul(s, tensor.Output, Cast(s, tfout, tensor.Dtype()))
	return tensor
}

// MatMul defines the matrix multiplication operation between the tensor
// and `tfout`.
// `tfout` dtype is converted to tensor.Dtype() before multiplying
func (tensor *Tensor) MatMul(tfout tf.Output) *Tensor {
	defer tensor.Check()
	s := tensor.Path.SubScope("MatMul")
	tensor.Output = op.MatMul(s, tensor.Output, Cast(s, tfout, tensor.Dtype()))
	return tensor
}

// Pow defines the pow operation x^y, where x are the tensor values
// y dtype is converted to tensor.Dtype() before executing Pow
func (tensor *Tensor) Pow(y tf.Output) *Tensor {
	defer tensor.Check()
	s := tensor.Path.SubScope("Pow")
	tensor.Output = op.Pow(s, tensor.Output, Cast(s, y, tensor.Dtype()))
	return tensor
}

// Square defines the square operation for the tensor values
func (tensor *Tensor) Square() *Tensor {
	defer tensor.Check()
	return tensor.Pow(op.Const(tensor.Path.SubScope("y"), 2.))
}

// Sqrt defines the square root operation for the tensor values
func (tensor *Tensor) Sqrt() *Tensor {
	defer tensor.Check()
	return tensor.Pow(op.Const(tensor.Path.SubScope("y"), 0.5))
}
<|endoftext|>


<commit_before>package maxminddb

import (
	"fmt"
	"log"
	"net"
)

type onlyCountry struct {
	Country struct {
		IsoCode string `maxminddb:"iso_code"`
	} `maxminddb:"country"`
}

// Example_struct shows how to decode to a struct
func Example_struct() {
	db, err := Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	if err != nil {
		log.Fatal(err)
	}
	ip := net.ParseIP("81.2.69.142")

	var record onlyCountry // Or any appropriate struct
	err = db.Lookup(ip, &record)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Print(record.Country.IsoCode)
	// Output:
	// GB
	db.Close()
}

// Example_interface demonstrates how to decode to an interface{}
func Example_interface() {
	db, err := Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	if err != nil {
		log.Fatal(err)
	}
	ip := net.ParseIP("81.2.69.142")

	var record interface{}
	err = db.Lookup(ip, &record)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%v", record)
	db.Close()
}
<commit_msg>Try putting examples back in test package<commit_after>package maxminddb_test

import (
	"fmt"
	"github.com/oschwald/maxminddb-golang"
	"log"
	"net"
)

type onlyCountry struct {
	Country struct {
		IsoCode string `maxminddb:"iso_code"`
	} `maxminddb:"country"`
}

// Example_struct shows how to decode to a struct
func Example_struct() {
	db, err := maxminddb.Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	if err != nil {
		log.Fatal(err)
	}
	ip := net.ParseIP("81.2.69.142")

	var record onlyCountry // Or any appropriate struct
	err = db.Lookup(ip, &record)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Print(record.Country.IsoCode)
	// Output:
	// GB
	db.Close()
}

// Example_interface demonstrates how to decode to an interface{}
func Example_interface() {
	db, err := maxminddb.Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	if err != nil {
		log.Fatal(err)
	}
	ip := net.ParseIP("81.2.69.142")

	var record interface{}
	err = db.Lookup(ip, &record)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%v", record)
	db.Close()
}
<|endoftext|>


<commit_before>package epub_test

import (
	"fmt"
	"log"

	"github.com/bmaupin/go-epub"
)

func ExampleEpub_AddCSS() {
	e := epub.NewEpub("My title")

	// Add CSS
	css1Path, err := e.AddCSS("cover.css", "epub.css")
	if err != nil {
		log.Fatal(err)
	}

	// The filename is optional
	css2Path, err := e.AddCSS("cover.css", "")
	if err != nil {
		log.Fatal(err)
	}

	// Use the CSS in a section
	sectionBody := `
Section 1

This is a paragraph.
`
	e.AddSection("Section 1", sectionBody, "", css1Path)

	fmt.Println(css1Path)
	fmt.Println(css2Path)

	// Output:
	// ../css/epub.css
	// ../css/cover.css
}

func ExampleEpub_AddFont() {
	e := epub.NewEpub("My title")

	// Add a font from a local file
	font1Path, err := e.AddFont("testdata/redacted-script-regular.ttf", "font.ttf")
	if err != nil {
		log.Fatal(err)
	}

	// The filename is optional
	font2Path, err := e.AddFont("testdata/redacted-script-regular.ttf", "")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(font1Path)
	fmt.Println(font2Path)

	// Output:
	// ../fonts/font.ttf
	// ../fonts/redacted-script-regular.ttf
}

func ExampleEpub_AddImage() {
	e := epub.NewEpub("My title")

	// Add an image from a local file
	img1Path, err := e.AddImage("testdata/gophercolor16x16.png", "go-gopher.png")
	if err != nil {
		log.Fatal(err)
	}

	// Add an image from a URL. The filename is optional
	img2Path, err := e.AddImage("https://golang.org/doc/gopher/gophercolor16x16.png", "")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(img1Path)
	fmt.Println(img2Path)

	// Output:
	// ../images/go-gopher.png
	// ../images/gophercolor16x16.png
}

func ExampleEpub_AddSection() {
	e := epub.NewEpub("My title")

	// Add a section. The CSS path is optional
	section1Body := `
Section 1

This is a paragraph.
`
	section1Path, err := e.AddSection("Section 1", section1Body, "firstsection.xhtml", "")
	if err != nil {
		log.Fatal(err)
	}

	// Link to the first section
	section2Body := fmt.Sprintf(`
Section 2

	Link to section 1`,
		section1Path)
	// The filename is also optional
	section2Path, err := e.AddSection("Section 2", section2Body, "", "")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(section1Path)
	fmt.Println(section2Path)

	// Output:
	// firstsection.xhtml
	// section0002.xhtml
}

func ExampleEpub_SetCover() {
	e := epub.NewEpub("My title")

	// Set the cover. The CSS content is optional
	e.SetCover("testdata/gophercolor16x16.png", "")

	// Update the cover using custom CSS
	coverCSSContent := `h1 {
  text-align: center;
}`
	e.SetCover("testdata/gophercolor16x16.png", coverCSSContent)
}
<commit_msg>Update example tests<commit_after>package epub_test

import (
	"fmt"
	"log"

	"github.com/bmaupin/go-epub"
)

func ExampleEpub_AddCSS() {
	e := epub.NewEpub("My title")

	// Add CSS
	css1Path, err := e.AddCSS("cover.css", "epub.css")
	if err != nil {
		log.Fatal(err)
	}

	// The filename is optional
	css2Path, err := e.AddCSS("cover.css", "")
	if err != nil {
		log.Fatal(err)
	}

	// Use the CSS in a section
	sectionBody := `
Section 1

This is a paragraph.
`
	e.AddSection(sectionBody, "Section 1", "", css1Path)

	fmt.Println(css1Path)
	fmt.Println(css2Path)

	// Output:
	// ../css/epub.css
	// ../css/cover.css
}

func ExampleEpub_AddFont() {
	e := epub.NewEpub("My title")

	// Add a font from a local file
	font1Path, err := e.AddFont("testdata/redacted-script-regular.ttf", "font.ttf")
	if err != nil {
		log.Fatal(err)
	}

	// The filename is optional
	font2Path, err := e.AddFont("testdata/redacted-script-regular.ttf", "")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(font1Path)
	fmt.Println(font2Path)

	// Output:
	// ../fonts/font.ttf
	// ../fonts/redacted-script-regular.ttf
}

func ExampleEpub_AddImage() {
	e := epub.NewEpub("My title")

	// Add an image from a local file
	img1Path, err := e.AddImage("testdata/gophercolor16x16.png", "go-gopher.png")
	if err != nil {
		log.Fatal(err)
	}

	// Add an image from a URL. The filename is optional
	img2Path, err := e.AddImage("https://golang.org/doc/gopher/gophercolor16x16.png", "")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(img1Path)
	fmt.Println(img2Path)

	// Output:
	// ../images/go-gopher.png
	// ../images/gophercolor16x16.png
}

func ExampleEpub_AddSection() {
	e := epub.NewEpub("My title")

	// Add a section. The CSS path is optional
	section1Body := `
Section 1

This is a paragraph.
`
	section1Path, err := e.AddSection(section1Body, "Section 1", "firstsection.xhtml", "")
	if err != nil {
		log.Fatal(err)
	}

	// Link to the first section
	section2Body := fmt.Sprintf(`
Section 2

	Link to section 1`,
		section1Path)
	// The title and filename are also optional
	section2Path, err := e.AddSection(section2Body, "", "", "")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(section1Path)
	fmt.Println(section2Path)

	// Output:
	// firstsection.xhtml
	// section0002.xhtml
}

func ExampleEpub_SetCover() {
	e := epub.NewEpub("My title")

	// Set the cover. The CSS content is optional
	e.SetCover("testdata/gophercolor16x16.png", "")

	// Update the cover using custom CSS
	coverCSSContent := `h1 {
  text-align: center;
}`
	e.SetCover("testdata/gophercolor16x16.png", coverCSSContent)
}
<|endoftext|>


<commit_before>package fbapi_test

import (
	"fmt"
	"net/http"
	"net/url"
	"os"

	"github.com/daaku/go.fbapi"
)

func Example() {
	client := &fbapi.Client{}
	var show struct {
		ID      uint64 `json:"id,string"`
		Name    string `json:"name"`
		Network string `json:"network"`
	}
	_, err := client.Do(
		&http.Request{Method: "GET", URL: &url.URL{Path: "DoctorWho"}},
		&show,
	)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	fmt.Printf("%+v", show)

	// Output: {ID:127031120644257 Name:Naitik Shah Network:BBC Worldwide}
}
<commit_msg>fix example output<commit_after>package fbapi_test

import (
	"fmt"
	"net/http"
	"net/url"
	"os"

	"github.com/daaku/go.fbapi"
)

func Example() {
	client := &fbapi.Client{}
	var show struct {
		ID      uint64 `json:"id,string"`
		Name    string `json:"name"`
		Network string `json:"network"`
	}
	_, err := client.Do(
		&http.Request{Method: "GET", URL: &url.URL{Path: "DoctorWho"}},
		&show,
	)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	fmt.Printf("%+v", show)

	// Output: {ID:127031120644257 Name:Doctor Who Network:BBC Worldwide}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"gopkg.in/AlecAivazis/survey.v1"
)

// the questions to ask
var simpleQs = []*survey.Question{
	{
		Name: "name",
		Prompt: &survey.Input{
			Message: "What is your name?",
		},
		Validate: survey.Required,
	},
	{
		Name: "color",
		Prompt: &survey.Select{
			Message: "Choose a color:",
			Options: []string{"red", "blue", "green"},
		},
		Validate: survey.Required,
	},
}

func main() {
	ansmap := make(map[string]string)

	// ask the question
	err := survey.Ask(simpleQs, &ansmap)

	if err != nil {
		fmt.Println(err.Error())
		return
	}
	// print the answers
	fmt.Printf("%s chose %s.\n", ansmap["name"], ansmap["color"])
}
<commit_msg>Fixed map example (#100)<commit_after>package main

import (
	"fmt"

	"gopkg.in/AlecAivazis/survey.v1"
)

// the questions to ask
var simpleQs = []*survey.Question{
	{
		Name: "name",
		Prompt: &survey.Input{
			Message: "What is your name?",
		},
		Validate: survey.Required,
	},
	{
		Name: "color",
		Prompt: &survey.Select{
			Message: "Choose a color:",
			Options: []string{"red", "blue", "green"},
		},
		Validate: survey.Required,
	},
}

func main() {
	ansmap := make(map[string]interface{})

	// ask the question
	err := survey.Ask(simpleQs, &ansmap)

	if err != nil {
		fmt.Println(err.Error())
		return
	}
	// print the answers
	fmt.Printf("%s chose %s.\n", ansmap["name"], ansmap["color"])
}
<|endoftext|>


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

package config

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"kythe.io/kythe/go/extractors/config/default/mvn"

	ecpb "kythe.io/kythe/proto/extraction_config_go_proto"
)

// kytheConfigFileName The name of the Kythe extraction config
const kytheExtractionConfigFile = ".kythe-extraction-config"

// Repo is a container of input/output parameters for doing extraction on remote
// repositories.
type Repo struct {
	// Clone extracts a copy of the repo to the specified output Directory.
	Clone func(ctx context.Context, outputDir string) error
	// Where to write from an extraction.
	OutputPath string
	// An optional path to a file containing a
	// kythe.proto.ExtractionConfiguration encoded as JSON that details how
	// to perform extraction. If this is unset, the extractor will first try
	// to find a config defined in the repo, or finally use a hard coded
	// default config.
	ConfigPath string
}

// GitCopier returns a function that clones a repository via git command line.
func GitCopier(repoURI string) func(ctx context.Context, outputDir string) error {
	return func(ctx context.Context, outputDir string) error {
		// TODO(danielmoy): strongly consider go-git instead of os.exec
		return exec.CommandContext(ctx, "git", "clone", repoURI, outputDir).Run()
	}
}

// LocalCopier returns a function that copies a local repository.
// This function assumes the eventual output directory is already created.
func LocalCopier(repoPath string) func(ctx context.Context, outputDir string) error {
	return func(ctx context.Context, outputDir string) error {
		gitDir := filepath.Join(repoPath, ".git")
		// TODO(danielmoy): consider extracting all or part of this
		// to a more common place.
		return filepath.Walk(repoPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if repoPath == path {
				// Intentionally do nothing for base dir.
				return nil
			}
			if filepath.HasPrefix(path, gitDir) {
				return filepath.SkipDir
			}
			rel, err := filepath.Rel(repoPath, path)
			if err != nil {
				return err
			}
			outPath := filepath.Join(outputDir, rel)
			if info.Mode().IsRegular() {
				if err := os.MkdirAll(filepath.Dir(outPath), 0755); err != nil {
					return fmt.Errorf("failed to make dir: %v", err)
				}
				inf, err := os.Open(path)
				if err != nil {
					return fmt.Errorf("failed to open input file from repo: %v", err)
				}
				defer inf.Close()
				of, err := os.Create(outPath)
				if err != nil {
					return fmt.Errorf("failed to open output file for repo copy: %v", err)
				}
				if _, err := io.Copy(of, inf); err != nil {
					of.Close()
					return fmt.Errorf("failed to copy repo file: %v", err)
				}
				return of.Close()
			} else if !info.IsDir() {
				// Notably in here are any links or other odd things.
				log.Printf("Unsupported file %s with mode %s\n", path, info.Mode())
			}
			return nil
		})
	}
}

// Extractor is the interface for handling kindex generation on repos.
//
// ExtractRepo takes an input repo, output path to a directory, and optional
// kythe.proto.ExtractionConfiguration file path, and performs kythe extraction
// on the repo, depositing results in the output directory path.
type Extractor func(ctx context.Context, repo Repo) error

// ExtractRepo extracts a given code repository and outputs kindex files.
//
// It makes a local clone of the repository. It optionally uses a passed
// extraction config, otherwise it attempts to find a Kythe config named
// ".kythe-extraction-config".
//
// It builds a one-off customized Docker image for extraction, and then runs it,
// generating kindex files (format defined here:
// http://kythe.io/docs/kythe-index-pack.html).
//
// This function requires both Git and Docker to be in $PATH during execution.
func ExtractRepo(ctx context.Context, repo Repo) error {
	if err := verifyRequiredTools(); err != nil {
		return fmt.Errorf("ExtractRepo requires git and docker to be in $PATH: %v", err)
	}

	// create a temporary directory for the repo clone
	repoDir, err := ioutil.TempDir("", "repoDir")
	if err != nil {
		return fmt.Errorf("creating tmp repo dir: %v", err)
	}
	defer os.RemoveAll(repoDir)

	// create a temporary directory for the extraction output
	tmpOutDir, err := ioutil.TempDir("", "tmpOutDir")
	if err != nil {
		return fmt.Errorf("creating tmp out dir: %v", err)
	}
	defer os.RemoveAll(tmpOutDir)

	// copy the repo into our temp directory, so we can mutate its
	// build config without affecting the original source.
	if err := repo.Clone(ctx, repoDir); err != nil {
		return fmt.Errorf("copying repo: %v", err)
	}

	log.Printf("Using configuration file: %q", repo.ConfigPath)
	extractionDockerFile, err := ioutil.TempFile(tmpOutDir, "extractionDockerFile")
	if err != nil {
		return fmt.Errorf("creating tmp Dockerfile: %v", err)
	}

	// generate an extraction image from the config
	extractionConfig, err := findConfig(repo.ConfigPath, repoDir)
	if err != nil {
		return fmt.Errorf("reading config file: %v", err)
	}

	err = CreateImage(extractionDockerFile.Name(), extractionConfig)
	if err != nil {
		return fmt.Errorf("creating extraction image: %v", err)
	}

	// use Docker to build the extraction image
	imageTag := strings.ToLower(filepath.Base(extractionDockerFile.Name()))
	output, err := exec.CommandContext(ctx, "docker", "build", "-f", extractionDockerFile.Name(), "-t", imageTag, tmpOutDir).CombinedOutput()
	defer mustCleanUpImage(ctx, imageTag)
	if err != nil {
		return fmt.Errorf("building docker image: %v\nCommand output %s", err, string(output))
	}

	// run the extraction
	output, err = exec.CommandContext(ctx, "docker", "run", "--rm", "-v", fmt.Sprintf("%s:%s", repoDir, DefaultRepoVolume), "-v", fmt.Sprintf("%s:%s", repo.OutputPath, DefaultOutputVolume), "-t", imageTag).CombinedOutput()
	if err != nil {
		return fmt.Errorf("extracting repo: %v\nCommand output: %s", err, string(output))
	}

	return nil
}

func verifyRequiredTools() error {
	if _, err := exec.LookPath("git"); err != nil {
		return err
	}

	if _, err := exec.LookPath("docker"); err != nil {
		return err
	}
	return nil
}

func findConfig(configPath, repoDir string) (*ecpb.ExtractionConfiguration, error) {
	// if a config was passed in, use the specified config, otherwise go
	// hunt for one in the repository.
	if configPath == "" {
		// otherwise, use a Kythe config within the repo (if it exists)
		configPath = filepath.Join(repoDir, kytheExtractionConfigFile)
	}

	f, err := os.Open(configPath)
	if os.IsNotExist(err) {
		// TODO(danielmoy): This needs to be configurable by builder, language, etc.
		return Load(mvn.DefaultConfig())
	} else {
		return nil, fmt.Errorf("opening config file: %v", err)
	}

	defer f.Close()
	return Load(f)
}

func mustCleanUpImage(ctx context.Context, tmpImageTag string) {
	cmd := exec.CommandContext(ctx, "docker", "image", "rm", tmpImageTag)
	err := cmd.Run()
	if err != nil {
		log.Printf("Failed to clean up docker image: %v", err)
	}
}
<commit_msg>Fix bug.  I should not write code while eating and also watching TV. (#149)<commit_after>/*
 * Copyright 2018 Google Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package config

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"kythe.io/kythe/go/extractors/config/default/mvn"

	ecpb "kythe.io/kythe/proto/extraction_config_go_proto"
)

// kytheConfigFileName The name of the Kythe extraction config
const kytheExtractionConfigFile = ".kythe-extraction-config"

// Repo is a container of input/output parameters for doing extraction on remote
// repositories.
type Repo struct {
	// Clone extracts a copy of the repo to the specified output Directory.
	Clone func(ctx context.Context, outputDir string) error
	// Where to write from an extraction.
	OutputPath string
	// An optional path to a file containing a
	// kythe.proto.ExtractionConfiguration encoded as JSON that details how
	// to perform extraction. If this is unset, the extractor will first try
	// to find a config defined in the repo, or finally use a hard coded
	// default config.
	ConfigPath string
}

// GitCopier returns a function that clones a repository via git command line.
func GitCopier(repoURI string) func(ctx context.Context, outputDir string) error {
	return func(ctx context.Context, outputDir string) error {
		// TODO(danielmoy): strongly consider go-git instead of os.exec
		return exec.CommandContext(ctx, "git", "clone", repoURI, outputDir).Run()
	}
}

// LocalCopier returns a function that copies a local repository.
// This function assumes the eventual output directory is already created.
func LocalCopier(repoPath string) func(ctx context.Context, outputDir string) error {
	return func(ctx context.Context, outputDir string) error {
		gitDir := filepath.Join(repoPath, ".git")
		// TODO(danielmoy): consider extracting all or part of this
		// to a more common place.
		return filepath.Walk(repoPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if repoPath == path {
				// Intentionally do nothing for base dir.
				return nil
			}
			if filepath.HasPrefix(path, gitDir) {
				return filepath.SkipDir
			}
			rel, err := filepath.Rel(repoPath, path)
			if err != nil {
				return err
			}
			outPath := filepath.Join(outputDir, rel)
			if info.Mode().IsRegular() {
				if err := os.MkdirAll(filepath.Dir(outPath), 0755); err != nil {
					return fmt.Errorf("failed to make dir: %v", err)
				}
				inf, err := os.Open(path)
				if err != nil {
					return fmt.Errorf("failed to open input file from repo: %v", err)
				}
				defer inf.Close()
				of, err := os.Create(outPath)
				if err != nil {
					return fmt.Errorf("failed to open output file for repo copy: %v", err)
				}
				if _, err := io.Copy(of, inf); err != nil {
					of.Close()
					return fmt.Errorf("failed to copy repo file: %v", err)
				}
				return of.Close()
			} else if !info.IsDir() {
				// Notably in here are any links or other odd things.
				log.Printf("Unsupported file %s with mode %s\n", path, info.Mode())
			}
			return nil
		})
	}
}

// Extractor is the interface for handling kindex generation on repos.
//
// ExtractRepo takes an input repo, output path to a directory, and optional
// kythe.proto.ExtractionConfiguration file path, and performs kythe extraction
// on the repo, depositing results in the output directory path.
type Extractor func(ctx context.Context, repo Repo) error

// ExtractRepo extracts a given code repository and outputs kindex files.
//
// It makes a local clone of the repository. It optionally uses a passed
// extraction config, otherwise it attempts to find a Kythe config named
// ".kythe-extraction-config".
//
// It builds a one-off customized Docker image for extraction, and then runs it,
// generating kindex files (format defined here:
// http://kythe.io/docs/kythe-index-pack.html).
//
// This function requires both Git and Docker to be in $PATH during execution.
func ExtractRepo(ctx context.Context, repo Repo) error {
	if err := verifyRequiredTools(); err != nil {
		return fmt.Errorf("ExtractRepo requires git and docker to be in $PATH: %v", err)
	}

	// create a temporary directory for the repo clone
	repoDir, err := ioutil.TempDir("", "repoDir")
	if err != nil {
		return fmt.Errorf("creating tmp repo dir: %v", err)
	}
	defer os.RemoveAll(repoDir)

	// create a temporary directory for the extraction output
	tmpOutDir, err := ioutil.TempDir("", "tmpOutDir")
	if err != nil {
		return fmt.Errorf("creating tmp out dir: %v", err)
	}
	defer os.RemoveAll(tmpOutDir)

	// copy the repo into our temp directory, so we can mutate its
	// build config without affecting the original source.
	if err := repo.Clone(ctx, repoDir); err != nil {
		return fmt.Errorf("copying repo: %v", err)
	}

	log.Printf("Using configuration file: %q", repo.ConfigPath)
	extractionDockerFile, err := ioutil.TempFile(tmpOutDir, "extractionDockerFile")
	if err != nil {
		return fmt.Errorf("creating tmp Dockerfile: %v", err)
	}

	// generate an extraction image from the config
	extractionConfig, err := findConfig(repo.ConfigPath, repoDir)
	if err != nil {
		return fmt.Errorf("reading config file: %v", err)
	}

	err = CreateImage(extractionDockerFile.Name(), extractionConfig)
	if err != nil {
		return fmt.Errorf("creating extraction image: %v", err)
	}

	// use Docker to build the extraction image
	imageTag := strings.ToLower(filepath.Base(extractionDockerFile.Name()))
	output, err := exec.CommandContext(ctx, "docker", "build", "-f", extractionDockerFile.Name(), "-t", imageTag, tmpOutDir).CombinedOutput()
	defer mustCleanUpImage(ctx, imageTag)
	if err != nil {
		return fmt.Errorf("building docker image: %v\nCommand output %s", err, string(output))
	}

	// run the extraction
	output, err = exec.CommandContext(ctx, "docker", "run", "--rm", "-v", fmt.Sprintf("%s:%s", repoDir, DefaultRepoVolume), "-v", fmt.Sprintf("%s:%s", repo.OutputPath, DefaultOutputVolume), "-t", imageTag).CombinedOutput()
	if err != nil {
		return fmt.Errorf("extracting repo: %v\nCommand output: %s", err, string(output))
	}

	return nil
}

func verifyRequiredTools() error {
	if _, err := exec.LookPath("git"); err != nil {
		return err
	}

	if _, err := exec.LookPath("docker"); err != nil {
		return err
	}
	return nil
}

func findConfig(configPath, repoDir string) (*ecpb.ExtractionConfiguration, error) {
	// if a config was passed in, use the specified config, otherwise go
	// hunt for one in the repository.
	if configPath == "" {
		// otherwise, use a Kythe config within the repo (if it exists)
		configPath = filepath.Join(repoDir, kytheExtractionConfigFile)
	}

	f, err := os.Open(configPath)
	if os.IsNotExist(err) {
		// TODO(danielmoy): This needs to be configurable by builder, language, etc.
		return Load(mvn.DefaultConfig())
	} else if err != nil {
		return nil, fmt.Errorf("opening config file: %v", err)
	}

	defer f.Close()
	return Load(f)
}

func mustCleanUpImage(ctx context.Context, tmpImageTag string) {
	cmd := exec.CommandContext(ctx, "docker", "image", "rm", tmpImageTag)
	err := cmd.Run()
	if err != nil {
		log.Printf("Failed to clean up docker image: %v", err)
	}
}
<|endoftext|>