<commit_before>package upload

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/aws/signer/v4"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

type AwsUploadF interface {
	Upload(io.Reader) (*s3manager.UploadOutput, error)
}

type UploadParameters struct {
	Action         string `json:"action"`
	AwsAccessKeyId string `json:"AWSAccessKeyId"`
	ContentType    string `json:"Content-Type"`
	Policy         string `json:"policy"`
	Signature      string `json:"signature"`
	Acl            string `json:"acl"`
	Key            string `json:"key"`
	SuccessStatus  string `json:"success_action_status"`
	SignatureUrl   string `json:"signature_url"`
	VideoId        string `json:"video_id"`
	AssetId        string `json:"asset_id"`
}

type V4Request struct {
	Method   string            `json:"method"`
	Action   string            `json:"action"`
	Path     string            `json:"path"`
	RawQuery string            `json:"raw_query"`
	Headers  map[string]string `json:"headers"`
}

type V4Response struct {
	Authorization string `json:"authorization"`
	Date          string `json:"date"`
}

func CreateV4Request(action string, req *request.Request) V4Request {
	r := V4Request{}
	hreq := req.HTTPRequest
	r.Method = hreq.Method
	r.Action = action
	r.Path = hreq.URL.Path
	r.RawQuery = hreq.URL.RawQuery
	for header, _ := range req.SignedHeaderVals {
		r.Headers[header] = hreq.Header.Get(header)
	}
	return r
}

func (r *V4Request) BuildRequest() *http.Request {
	req, _ := http.NewRequest(r.Method, r.Action, nil)
	req.URL.Path = r.Path
	req.URL.RawQuery = r.RawQuery
	for header, val := range r.Headers {
		req.Header.Set(header, val)
	}
	return req
}

func (r *V4Request) Sign(awsKey, awsSecret string) (resp V4Response) {
	// use the v4 signer automatically
	provider := credentials.StaticProvider{}
	provider.Value.AccessKeyID = awsKey
	provider.Value.SecretAccessKey = awsSecret
	cred := credentials.NewCredentials(&provider)
	signer := v4.NewSigner(cred)
	req := r.BuildRequest()
	_, err := signer.Sign(req, nil, "s3", "us-east-1", time.Now())
	if err != nil {
		log.Printf("could not sign url %s", err.Error())
		return resp
	}
	date := req.Header.Get("X-Amz-Date")
	auth := req.Header.Get("Authorization")
	return V4Response{
		Authorization: auth,
		Date:          date,
	}
}

var CreatorFn func(UploadParameters) (AwsUploadF, error)

func init() {
	CreatorFn = NewAwsUpload
}

// This is the struct that contains all the AWS settings
type AwsUpload struct {
	UploadParams UploadParameters
	Uploader     *s3manager.Uploader
}

// UploadParameters is retrieved from the Unicorn API, so we're creating an AwsUpload from the settings
func NewAwsUpload(params UploadParameters) (AwsUploadF, error) {
	au := &AwsUpload{
		UploadParams: params,
	}
	provider := credentials.StaticProvider{}
	// use dummy values
	provider.Value.AccessKeyID = multipartUploadAwsAccessKeyId
	provider.Value.SecretAccessKey = multipartUploadAwsSecretAccessKey
	credentials := credentials.NewCredentials(&provider)

	region, e := au.GetRegion()
	if e != nil {
		return au, e
	}
	// session
	sess, err := session.NewSession(&aws.Config{
		Credentials: credentials,
		Region:      aws.String(region),
	})
	if err != nil {
		return au, err
	}

	svc := s3.New(sess)

	customSigner := true
	if customSigner {
		// sign handler
		signer := au.Signer()
		svc.Handlers.Sign.PushBack(signer)
	}

	// s3manager uploader
	au.Uploader = s3manager.NewUploaderWithClient(svc)
	return au, nil
}

func (a *AwsUpload) Url() string {
	return a.UploadParams.Action
}

func (a *AwsUpload) Key() string {
	return a.UploadParams.Key
}

func (a *AwsUpload) Acl() string {
	return a.UploadParams.Acl
}

func (a *AwsUpload) ContentType() string {
	return a.UploadParams.ContentType
}

func (a *AwsUpload) AwsKeyId() string {
	return a.UploadParams.AwsAccessKeyId
}

func (a *AwsUpload) UploaderSigUrl() string {
	// take the UploadParams signature and append it to the uploader url
	return a.UploadParams.SignatureUrl
}

// bucketOfUploadAction parses an "action" URL as received with GetUploadInfo,
// and returns the bucket name part of that URL.
//
// Example:
//         const a = "https://synqfm.s3.amazonaws.com"
//
//         bucket, err := bucketOfUploadAction(a)
//         if err != nil {
//                 log.Fatal(err)
//         }
//         fmt.Println(bucket) // prints synqfm
func (a AwsUpload) GetBucket() (string, error) {
	return BucketOfUploadAction(a.Url())
}

// regionOfUploadAction parses an "action" URL as received with GetUploadInfo,
// and returns the region of the bucket that is the URL refers to.
//
// This relies on heuristics, and will not work with certain styles of URLs.
//
// See: http://docs.aws.amazon.com/general/latest/gr/rande.html#s3_region
//
// Example:
//         const a = "https://synqfm.s3.amazonaws.com"
//
//         region, err := regionOfUploadAction(a)
//         if err != nil {
//                 log.Fatal(err)
//         }
//         fmt.Println(region) // prints us-east-1
func (a AwsUpload) GetRegion() (string, error) {
	return RegionOfUploadAction(a.Url())
}

func (a *AwsUpload) Upload(body io.Reader) (*s3manager.UploadOutput, error) {
	// upload parameters
	acl := a.Acl()
	bucket, err := a.GetBucket()
	if err != nil {
		return nil, err
	}
	contentType := a.ContentType()
	key := a.Key()
	uploadInput := &s3manager.UploadInput{
		ACL:         &acl,
		Bucket:      &bucket,
		Body:        body,
		ContentType: &contentType,
		Key:         &key,
	}
	return a.Uploader.Upload(uploadInput)
}

// MultipartUploadSigner returns a function that can be added to an s3 client's
// list of handlers. The function will take over signing of requests from
// aws-sdk-go.
//
// The signer function uses SYNQ's embeddable web uploader's remote procedure
// to sign requests.
//
// This function is used by the internal multipartUpload function.
//
// Example:
//         // AWS session.
//         sess := session.Must(session.NewSession())
//
//         // S3 service client.
//         svc := s3.New(sess)
//
//         // Signer function. Determine the arguments somehow.
//         signer := a.Signer()
//
//         // Register handler as the last handler of the signing phase.
//         svc.Handlers.Sign.PushBack(signer)
//
//         // S3 requests are now signed by signer().

func (a AwsUpload) Signer() func(r *request.Request) {
	signer := func(r *request.Request) {
		err := a.SignRequest(r)
		if err != nil {
			return // TODO(mastensg): how to report errors from handlers?
		}
	}

	return signer
}

func (a *AwsUpload) ServerSignV2(r *request.Request) (string, error) {
	v4 := CreateV4Request(a.UploadParams.Action, r)
	resp, err := a.V4Sig(v4)
	if err != nil {
		return "", err
	}
	if resp.Date != "" {
		// reset the data
		r.HTTPRequest.Header.Set("X-Amz-Date", resp.Date)
	}
	return resp.Authorization, nil
}

func (a *AwsUpload) ServerSignV1(r *http.Request) (string, error) {
	bucket, err := a.GetBucket()
	if err != nil {
		return "", err
	}
	if err := RewriteXAmzDateHeader(r.Header); err != nil {
		return "", err
	}
	headers := ""
	x_amz_date := r.Header.Get("X-Amz-Date")
	if r.URL.RawQuery == "uploads=" {
		// Initiate multi-part upload

		headers = fmt.Sprintf("%s\n\n%s\n\nx-amz-acl:%s\nx-amz-date:%s\n/%s%s",
			r.Method,
			a.ContentType(),
			a.Acl(),
			x_amz_date,
			bucket,
			r.URL.Path+"?uploads",
		)
	} else if r.Method == "PUT" {
		// Upload one part
		headers = fmt.Sprintf("%s\n\n%s\n\nx-amz-date:%s\n/%s%s",
			r.Method,
			"",
			x_amz_date,
			bucket,
			r.URL.Path+"?"+r.URL.RawQuery,
		)
	} else if r.Method == "POST" {
		// Finish multi-part upload

		// TODO(mastensg): the content-type header set by
		// aws-sdk-go is not exactly the one expected by
		// uploader/signature, maybe
		r.Header.Set("Content-Type", "application/xml; charset=UTF-8")

		headers = fmt.Sprintf("%s\n\n%s\n\nx-amz-date:%s\n/%s%s",
			r.Method,
			r.Header.Get("Content-Type"),
			x_amz_date,
			bucket,
			r.URL.Path+"?"+r.URL.RawQuery,
		)
	} else {
		return "", errors.New("Unknown request type.")
	}
	sig, err := a.UploaderSignature(headers)
	if err != nil {
		return "", err
	}
	delete(r.Header, "X-Amz-Content-Sha256")
	return fmt.Sprintf("AWS %s:%s", a.AwsKeyId(), sig), nil
}

// This runs as a handler within the Sign HandlerList and uses unicorn to sign the request
// This will replace whats in awssdk-go/aws/signer/v4/v4.go and its own "signWithBody" method
// https://docs.aws.amazon.com/general/latest/gr/signature-version-2.html
// https://docs.aws.amazon.com/general/latest/gr/sigv4_signing.html
// https://docs.aws.amazon.com/AmazonS3/latest/API/sig-v4-authenticating-requests.html
func (a *AwsUpload) SignRequest(r *request.Request) error {
	// construct "headers" string to send to
	// https://uploader.synq.fm/uploader/signature
	auth, err := a.ServerSignV2(r)
	if err != nil {
		return err
	}

	// rewrite authorization header(s)
	delete(r.HTTPRequest.Header, "Authorization")
	r.HTTPRequest.Header.Set("Authorization", auth)

	return nil
}

// UploaderSignature uses the backend of the embeddable web uploader to sign
// multipart upload requests.
func (a *AwsUpload) UploaderSignature(headers string) ([]byte, error) {
	reqStruct := UploaderSignatureRequest{Headers: headers}
	reqBody, err := json.Marshal(reqStruct)
	if err != nil {
		return nil, err
	}

	respBody, err := a.Request(reqBody)

	// parse response
	respStruct := UploaderSignatureResponse{}
	err = json.Unmarshal(respBody, &respStruct)
	if err != nil {
		log.Println("error unmarshaling", err.Error())
		return nil, err
	}
	return []byte(respStruct.Signature), nil
}

func (a *AwsUpload) V4Sig(req V4Request) (resp V4Response, err error) {
	reqBody, err := json.Marshal(req)
	if err != nil {
		return resp, err
	}
	respBody, err := a.Request(reqBody)
	if err != nil {
		return resp, err
	}
	err = json.Unmarshal(respBody, &resp)
	if err != nil {
		return resp, err
	}
	return resp, nil
}

func (a *AwsUpload) Request(body []byte) ([]byte, error) {
	url := a.UploaderSigUrl()

	resp, err := http.Post(url, "application/json", bytes.NewReader(body))
	if err != nil {
		log.Printf("could not call %s : %s\n", url, err.Error())
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		// TODO(mastensg): report status and maybe body
		// TODO(mastensg): handle known error responses specifically
		log.Printf("invalid response code %d from response\n", resp.StatusCode)
		return nil, errors.New("HTTP response status not OK.")
	}
	// read response
	respBody, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("error reading response body", err.Error())
		return nil, err
	}
	return respBody, nil
}
<commit_msg>if AWS_SECRET is set, just use local signing<commit_after>package upload

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/aws/signer/v4"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

type AwsUploadF interface {
	Upload(io.Reader) (*s3manager.UploadOutput, error)
}

type UploadParameters struct {
	Action         string `json:"action"`
	AwsAccessKeyId string `json:"AWSAccessKeyId"`
	ContentType    string `json:"Content-Type"`
	Policy         string `json:"policy"`
	Signature      string `json:"signature"`
	Acl            string `json:"acl"`
	Key            string `json:"key"`
	SuccessStatus  string `json:"success_action_status"`
	SignatureUrl   string `json:"signature_url"`
	VideoId        string `json:"video_id"`
	AssetId        string `json:"asset_id"`
}

type V4Request struct {
	Method   string            `json:"method"`
	Action   string            `json:"action"`
	Path     string            `json:"path"`
	RawQuery string            `json:"raw_query"`
	Headers  map[string]string `json:"headers"`
}

type V4Response struct {
	Authorization string `json:"authorization"`
	Date          string `json:"date"`
}

func CreateV4Request(action string, req *request.Request) V4Request {
	r := V4Request{}
	hreq := req.HTTPRequest
	r.Method = hreq.Method
	r.Action = action
	r.Path = hreq.URL.Path
	r.RawQuery = hreq.URL.RawQuery
	r.Headers = make(map[string]string)
	for header, _ := range req.SignedHeaderVals {
		r.Headers[header] = hreq.Header.Get(header)
	}
	return r
}

func (r *V4Request) BuildRequest() *http.Request {
	req, _ := http.NewRequest(r.Method, r.Action, nil)
	req.URL.Path = r.Path
	req.URL.RawQuery = r.RawQuery
	for header, val := range r.Headers {
		req.Header.Set(header, val)
	}
	return req
}

func (r *V4Request) Sign(awsKey, awsSecret string) (resp V4Response) {
	// use the v4 signer automatically
	provider := credentials.StaticProvider{}
	provider.Value.AccessKeyID = awsKey
	provider.Value.SecretAccessKey = awsSecret
	cred := credentials.NewCredentials(&provider)
	signer := v4.NewSigner(cred)
	req := r.BuildRequest()
	_, err := signer.Sign(req, nil, "s3", "us-east-1", time.Now())
	if err != nil {
		log.Printf("could not sign url %s", err.Error())
		return resp
	}
	date := req.Header.Get("X-Amz-Date")
	auth := req.Header.Get("Authorization")
	return V4Response{
		Authorization: auth,
		Date:          date,
	}
}

var CreatorFn func(UploadParameters) (AwsUploadF, error)

func init() {
	CreatorFn = NewAwsUpload
}

// This is the struct that contains all the AWS settings
type AwsUpload struct {
	UploadParams UploadParameters
	Uploader     *s3manager.Uploader
}

// UploadParameters is retrieved from the Unicorn API, so we're creating an AwsUpload from the settings
func NewAwsUpload(params UploadParameters) (AwsUploadF, error) {
	au := &AwsUpload{
		UploadParams: params,
	}
	provider := credentials.StaticProvider{}
	// use dummy values
	provider.Value.AccessKeyID = multipartUploadAwsAccessKeyId
	provider.Value.SecretAccessKey = multipartUploadAwsSecretAccessKey
	credentials := credentials.NewCredentials(&provider)

	region, e := au.GetRegion()
	if e != nil {
		return au, e
	}
	// session
	sess, err := session.NewSession(&aws.Config{
		Credentials: credentials,
		Region:      aws.String(region),
	})
	if err != nil {
		return au, err
	}

	svc := s3.New(sess)

	customSigner := true
	if customSigner {
		// sign handler
		signer := au.Signer()
		svc.Handlers.Sign.PushBack(signer)
	}

	// s3manager uploader
	au.Uploader = s3manager.NewUploaderWithClient(svc)
	return au, nil
}

func (a *AwsUpload) Url() string {
	return a.UploadParams.Action
}

func (a *AwsUpload) Key() string {
	return a.UploadParams.Key
}

func (a *AwsUpload) Acl() string {
	return a.UploadParams.Acl
}

func (a *AwsUpload) ContentType() string {
	return a.UploadParams.ContentType
}

func (a *AwsUpload) AwsKeyId() string {
	return a.UploadParams.AwsAccessKeyId
}

func (a *AwsUpload) UploaderSigUrl() string {
	// take the UploadParams signature and append it to the uploader url
	return a.UploadParams.SignatureUrl
}

// bucketOfUploadAction parses an "action" URL as received with GetUploadInfo,
// and returns the bucket name part of that URL.
//
// Example:
//         const a = "https://synqfm.s3.amazonaws.com"
//
//         bucket, err := bucketOfUploadAction(a)
//         if err != nil {
//                 log.Fatal(err)
//         }
//         fmt.Println(bucket) // prints synqfm
func (a AwsUpload) GetBucket() (string, error) {
	return BucketOfUploadAction(a.Url())
}

// regionOfUploadAction parses an "action" URL as received with GetUploadInfo,
// and returns the region of the bucket that is the URL refers to.
//
// This relies on heuristics, and will not work with certain styles of URLs.
//
// See: http://docs.aws.amazon.com/general/latest/gr/rande.html#s3_region
//
// Example:
//         const a = "https://synqfm.s3.amazonaws.com"
//
//         region, err := regionOfUploadAction(a)
//         if err != nil {
//                 log.Fatal(err)
//         }
//         fmt.Println(region) // prints us-east-1
func (a AwsUpload) GetRegion() (string, error) {
	return RegionOfUploadAction(a.Url())
}

func (a *AwsUpload) Upload(body io.Reader) (*s3manager.UploadOutput, error) {
	// upload parameters
	acl := a.Acl()
	bucket, err := a.GetBucket()
	if err != nil {
		return nil, err
	}
	contentType := a.ContentType()
	key := a.Key()
	uploadInput := &s3manager.UploadInput{
		ACL:         &acl,
		Bucket:      &bucket,
		Body:        body,
		ContentType: &contentType,
		Key:         &key,
	}
	return a.Uploader.Upload(uploadInput)
}

// MultipartUploadSigner returns a function that can be added to an s3 client's
// list of handlers. The function will take over signing of requests from
// aws-sdk-go.
//
// The signer function uses SYNQ's embeddable web uploader's remote procedure
// to sign requests.
//
// This function is used by the internal multipartUpload function.
//
// Example:
//         // AWS session.
//         sess := session.Must(session.NewSession())
//
//         // S3 service client.
//         svc := s3.New(sess)
//
//         // Signer function. Determine the arguments somehow.
//         signer := a.Signer()
//
//         // Register handler as the last handler of the signing phase.
//         svc.Handlers.Sign.PushBack(signer)
//
//         // S3 requests are now signed by signer().

func (a AwsUpload) Signer() func(r *request.Request) {
	signer := func(r *request.Request) {
		err := a.SignRequest(r)
		if err != nil {
			return // TODO(mastensg): how to report errors from handlers?
		}
	}

	return signer
}

func (a *AwsUpload) ServerSignV2(r *request.Request) (string, error) {
	v4 := CreateV4Request(a.UploadParams.Action, r)
	resp, err := a.V4Sig(v4)
	if err != nil {
		return "", err
	}
	if resp.Date != "" {
		// reset the data
		r.HTTPRequest.Header.Set("X-Amz-Date", resp.Date)
	}
	return resp.Authorization, nil
}

func (a *AwsUpload) ServerSignV1(r *http.Request) (string, error) {
	bucket, err := a.GetBucket()
	if err != nil {
		return "", err
	}
	if err := RewriteXAmzDateHeader(r.Header); err != nil {
		return "", err
	}
	headers := ""
	x_amz_date := r.Header.Get("X-Amz-Date")
	if r.URL.RawQuery == "uploads=" {
		// Initiate multi-part upload

		headers = fmt.Sprintf("%s\n\n%s\n\nx-amz-acl:%s\nx-amz-date:%s\n/%s%s",
			r.Method,
			a.ContentType(),
			a.Acl(),
			x_amz_date,
			bucket,
			r.URL.Path+"?uploads",
		)
	} else if r.Method == "PUT" {
		// Upload one part
		headers = fmt.Sprintf("%s\n\n%s\n\nx-amz-date:%s\n/%s%s",
			r.Method,
			"",
			x_amz_date,
			bucket,
			r.URL.Path+"?"+r.URL.RawQuery,
		)
	} else if r.Method == "POST" {
		// Finish multi-part upload

		// TODO(mastensg): the content-type header set by
		// aws-sdk-go is not exactly the one expected by
		// uploader/signature, maybe
		r.Header.Set("Content-Type", "application/xml; charset=UTF-8")

		headers = fmt.Sprintf("%s\n\n%s\n\nx-amz-date:%s\n/%s%s",
			r.Method,
			r.Header.Get("Content-Type"),
			x_amz_date,
			bucket,
			r.URL.Path+"?"+r.URL.RawQuery,
		)
	} else {
		return "", errors.New("Unknown request type.")
	}
	sig, err := a.UploaderSignature(headers)
	if err != nil {
		return "", err
	}
	delete(r.Header, "X-Amz-Content-Sha256")
	return fmt.Sprintf("AWS %s:%s", a.AwsKeyId(), sig), nil
}

// This runs as a handler within the Sign HandlerList and uses unicorn to sign the request
// This will replace whats in awssdk-go/aws/signer/v4/v4.go and its own "signWithBody" method
// https://docs.aws.amazon.com/general/latest/gr/signature-version-2.html
// https://docs.aws.amazon.com/general/latest/gr/sigv4_signing.html
// https://docs.aws.amazon.com/AmazonS3/latest/API/sig-v4-authenticating-requests.html
func (a *AwsUpload) SignRequest(r *request.Request) error {
	// construct "headers" string to send to
	// https://uploader.synq.fm/uploader/signature
	auth, err := a.ServerSignV2(r)
	if err != nil {
		return err
	}

	// rewrite authorization header(s)
	delete(r.HTTPRequest.Header, "Authorization")
	r.HTTPRequest.Header.Set("Authorization", auth)

	return nil
}

// UploaderSignature uses the backend of the embeddable web uploader to sign
// multipart upload requests.
func (a *AwsUpload) UploaderSignature(headers string) ([]byte, error) {
	reqStruct := UploaderSignatureRequest{Headers: headers}
	reqBody, err := json.Marshal(reqStruct)
	if err != nil {
		return nil, err
	}

	respBody, err := a.Request(reqBody)

	// parse response
	respStruct := UploaderSignatureResponse{}
	err = json.Unmarshal(respBody, &respStruct)
	if err != nil {
		log.Println("error unmarshaling", err.Error())
		return nil, err
	}
	return []byte(respStruct.Signature), nil
}

func (a *AwsUpload) V4Sig(req V4Request) (resp V4Response, err error) {
	secret := os.Getenv("AWS_SECRET")
	if secret != "" {
		resp = req.Sign(a.UploadParams.AwsAccessKeyId, secret)
		return resp, nil
	}
	reqBody, err := json.Marshal(req)
	if err != nil {
		return resp, err
	}
	respBody, err := a.Request(reqBody)
	if err != nil {
		return resp, err
	}
	err = json.Unmarshal(respBody, &resp)
	if err != nil {
		return resp, err
	}
	return resp, nil
}

func (a *AwsUpload) Request(body []byte) ([]byte, error) {
	url := a.UploaderSigUrl()

	resp, err := http.Post(url, "application/json", bytes.NewReader(body))
	if err != nil {
		log.Printf("could not call %s : %s\n", url, err.Error())
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		// TODO(mastensg): report status and maybe body
		// TODO(mastensg): handle known error responses specifically
		log.Printf("invalid response code %d from response\n", resp.StatusCode)
		return nil, errors.New("HTTP response status not OK.")
	}
	// read response
	respBody, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("error reading response body", err.Error())
		return nil, err
	}
	return respBody, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"strings"

	"golang.org/x/net/html"
)

func init() {
	log.SetFlags(0)
	log.SetPrefix("elembyid: ")
}

func main() {
	url := flag.String("url", "", "URL for lookup.")
	id := flag.String("id", "", "Element id for lookup")
	flag.Parse()

	if *url == "" {
		log.Fatalln("please specify url")
	}
	if *id == "" {
		log.Fatalln("please specify element id")
	}

	resp, err := http.Get(*url)
	if err != nil {
		log.Fatalln(err)
	}
	defer resp.Body.Close()

	doc, err := html.Parse(resp.Body)
	if err != nil {
		log.Fatalln(err)
	}

	if n := elementByID(doc, *id); n != nil {
		fmt.Printf("<%s %s>\n", n.Data, joinAttrs(n.Attr))
	} else {
		fmt.Printf("no element with id: '%s'\n", *id)
	}
}

func joinAttrs(attrs []html.Attribute) string {
	asHTML := make([]string, 0, len(attrs))
	for _, a := range attrs {
		asHTML = append(asHTML, fmt.Sprintf("%s=\"%s\"", a.Key, a.Val))
	}
	return strings.Join(asHTML, " ")
}

func checkID(n *html.Node, id string) bool {
	if n == nil {
		return false
	}
	for _, a := range n.Attr {
		if a.Key == "id" && a.Val == id {
			return true
		}
	}
	return false
}

func startElement(n *html.Node, id string) bool {
	return checkID(n, id)
}

func endElement(n *html.Node, id string) bool {
	return checkID(n, id)
}

func elementByID(n *html.Node, id string) *html.Node {
	if startElement(n, id) {
		return n
	}

	for c := n.FirstChild; c != nil; c = c.NextSibling {
		if res := elementByID(c, id); res != nil {
			return res
		}
	}

	if endElement(n, id) {
		return n
	}

	return nil
}
<commit_msg>[5.8] Updated result for failed look up.<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"strings"

	"golang.org/x/net/html"
)

func init() {
	log.SetFlags(0)
	log.SetPrefix("elembyid: ")
}

func main() {
	url := flag.String("url", "", "URL for lookup.")
	id := flag.String("id", "", "Element id for lookup")
	flag.Parse()

	if *url == "" {
		log.Fatalln("please specify url")
	}
	if *id == "" {
		log.Fatalln("please specify element id")
	}

	resp, err := http.Get(*url)
	if err != nil {
		log.Fatalln(err)
	}
	defer resp.Body.Close()

	doc, err := html.Parse(resp.Body)
	if err != nil {
		log.Fatalln(err)
	}

	if n := elementByID(doc, *id); n != nil {
		fmt.Printf("<%s %s>\n", n.Data, joinAttrs(n.Attr))
	} else {
		log.Fatalf("no element with id: '%s'", *id)
	}
}

func joinAttrs(attrs []html.Attribute) string {
	asHTML := make([]string, 0, len(attrs))
	for _, a := range attrs {
		asHTML = append(asHTML, fmt.Sprintf("%s=\"%s\"", a.Key, a.Val))
	}
	return strings.Join(asHTML, " ")
}

func checkID(n *html.Node, id string) bool {
	if n == nil {
		return false
	}
	for _, a := range n.Attr {
		if a.Key == "id" && a.Val == id {
			return true
		}
	}
	return false
}

func startElement(n *html.Node, id string) bool {
	return checkID(n, id)
}

func endElement(n *html.Node, id string) bool {
	return checkID(n, id)
}

func elementByID(n *html.Node, id string) *html.Node {
	if startElement(n, id) {
		return n
	}

	for c := n.FirstChild; c != nil; c = c.NextSibling {
		if res := elementByID(c, id); res != nil {
			return res
		}
	}

	if endElement(n, id) {
		return n
	}

	return nil
}
<|endoftext|>


<commit_before>package siteengines

// TODO: Check for the visual likeness of two usernames when checking for availability! Generate images and compare pixels.
// TODO: Logging in should work case sensitively, or at least without concern for the case of the first letter
// TODO: Consider using "0" and "1" instead of "true" or "false" when setting values, while still understanding "true" or "false"

import (
	"crypto/sha256"
	"io"
	"math/rand"
	"strings"
	"time"

	. "github.com/xyproto/browserspeak"
	. "github.com/xyproto/genericsite"
	"github.com/xyproto/instapage"
	"github.com/xyproto/web"
)

// An Engine is a specific piece of a website
// This part handles the login/logout/registration/confirmation pages

const (
	MINIMUM_CONFIRMATION_CODE_LENGTH = 20
	USERNAME_ALLOWED_LETTERS         = "abcdefghijklmnopqrstuvwxyzæøåABCDEFGHIJKLMNOPQRSTUVWXYZÆØÅ_0123456789"
)

type UserEngine struct {
	state *UserState
}

func NewUserEngine(userState *UserState) *UserEngine {
	// For the secure cookies
	// This must happen before the random seeding, or
	// else people will have to log in again after every server restart
	web.Config.CookieSecret = RandomCookieFriendlyString(30)

	rand.Seed(time.Now().UnixNano())

	return &UserEngine{userState}
}

func (ue *UserEngine) GetState() *UserState {
	return ue.state
}

func CorrectPassword(state *UserState, username, password string) bool {
	passwordHash, err := state.GetPasswordHash(username)
	if err != nil {
		return false
	}
	if passwordHash == HashPasswordVersion2(password) {
		return true
	}
	return false
}

// Goes through all the confirmationCodes of all the unconfirmed users
// and checks if this confirmationCode already is in use
func AlreadyHasConfirmationCode(state *UserState, confirmationCode string) bool {
	unconfirmedUsernames, err := state.GetAllUnconfirmedUsernames()
	if err != nil {
		return false
	}
	for _, aUsername := range unconfirmedUsernames {
		aConfirmationCode, err := state.GetConfirmationCode(aUsername)
		if err != nil {
			panic("ERROR: Inconsistent user")
		}
		if confirmationCode == aConfirmationCode {
			// Found it
			return true
		}
	}
	return false
}

// Create a user by adding the username to the list of usernames
func GenerateConfirmUser(state *UserState) WebHandle {
	return func(ctx *web.Context, val string) string {
		confirmationCode := val

		unconfirmedUsernames, err := state.GetAllUnconfirmedUsernames()
		if err != nil {
			return instapage.MessageOKurl("Confirmation", "All users are confirmed already.", "/register")
		}

		// Find the username by looking up the confirmationCode on unconfirmed users
		username := ""
		for _, aUsername := range unconfirmedUsernames {
			aConfirmationCode, err := state.GetConfirmationCode(username)
			if err != nil {
				panic("ERROR: Inconsistent user!")
			}
			if confirmationCode == aConfirmationCode {
				// Found the right user
				username = aUsername
				break
			}
		}

		// Check that the user is there
		if username == "" {
			// Say "no longer" because we don't care about people that just try random confirmation links
			return instapage.MessageOKurl("Confirmation", "The confirmation link is no longer valid.", "/register")
		}
		hasUser := state.HasUser(username)
		if !hasUser {
			return instapage.MessageOKurl("Confirmation", "The user you wish to confirm does not exist anymore.", "/register")
		}

		// Remove from the list of unconfirmed usernames
		state.RemoveUnconfirmed(username)

		// Mark user as confirmed
		state.MarkConfirmed(username)

		return instapage.MessageOKurl("Confirmation", "Thank you "+username+", you can now log in.", "/login")
	}
}

// Log in a user by changing the loggedin value
func GenerateLoginUser(state *UserState) WebHandle {
	return func(ctx *web.Context, val string) string {
		// Fetch password from ctx
		password, found := ctx.Params["password"]
		if !found {
			return instapage.MessageOKback("Login", "Can't log in without a password.")
		}
		username := val
		if username == "" {
			return instapage.MessageOKback("Login", "Can't log in with a blank username.")
		}
		if !state.HasUser(username) {
			return instapage.MessageOKback("Login", "User "+username+" does not exist, could not log in.")
		}
		if !state.IsConfirmed(username) {
			return instapage.MessageOKback("Login", "The email for "+username+" has not been confirmed, check your email and follow the link.")
		}
		if !CorrectPassword(state, username, password) {
			return instapage.MessageOKback("Login", "Wrong password.")
		}

		// Log in the user by changing the database and setting a secure cookie
		state.SetLoggedIn(username)

		state.SetBrowserUsername(ctx, username)

		// TODO: Use a welcoming messageOK where the user can see when he/she last logged in and from which host

		if username == "admin" {
			ctx.SetHeader("Refresh", "0; url=/admin", true)
		} else {
			// TODO: Redirect to the page the user was at before logging in
			ctx.SetHeader("Refresh", "0; url=/", true)
		}

		return ""
	}
}

func HashPasswordVersion2(password string) string {
	hasher := sha256.New()
	// TODO: Read up on password hashing
	io.WriteString(hasher, password+"some salt is better than none")
	return string(hasher.Sum(nil))
}

// TODO: Forgot username? Enter email, send username.
// TODO: Lost confirmation link? Enter mail, Receive confirmation link.
// TODO: Forgot password? Enter mail, receive reset-password link.
// TODO: Make sure not two usernames can register at once before confirming
// TODO: Only one username per email address? (meh? can use more than one address?=
// TODO: Maximum 1 confirmation email per email adress
// TODO: Maximum 1 forgot username per email adress per day
// TODO: Maximum 1 forgot password per email adress per day
// TODO: Maximum 1 lost confirmation link per email adress per day
// TODO: Link for "Did you not request this email? Click here" i alle eposter som sendes.
// TODO: Rate limiting, maximum rate per minute or day

// Register a new user
func GenerateRegisterUser(state *UserState) WebHandle {
	return func(ctx *web.Context, val string) string {

		// Password checks
		password1, found := ctx.Params["password1"]
		if password1 == "" || !found {
			return instapage.MessageOKback("Register", "Can't register without a password.")
		}
		password2, found := ctx.Params["password2"]
		if password2 == "" || !found {
			return instapage.MessageOKback("Register", "Please confirm the password by typing it in twice.")
		}
		if password1 != password2 {
			return instapage.MessageOKback("Register", "The password and confirmation password must be equal.")
		}

		// Email checks
		email, found := ctx.Params["email"]
		if !found {
			return instapage.MessageOKback("Register", "Can't register without an email address.")
		}
		// must have @ and ., but no " "
		if !strings.Contains(email, "@") || !strings.Contains(email, ".") || strings.Contains(email, " ") {
			return instapage.MessageOKback("Register", "Please use a valid email address.")
		}

		// Username checks
		username := val
		if username == "" {
			return instapage.MessageOKback("Register", "Can't register without a username.")
		}
		if state.HasUser(username) {
			return instapage.MessageOKback("Register", "That user already exists, try another username.")
		}

		// Only some letters are allowed in the username
	NEXT:
		for _, letter := range username {
			for _, allowedLetter := range USERNAME_ALLOWED_LETTERS {
				if letter == allowedLetter {
					continue NEXT
				}
			}
			return instapage.MessageOKback("Register", "Only a-å, A-Å, 0-9 and _ are allowed in usernames.")
		}
		if username == password1 {
			return instapage.MessageOKback("Register", "Username and password must be different, try another password.")
		}

		adminuser := false
		// A special case
		if username == "admin" {
			// The first user to register with the username "admin" becomes the administrator
			adminuser = true
		}

		// Register the user
		passwordHash := HashPasswordVersion2(password1)
		state.AddUserUnchecked(username, passwordHash, email)

		// Mark user as administrator if that is the case
		if adminuser {
			// Set admin status
			state.SetAdminStatus(username)
		}

		// The confirmation code must be a minimum of 8 letters long
		length := MINIMUM_CONFIRMATION_CODE_LENGTH
		confirmationCode := RandomHumanFriendlyString(length)
		for AlreadyHasConfirmationCode(state, confirmationCode) {
			// Increase the length of the confirmationCode random string every time there is a collision
			length++
			confirmationCode = RandomHumanFriendlyString(length)
			if length > 100 {
				// This should never happen
				panic("ERROR: Too many generated confirmation codes are not unique, something is wrong")
			}
		}

		// Send confirmation email
		ConfirmationEmail("archlinux.no", "https://archlinux.no/confirm/"+confirmationCode, username, email)

		// Register the need to be confirmed
		state.AddUnconfirmed(username, confirmationCode)

		// Redirect
		return instapage.MessageOKurl("Registration complete", "Thanks for registering, the confirmation e-mail has been sent.", "/login")
	}
}

// Log out a user by changing the loggedin value
func GenerateLogoutCurrentUser(state *UserState) SimpleContextHandle {
	return func(ctx *web.Context) string {
		username := GetBrowserUsername(ctx)
		if username == "" {
			return instapage.MessageOKback("Logout", "No user to log out")
		}
		if !state.HasUser(username) {
			return instapage.MessageOKback("Logout", "user "+username+" does not exist, could not log out")
		}

		// Log out the user by changing the database, the cookie can stay
		state.SetLoggedOut(username)

		// Redirect
		//ctx.SetHeader("Refresh", "0; url=/login", true)
		return instapage.MessageOKurl("Logout", username+" is now logged out. Hope to see you soon!", "/login")
	}
}

func GenerateNoJavascriptMessage() SimpleContextHandle {
	return func(ctx *web.Context) string {
		return instapage.MessageOKback("JavaScript error", "Cookies and Javascriåt must be enabled.<br />Older browsers might be supported in the future.")
	}
}

func LoginCP(basecp BaseCP, state *UserState, url string) *ContentPage {
	cp := basecp(state)
	cp.ContentTitle = "Login"
	cp.ContentHTML = instapage.LoginForm()
	cp.ContentJS += OnClick("#loginButton", "$('#loginForm').get(0).setAttribute('action', '/login/' + $('#username').val());")
	//cp.ExtraCSSurls = append(cp.ExtraCSSurls, "/css/login.css")
	cp.Url = url

	// Hide the Login menu if we're on the Login page
	//cp.HeaderJS = strings.Replace(cp.HeaderJS, "menuLogin", "menuNop", 1)
	//cp.ContentJS += Hide("#menuLogin")

	return cp
}

func RegisterCP(basecp BaseCP, state *UserState, url string) *ContentPage {
	cp := basecp(state)
	cp.ContentTitle = "Register"
	cp.ContentHTML = instapage.RegisterForm()
	cp.ContentJS += OnClick("#registerButton", "$('#registerForm').get(0).setAttribute('action', '/register/' + $('#username').val());")
	//cp.ExtraCSSurls = append(cp.ExtraCSSurls, "/css/register.css")
	cp.Url = url

	// Hide the Register menu if we're on the Register page
	//cp.HeaderJS = strings.Replace(cp.HeaderJS, "menuRegister", "menuNop", 1)
	//cp.ContentJS += Hide("#menuRegister")

	return cp
}

func (ue *UserEngine) ServeSystem() {
	state := ue.state
	web.Post("/register/(.*)", GenerateRegisterUser(state))
	web.Post("/login/(.*)", GenerateLoginUser(state))
	web.Post("/login", GenerateNoJavascriptMessage())
	web.Get("/logout", GenerateLogoutCurrentUser(state))
	web.Get("/confirm/(.*)", GenerateConfirmUser(state))
}
<commit_msg>Refactoring<commit_after>package siteengines

// TODO: Check for the visual likeness of two usernames when checking for availability! Generate images and compare pixels.
// TODO: Logging in should work case sensitively, or at least without concern for the case of the first letter
// TODO: Consider using "0" and "1" instead of "true" or "false" when setting values, while still understanding "true" or "false"

import (
	"crypto/sha256"
	"io"
	"math/rand"
	"strings"
	"time"

	. "github.com/xyproto/browserspeak"
	. "github.com/xyproto/genericsite"
	"github.com/xyproto/instapage"
	"github.com/xyproto/web"
)

// An Engine is a specific piece of a website
// This part handles the login/logout/registration/confirmation pages

const (
	MINIMUM_CONFIRMATION_CODE_LENGTH = 20
	USERNAME_ALLOWED_LETTERS         = "abcdefghijklmnopqrstuvwxyzæøåABCDEFGHIJKLMNOPQRSTUVWXYZÆØÅ_0123456789"
)

type UserEngine struct {
	state *UserState
}

func NewUserEngine(userState *UserState) *UserEngine {
	// For the secure cookies
	// This must happen before the random seeding, or
	// else people will have to log in again after every server restart
	web.Config.CookieSecret = RandomCookieFriendlyString(30)

	rand.Seed(time.Now().UnixNano())

	return &UserEngine{userState}
}

func (ue *UserEngine) GetState() *UserState {
	return ue.state
}

func CorrectPassword(state *UserState, username, password string) bool {
	passwordHash, err := state.GetPasswordHash(username)
	if err != nil {
		return false
	}
	if passwordHash == HashPasswordVersion2(password) {
		return true
	}
	return false
}

// Goes through all the confirmationCodes of all the unconfirmed users
// and checks if this confirmationCode already is in use
func AlreadyHasConfirmationCode(state *UserState, confirmationCode string) bool {
	unconfirmedUsernames, err := state.GetAllUnconfirmedUsernames()
	if err != nil {
		return false
	}
	for _, aUsername := range unconfirmedUsernames {
		aConfirmationCode, err := state.GetConfirmationCode(aUsername)
		if err != nil {
			panic("ERROR: Inconsistent user")
		}
		if confirmationCode == aConfirmationCode {
			// Found it
			return true
		}
	}
	return false
}

// Create a user by adding the username to the list of usernames
func GenerateConfirmUser(state *UserState) WebHandle {
	return func(ctx *web.Context, val string) string {
		confirmationCode := val

		unconfirmedUsernames, err := state.GetAllUnconfirmedUsernames()
		if err != nil {
			return instapage.MessageOKurl("Confirmation", "All users are confirmed already.", "/register")
		}

		// Find the username by looking up the confirmationCode on unconfirmed users
		username := ""
		for _, aUsername := range unconfirmedUsernames {
			aConfirmationCode, err := state.GetConfirmationCode(username)
			if err != nil {
				panic("ERROR: Inconsistent user!")
			}
			if confirmationCode == aConfirmationCode {
				// Found the right user
				username = aUsername
				break
			}
		}

		// Check that the user is there
		if username == "" {
			// Say "no longer" because we don't care about people that just try random confirmation links
			return instapage.MessageOKurl("Confirmation", "The confirmation link is no longer valid.", "/register")
		}
		hasUser := state.HasUser(username)
		if !hasUser {
			return instapage.MessageOKurl("Confirmation", "The user you wish to confirm does not exist anymore.", "/register")
		}

		// Remove from the list of unconfirmed usernames
		state.RemoveUnconfirmed(username)

		// Mark user as confirmed
		state.MarkConfirmed(username)

		return instapage.MessageOKurl("Confirmation", "Thank you "+username+", you can now log in.", "/login")
	}
}

// Log in a user by changing the loggedin value
func GenerateLoginUser(state *UserState) WebHandle {
	return func(ctx *web.Context, val string) string {
		// Fetch password from ctx
		password, found := ctx.Params["password"]
		if !found {
			return instapage.MessageOKback("Login", "Can't log in without a password.")
		}
		username := val
		if username == "" {
			return instapage.MessageOKback("Login", "Can't log in with a blank username.")
		}
		if !state.HasUser(username) {
			return instapage.MessageOKback("Login", "User "+username+" does not exist, could not log in.")
		}
		if !state.IsConfirmed(username) {
			return instapage.MessageOKback("Login", "The email for "+username+" has not been confirmed, check your email and follow the link.")
		}
		if !CorrectPassword(state, username, password) {
			return instapage.MessageOKback("Login", "Wrong password.")
		}

		// Log in the user by changing the database and setting a secure cookie
		state.SetLoggedIn(username)

		state.SetBrowserUsername(ctx, username)

		// TODO: Use a welcoming messageOK where the user can see when he/she last logged in and from which host

		if username == "admin" {
			ctx.SetHeader("Refresh", "0; url=/admin", true)
		} else {
			// TODO: Redirect to the page the user was at before logging in
			ctx.SetHeader("Refresh", "0; url=/", true)
		}

		return ""
	}
}

func HashPasswordVersion2(password string) string {
	hasher := sha256.New()
	// TODO: Read up on password hashing
	io.WriteString(hasher, password+"some salt is better than none")
	return string(hasher.Sum(nil))
}

// TODO: Forgot username? Enter email, send username.
// TODO: Lost confirmation link? Enter mail, Receive confirmation link.
// TODO: Forgot password? Enter mail, receive reset-password link.
// TODO: Make sure not two usernames can register at once before confirming
// TODO: Only one username per email address? (meh? can use more than one address?=
// TODO: Maximum 1 confirmation email per email adress
// TODO: Maximum 1 forgot username per email adress per day
// TODO: Maximum 1 forgot password per email adress per day
// TODO: Maximum 1 lost confirmation link per email adress per day
// TODO: Link for "Did you not request this email? Click here" i alle eposter som sendes.
// TODO: Rate limiting, maximum rate per minute or day

// Register a new user
func GenerateRegisterUser(state *UserState) WebHandle {
	return func(ctx *web.Context, val string) string {

		// Password checks
		password1, found := ctx.Params["password1"]
		if password1 == "" || !found {
			return instapage.MessageOKback("Register", "Can't register without a password.")
		}
		password2, found := ctx.Params["password2"]
		if password2 == "" || !found {
			return instapage.MessageOKback("Register", "Please confirm the password by typing it in twice.")
		}
		if password1 != password2 {
			return instapage.MessageOKback("Register", "The password and confirmation password must be equal.")
		}

		// Email checks
		email, found := ctx.Params["email"]
		if !found {
			return instapage.MessageOKback("Register", "Can't register without an email address.")
		}
		// must have @ and ., but no " "
		if !strings.Contains(email, "@") || !strings.Contains(email, ".") || strings.Contains(email, " ") {
			return instapage.MessageOKback("Register", "Please use a valid email address.")
		}

		// Username checks
		username := val
		if username == "" {
			return instapage.MessageOKback("Register", "Can't register without a username.")
		}
		if state.HasUser(username) {
			return instapage.MessageOKback("Register", "That user already exists, try another username.")
		}

		// Only some letters are allowed in the username
	NEXT:
		for _, letter := range username {
			for _, allowedLetter := range USERNAME_ALLOWED_LETTERS {
				if letter == allowedLetter {
					continue NEXT
				}
			}
			return instapage.MessageOKback("Register", "Only a-å, A-Å, 0-9 and _ are allowed in usernames.")
		}
		if username == password1 {
			return instapage.MessageOKback("Register", "Username and password must be different, try another password.")
		}

		adminuser := false
		// A special case
		if username == "admin" {
			// The first user to register with the username "admin" becomes the administrator
			adminuser = true
		}

		// Register the user
		passwordHash := HashPasswordVersion2(password1)
		state.AddUserUnchecked(username, passwordHash, email)

		// Mark user as administrator if that is the case
		if adminuser {
			// Set admin status
			state.SetAdminStatus(username)
		}

		// The confirmation code must be a minimum of 8 letters long
		length := MINIMUM_CONFIRMATION_CODE_LENGTH
		confirmationCode := RandomHumanFriendlyString(length)
		for AlreadyHasConfirmationCode(state, confirmationCode) {
			// Increase the length of the confirmationCode random string every time there is a collision
			length++
			confirmationCode = RandomHumanFriendlyString(length)
			if length > 100 {
				// This should never happen
				panic("ERROR: Too many generated confirmation codes are not unique, something is wrong")
			}
		}

		// Send confirmation email
		ConfirmationEmail("archlinux.no", "https://archlinux.no/confirm/"+confirmationCode, username, email)

		// Register the need to be confirmed
		state.AddUnconfirmed(username, confirmationCode)

		// Redirect
		return instapage.MessageOKurl("Registration complete", "Thanks for registering, the confirmation e-mail has been sent.", "/login")
	}
}

// Log out a user by changing the loggedin value
func GenerateLogoutCurrentUser(state *UserState) SimpleContextHandle {
	return func(ctx *web.Context) string {
		username := GetBrowserUsername(ctx)
		if username == "" {
			return instapage.MessageOKback("Logout", "No user to log out")
		}
		if !state.HasUser(username) {
			return instapage.MessageOKback("Logout", "user "+username+" does not exist, could not log out")
		}

		// Log out the user by changing the database, the cookie can stay
		state.SetLoggedOut(username)

		// Redirect
		//ctx.SetHeader("Refresh", "0; url=/login", true)
		return instapage.MessageOKurl("Logout", username+" is now logged out. Hope to see you soon!", "/login")
	}
}

func GenerateNoJavascriptMessage() SimpleContextHandle {
	return func(ctx *web.Context) string {
		return instapage.MessageOKback("JavaScript error", "Cookies and Javascriåt must be enabled.<br />Older browsers might be supported in the future.")
	}
}

func LoginCP(basecp BaseCP, state *UserState, url string) *ContentPage {
	cp := basecp(state)
	cp.ContentTitle = "Login"
	cp.ContentHTML = instapage.LoginForm()
	cp.ContentJS += OnClick("#loginButton", "$('#loginForm').get(0).setAttribute('action', '/login/' + $('#username').val());")
	//cp.ExtraCSSurls = append(cp.ExtraCSSurls, "/css/login.css")
	cp.Url = url

	// Hide the Login menu if we're on the Login page
	//cp.HeaderJS = strings.Replace(cp.HeaderJS, "menuLogin", "menuNop", 1)
	//cp.ContentJS += Hide("#menuLogin")

	return cp
}

func RegisterCP(basecp BaseCP, state *UserState, url string) *ContentPage {
	cp := basecp(state)
	cp.ContentTitle = "Register"
	cp.ContentHTML = instapage.RegisterForm()
	cp.ContentJS += OnClick("#registerButton", "$('#registerForm').get(0).setAttribute('action', '/register/' + $('#username').val());")
	//cp.ExtraCSSurls = append(cp.ExtraCSSurls, "/css/register.css")
	cp.Url = url

	// Hide the Register menu if we're on the Register page
	//cp.HeaderJS = strings.Replace(cp.HeaderJS, "menuRegister", "menuNop", 1)
	//cp.ContentJS += Hide("#menuRegister")

	return cp
}

func (ue *UserEngine) ServePages() {
	state := ue.state
	web.Post("/register/(.*)", GenerateRegisterUser(state))
	web.Post("/login/(.*)", GenerateLoginUser(state))
	web.Post("/login", GenerateNoJavascriptMessage())
	web.Get("/logout", GenerateLogoutCurrentUser(state))
	web.Get("/confirm/(.*)", GenerateConfirmUser(state))
}
<|endoftext|>


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

package loader

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrrpcclient"
	"github.com/decred/dcrwallet/ticketbuyer"
	"github.com/decred/dcrwallet/wallet"
	"github.com/decred/dcrwallet/walletdb"
	_ "github.com/decred/dcrwallet/walletdb/bdb" // driver loaded during init
)

const (
	walletDbName = "wallet.db"
)

// Loader implements the creating of new and opening of existing wallets, while
// providing a callback system for other subsystems to handle the loading of a
// wallet.  This is primarely intended for use by the RPC servers, to enable
// methods and services which require the wallet when the wallet is loaded by
// another subsystem.
//
// Loader is safe for concurrent access.
type Loader struct {
	callbacks   []func(*wallet.Wallet)
	chainClient *dcrrpcclient.Client
	chainParams *chaincfg.Params
	dbDirPath   string
	wallet      *wallet.Wallet
	db          walletdb.DB
	mu          sync.Mutex

	purchaseManager *ticketbuyer.PurchaseManager
	ntfnClient      wallet.MainTipChangedNotificationsClient
	stakeOptions    *StakeOptions
	addrIdxScanLen  int
	allowHighFees   bool
	relayFee        float64
}

// StakeOptions contains the various options necessary for stake mining.
type StakeOptions struct {
	TicketPurchasingEnabled bool
	VotingEnabled           bool
	TicketFee               float64
	AddressReuse            bool
	TicketAddress           string
	PoolAddress             string
	PoolFees                float64
	StakePoolColdExtKey     string
}

// NewLoader constructs a Loader.
func NewLoader(chainParams *chaincfg.Params, dbDirPath string, stakeOptions *StakeOptions, addrIdxScanLen int,
	allowHighFees bool, relayFee float64) *Loader {

	return &Loader{
		chainParams:    chainParams,
		dbDirPath:      dbDirPath,
		stakeOptions:   stakeOptions,
		addrIdxScanLen: addrIdxScanLen,
		allowHighFees:  allowHighFees,
		relayFee:       relayFee,
	}
}

// onLoaded executes each added callback and prevents loader from loading any
// additional wallets.  Requires mutex to be locked.
func (l *Loader) onLoaded(w *wallet.Wallet, db walletdb.DB) {
	for _, fn := range l.callbacks {
		fn(w)
	}

	l.wallet = w
	l.db = db
	l.callbacks = nil // not needed anymore
}

// RunAfterLoad adds a function to be executed when the loader creates or opens
// a wallet.  Functions are executed in a single goroutine in the order they are
// added.
func (l *Loader) RunAfterLoad(fn func(*wallet.Wallet)) {
	l.mu.Lock()
	if l.wallet != nil {
		w := l.wallet
		l.mu.Unlock()
		fn(w)
	} else {
		l.callbacks = append(l.callbacks, fn)
		l.mu.Unlock()
	}
}

// CreateNewWallet creates a new wallet using the provided public and private
// passphrases.  The seed is optional.  If non-nil, addresses are derived from
// this seed.  If nil, a secure random seed is generated.
func (l *Loader) CreateNewWallet(pubPassphrase, privPassphrase, seed []byte) (w *wallet.Wallet, err error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrWalletLoaded
	}

	// Ensure that the network directory exists.
	if fi, err := os.Stat(l.dbDirPath); err != nil {
		if os.IsNotExist(err) {
			// Attempt data directory creation
			if err = os.MkdirAll(l.dbDirPath, 0700); err != nil {
				return nil, fmt.Errorf("cannot create directory: %s", err)
			}
		} else {
			return nil, fmt.Errorf("error checking directory: %s", err)
		}
	} else {
		if !fi.IsDir() {
			return nil, fmt.Errorf("path '%s' is not a directory", l.dbDirPath)
		}
	}

	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	exists, err := fileExists(dbPath)
	if err != nil {
		return nil, err
	}
	if exists {
		return nil, ErrWalletExists
	}

	// Create the wallet database backed by bolt db.
	err = os.MkdirAll(l.dbDirPath, 0700)
	if err != nil {
		return nil, err
	}
	db, err := walletdb.Create("bdb", dbPath)
	if err != nil {
		return nil, err
	}
	// Attempt to remove database file if this function errors.
	defer func() {
		if err != nil {
			_ = os.Remove(dbPath)
		}
	}()

	// Initialize the newly created database for the wallet before opening.
	err = wallet.Create(db, pubPassphrase, privPassphrase, seed, l.chainParams)
	if err != nil {
		return nil, err
	}

	// Open the newly-created wallet.
	so := l.stakeOptions
	w, err = wallet.Open(db, pubPassphrase, so.VotingEnabled, so.AddressReuse,
		so.TicketAddress, so.PoolAddress, so.PoolFees, so.TicketFee,
		l.addrIdxScanLen, so.StakePoolColdExtKey, l.allowHighFees,
		l.relayFee, l.chainParams)
	if err != nil {
		return nil, err
	}
	w.Start()

	l.onLoaded(w, db)
	return w, nil
}

// OpenExistingWallet opens the wallet from the loader's wallet database path
// and the public passphrase.  If the loader is being called by a context where
// standard input prompts may be used during wallet upgrades, setting
// canConsolePrompt will enable these prompts.
func (l *Loader) OpenExistingWallet(pubPassphrase []byte) (w *wallet.Wallet, rerr error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrWalletLoaded
	}

	// Open the database using the boltdb backend.
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	db, err := walletdb.Open("bdb", dbPath)
	if err != nil {
		log.Errorf("Failed to open database: %v", err)
		return nil, err
	}
	// If this function does not return to completion the database must be
	// closed.  Otherwise, because the database is locked on opens, any
	// other attempts to open the wallet will hang, and there is no way to
	// recover since this db handle would be leaked.
	defer func() {
		if rerr != nil {
			db.Close()
		}
	}()

	so := l.stakeOptions
	w, err = wallet.Open(db, pubPassphrase, so.VotingEnabled, so.AddressReuse,
		so.TicketAddress, so.PoolAddress, so.PoolFees, so.TicketFee,
		l.addrIdxScanLen, so.StakePoolColdExtKey, l.allowHighFees,
		l.relayFee, l.chainParams)
	if err != nil {
		return nil, err
	}

	w.Start()
	l.onLoaded(w, db)
	return w, nil
}

// WalletExists returns whether a file exists at the loader's database path.
// This may return an error for unexpected I/O failures.
func (l *Loader) WalletExists() (bool, error) {
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	return fileExists(dbPath)
}

// LoadedWallet returns the loaded wallet, if any, and a bool for whether the
// wallet has been loaded or not.  If true, the wallet pointer should be safe to
// dereference.
func (l *Loader) LoadedWallet() (*wallet.Wallet, bool) {
	l.mu.Lock()
	w := l.wallet
	l.mu.Unlock()
	return w, w != nil
}

// UnloadWallet stops the loaded wallet, if any, and closes the wallet
// database.  This returns ErrWalletNotLoaded if the wallet has not been loaded
// with CreateNewWallet or LoadExistingWallet.  The Loader may be reused if
// this function returns without error.
func (l *Loader) UnloadWallet() error {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet == nil {
		return ErrWalletNotLoaded
	}

	// Ignore err already stopped.
	l.stopTicketPurchase()

	l.wallet.Stop()
	l.wallet.WaitForShutdown()
	err := l.db.Close()
	if err != nil {
		return err
	}

	l.wallet = nil
	l.db = nil
	return nil
}

// SetChainClient sets the chain server client.
func (l *Loader) SetChainClient(chainClient *dcrrpcclient.Client) {
	l.mu.Lock()
	l.chainClient = chainClient
	l.mu.Unlock()
}

// StartTicketPurchase launches the ticketbuyer to start purchasing tickets.
func (l *Loader) StartTicketPurchase(passphrase []byte, ticketbuyerCfg *ticketbuyer.Config) error {
	defer l.mu.Unlock()
	l.mu.Lock()

	// Already running?
	if l.purchaseManager != nil {
		return ErrTicketBuyerStarted
	}

	if l.wallet == nil {
		return ErrWalletNotLoaded
	}

	if l.chainClient == nil {
		return ErrNoChainClient
	}

	w := l.wallet
	p, err := ticketbuyer.NewTicketPurchaser(ticketbuyerCfg, l.chainClient, w, l.chainParams)
	if err != nil {
		return err
	}
	n := w.NtfnServer.MainTipChangedNotifications()
	pm := ticketbuyer.NewPurchaseManager(w, p, n.C, passphrase)
	l.ntfnClient = n
	l.purchaseManager = pm
	pm.Start()
	l.wallet.SetTicketPurchasingEnabled(true)
	return nil
}

// stopTicketPurchase stops the ticket purchaser, waiting until it has finished.
// It must be called with the mutex lock held.
func (l *Loader) stopTicketPurchase() error {
	if l.purchaseManager == nil {
		return ErrTicketBuyerStopped
	}

	l.ntfnClient.Done()
	l.purchaseManager.Stop()
	l.purchaseManager.WaitForShutdown()
	l.purchaseManager = nil
	l.wallet.SetTicketPurchasingEnabled(false)
	return nil
}

// StopTicketPurchase stops the ticket purchaser, waiting until it has finished.
// If no ticket purchaser is running, it returns ErrTicketBuyerStopped.
func (l *Loader) StopTicketPurchase() error {
	defer l.mu.Unlock()
	l.mu.Lock()

	return l.stopTicketPurchase()
}

// PurchaseManager returns the ticket purchaser instance. If ticket purchasing
// has been disabled, it returns nil.
func (l *Loader) PurchaseManager() *ticketbuyer.PurchaseManager {
	l.mu.Lock()
	pm := l.purchaseManager
	l.mu.Unlock()
	return pm
}

func fileExists(filePath string) (bool, error) {
	_, err := os.Stat(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<commit_msg>Defer --create error cleanup sooner.<commit_after>// Copyright (c) 2015-2017 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package loader

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrrpcclient"
	"github.com/decred/dcrwallet/ticketbuyer"
	"github.com/decred/dcrwallet/wallet"
	"github.com/decred/dcrwallet/walletdb"
	_ "github.com/decred/dcrwallet/walletdb/bdb" // driver loaded during init
)

const (
	walletDbName = "wallet.db"
)

// Loader implements the creating of new and opening of existing wallets, while
// providing a callback system for other subsystems to handle the loading of a
// wallet.  This is primarely intended for use by the RPC servers, to enable
// methods and services which require the wallet when the wallet is loaded by
// another subsystem.
//
// Loader is safe for concurrent access.
type Loader struct {
	callbacks   []func(*wallet.Wallet)
	chainClient *dcrrpcclient.Client
	chainParams *chaincfg.Params
	dbDirPath   string
	wallet      *wallet.Wallet
	db          walletdb.DB
	mu          sync.Mutex

	purchaseManager *ticketbuyer.PurchaseManager
	ntfnClient      wallet.MainTipChangedNotificationsClient
	stakeOptions    *StakeOptions
	addrIdxScanLen  int
	allowHighFees   bool
	relayFee        float64
}

// StakeOptions contains the various options necessary for stake mining.
type StakeOptions struct {
	TicketPurchasingEnabled bool
	VotingEnabled           bool
	TicketFee               float64
	AddressReuse            bool
	TicketAddress           string
	PoolAddress             string
	PoolFees                float64
	StakePoolColdExtKey     string
}

// NewLoader constructs a Loader.
func NewLoader(chainParams *chaincfg.Params, dbDirPath string, stakeOptions *StakeOptions, addrIdxScanLen int,
	allowHighFees bool, relayFee float64) *Loader {

	return &Loader{
		chainParams:    chainParams,
		dbDirPath:      dbDirPath,
		stakeOptions:   stakeOptions,
		addrIdxScanLen: addrIdxScanLen,
		allowHighFees:  allowHighFees,
		relayFee:       relayFee,
	}
}

// onLoaded executes each added callback and prevents loader from loading any
// additional wallets.  Requires mutex to be locked.
func (l *Loader) onLoaded(w *wallet.Wallet, db walletdb.DB) {
	for _, fn := range l.callbacks {
		fn(w)
	}

	l.wallet = w
	l.db = db
	l.callbacks = nil // not needed anymore
}

// RunAfterLoad adds a function to be executed when the loader creates or opens
// a wallet.  Functions are executed in a single goroutine in the order they are
// added.
func (l *Loader) RunAfterLoad(fn func(*wallet.Wallet)) {
	l.mu.Lock()
	if l.wallet != nil {
		w := l.wallet
		l.mu.Unlock()
		fn(w)
	} else {
		l.callbacks = append(l.callbacks, fn)
		l.mu.Unlock()
	}
}

// CreateNewWallet creates a new wallet using the provided public and private
// passphrases.  The seed is optional.  If non-nil, addresses are derived from
// this seed.  If nil, a secure random seed is generated.
func (l *Loader) CreateNewWallet(pubPassphrase, privPassphrase, seed []byte) (w *wallet.Wallet, err error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrWalletLoaded
	}

	// Ensure that the network directory exists.
	if fi, err := os.Stat(l.dbDirPath); err != nil {
		if os.IsNotExist(err) {
			// Attempt data directory creation
			if err = os.MkdirAll(l.dbDirPath, 0700); err != nil {
				return nil, fmt.Errorf("cannot create directory: %s", err)
			}
		} else {
			return nil, fmt.Errorf("error checking directory: %s", err)
		}
	} else {
		if !fi.IsDir() {
			return nil, fmt.Errorf("path '%s' is not a directory", l.dbDirPath)
		}
	}

	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	exists, err := fileExists(dbPath)
	if err != nil {
		return nil, err
	}
	if exists {
		return nil, ErrWalletExists
	}

	// At this point it is asserted that there is no existing database file, and
	// deleting anything won't destroy a wallet in use.  Defer a function that
	// attempts to remove any written database file if this function errors.
	defer func() {
		if err != nil {
			_ = os.Remove(dbPath)
		}
	}()

	// Create the wallet database backed by bolt db.
	err = os.MkdirAll(l.dbDirPath, 0700)
	if err != nil {
		return nil, err
	}
	db, err := walletdb.Create("bdb", dbPath)
	if err != nil {
		return nil, err
	}

	// Initialize the newly created database for the wallet before opening.
	err = wallet.Create(db, pubPassphrase, privPassphrase, seed, l.chainParams)
	if err != nil {
		return nil, err
	}

	// Open the newly-created wallet.
	so := l.stakeOptions
	w, err = wallet.Open(db, pubPassphrase, so.VotingEnabled, so.AddressReuse,
		so.TicketAddress, so.PoolAddress, so.PoolFees, so.TicketFee,
		l.addrIdxScanLen, so.StakePoolColdExtKey, l.allowHighFees,
		l.relayFee, l.chainParams)
	if err != nil {
		return nil, err
	}
	w.Start()

	l.onLoaded(w, db)
	return w, nil
}

// OpenExistingWallet opens the wallet from the loader's wallet database path
// and the public passphrase.  If the loader is being called by a context where
// standard input prompts may be used during wallet upgrades, setting
// canConsolePrompt will enable these prompts.
func (l *Loader) OpenExistingWallet(pubPassphrase []byte) (w *wallet.Wallet, rerr error) {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet != nil {
		return nil, ErrWalletLoaded
	}

	// Open the database using the boltdb backend.
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	db, err := walletdb.Open("bdb", dbPath)
	if err != nil {
		log.Errorf("Failed to open database: %v", err)
		return nil, err
	}
	// If this function does not return to completion the database must be
	// closed.  Otherwise, because the database is locked on opens, any
	// other attempts to open the wallet will hang, and there is no way to
	// recover since this db handle would be leaked.
	defer func() {
		if rerr != nil {
			db.Close()
		}
	}()

	so := l.stakeOptions
	w, err = wallet.Open(db, pubPassphrase, so.VotingEnabled, so.AddressReuse,
		so.TicketAddress, so.PoolAddress, so.PoolFees, so.TicketFee,
		l.addrIdxScanLen, so.StakePoolColdExtKey, l.allowHighFees,
		l.relayFee, l.chainParams)
	if err != nil {
		return nil, err
	}

	w.Start()
	l.onLoaded(w, db)
	return w, nil
}

// WalletExists returns whether a file exists at the loader's database path.
// This may return an error for unexpected I/O failures.
func (l *Loader) WalletExists() (bool, error) {
	dbPath := filepath.Join(l.dbDirPath, walletDbName)
	return fileExists(dbPath)
}

// LoadedWallet returns the loaded wallet, if any, and a bool for whether the
// wallet has been loaded or not.  If true, the wallet pointer should be safe to
// dereference.
func (l *Loader) LoadedWallet() (*wallet.Wallet, bool) {
	l.mu.Lock()
	w := l.wallet
	l.mu.Unlock()
	return w, w != nil
}

// UnloadWallet stops the loaded wallet, if any, and closes the wallet
// database.  This returns ErrWalletNotLoaded if the wallet has not been loaded
// with CreateNewWallet or LoadExistingWallet.  The Loader may be reused if
// this function returns without error.
func (l *Loader) UnloadWallet() error {
	defer l.mu.Unlock()
	l.mu.Lock()

	if l.wallet == nil {
		return ErrWalletNotLoaded
	}

	// Ignore err already stopped.
	l.stopTicketPurchase()

	l.wallet.Stop()
	l.wallet.WaitForShutdown()
	err := l.db.Close()
	if err != nil {
		return err
	}

	l.wallet = nil
	l.db = nil
	return nil
}

// SetChainClient sets the chain server client.
func (l *Loader) SetChainClient(chainClient *dcrrpcclient.Client) {
	l.mu.Lock()
	l.chainClient = chainClient
	l.mu.Unlock()
}

// StartTicketPurchase launches the ticketbuyer to start purchasing tickets.
func (l *Loader) StartTicketPurchase(passphrase []byte, ticketbuyerCfg *ticketbuyer.Config) error {
	defer l.mu.Unlock()
	l.mu.Lock()

	// Already running?
	if l.purchaseManager != nil {
		return ErrTicketBuyerStarted
	}

	if l.wallet == nil {
		return ErrWalletNotLoaded
	}

	if l.chainClient == nil {
		return ErrNoChainClient
	}

	w := l.wallet
	p, err := ticketbuyer.NewTicketPurchaser(ticketbuyerCfg, l.chainClient, w, l.chainParams)
	if err != nil {
		return err
	}
	n := w.NtfnServer.MainTipChangedNotifications()
	pm := ticketbuyer.NewPurchaseManager(w, p, n.C, passphrase)
	l.ntfnClient = n
	l.purchaseManager = pm
	pm.Start()
	l.wallet.SetTicketPurchasingEnabled(true)
	return nil
}

// stopTicketPurchase stops the ticket purchaser, waiting until it has finished.
// It must be called with the mutex lock held.
func (l *Loader) stopTicketPurchase() error {
	if l.purchaseManager == nil {
		return ErrTicketBuyerStopped
	}

	l.ntfnClient.Done()
	l.purchaseManager.Stop()
	l.purchaseManager.WaitForShutdown()
	l.purchaseManager = nil
	l.wallet.SetTicketPurchasingEnabled(false)
	return nil
}

// StopTicketPurchase stops the ticket purchaser, waiting until it has finished.
// If no ticket purchaser is running, it returns ErrTicketBuyerStopped.
func (l *Loader) StopTicketPurchase() error {
	defer l.mu.Unlock()
	l.mu.Lock()

	return l.stopTicketPurchase()
}

// PurchaseManager returns the ticket purchaser instance. If ticket purchasing
// has been disabled, it returns nil.
func (l *Loader) PurchaseManager() *ticketbuyer.PurchaseManager {
	l.mu.Lock()
	pm := l.purchaseManager
	l.mu.Unlock()
	return pm
}

func fileExists(filePath string) (bool, error) {
	_, err := os.Stat(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<|endoftext|>


<commit_before>package filter

import (
	"errors"
	"github.com/gogo/protobuf/proto"
)

var (
	ErrNotFound       = errors.New("Not found")
	ErrTooManyResults = errors.New("Too many results")
)

// Filter is used to query a State and match
// a single protobuf message.
type Filter func(proto.Message) bool

// Messages is a filterable array of protobuf.Message.
type Messages []proto.Message

/*
type State struct {
	messages []proto.Message
}

// StateFromAgent loads a State with proto.messages
// from a GET_STATE call.
// TODO: Consider using reflect to reduce typing.
*/
// FindAny will return the first message
// where all filters return true. If no
// messages match we will return ErrNotFound.
func (messages Messages) FindAny(filters ...Filter) (proto.Message, error) {
	var match bool
loop:
	for _, message := range messages {
		for _, filter := range filters {
			match = filter(message)
			if !match {
				continue loop
			}
		}
		if match {
			return message, nil
		}
	}
	return nil, ErrNotFound
}

// FindOne will a single message where all
// filters return true. If more than one
// message matches return ErrTooManyResults.
// If no messages match all filters we will
// return ErrNotFound.
func (messages Messages) FindOne(filters ...Filter) (proto.Message, error) {
	var (
		position int
		match    bool
	)
loop:
	for i, message := range messages {
		for _, filter := range filters {
			if !filter(message) {
				continue loop
			}
		}
		// Already matched a message
		if match {
			return nil, ErrTooManyResults
		}
		// Mark this message as matched
		match = true
		position = i
	}
	if !match {
		return nil, ErrNotFound
	}
	return messages[position], nil
}

// FindMany will return as many messages
// in which all of the filters match true.
// If no messages match all the filters
// we will return an empty []proto.Message.
func (messages Messages) FindMany(filters ...Filter) Messages {
	matches := Messages{}
loop:
	for _, message := range messages {
		for _, filter := range filters {
			if !filter(message) {
				continue loop
			}
		}
		matches = append(matches, message)
	}
	return matches
}
<commit_msg>remove comment<commit_after>package filter

import (
	"errors"
	"github.com/gogo/protobuf/proto"
)

var (
	ErrNotFound       = errors.New("Not found")
	ErrTooManyResults = errors.New("Too many results")
)

// Filter is used to query a State and match
// a single protobuf message.
type Filter func(proto.Message) bool

// Messages is a filterable array of protobuf.Message.
type Messages []proto.Message

// FindAny will return the first message
// where all filters return true. If no
// messages match we will return ErrNotFound.
func (messages Messages) FindAny(filters ...Filter) (proto.Message, error) {
	var match bool
loop:
	for _, message := range messages {
		for _, filter := range filters {
			match = filter(message)
			if !match {
				continue loop
			}
		}
		if match {
			return message, nil
		}
	}
	return nil, ErrNotFound
}

// FindOne will a single message where all
// filters return true. If more than one
// message matches return ErrTooManyResults.
// If no messages match all filters we will
// return ErrNotFound.
func (messages Messages) FindOne(filters ...Filter) (proto.Message, error) {
	var (
		position int
		match    bool
	)
loop:
	for i, message := range messages {
		for _, filter := range filters {
			if !filter(message) {
				continue loop
			}
		}
		// Already matched a message
		if match {
			return nil, ErrTooManyResults
		}
		// Mark this message as matched
		match = true
		position = i
	}
	if !match {
		return nil, ErrNotFound
	}
	return messages[position], nil
}

// FindMany will return as many messages
// in which all of the filters match true.
// If no messages match all the filters
// we will return an empty []proto.Message.
func (messages Messages) FindMany(filters ...Filter) Messages {
	matches := Messages{}
loop:
	for _, message := range messages {
		for _, filter := range filters {
			if !filter(message) {
				continue loop
			}
		}
		matches = append(matches, message)
	}
	return matches
}
<|endoftext|>


<commit_before>package loader

import (
	"encoding/csv"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/brnstz/bus/internal/conf"
	"github.com/brnstz/bus/models"
	"github.com/jmoiron/sqlx"
)

var days = []string{"monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"}

var datefmt = "20060102"

type Loader struct {
	// the dir from which we load google transit files
	dir string

	// mapping from trip id to a trip object
	trips map[string]*models.Trip

	// mapping from stop_id to a slice of trip_ids
	stopTrips map[string][]string

	// mapping trip_id to route_id
	tripRoute map[string]string

	// a map of "{stop_id}-{route_id}" to stop objects. Essentially a
	// list of unique stops by route.
	uniqueStop map[string]*models.Stop

	// mapping of trip_id to service object
	tripService map[string]*models.Service

	// mapping of service_id to map of unique route_id
	serviceRoute map[string]map[string]bool

	routes map[string]bool

	Stops []*models.Stop

	ScheduledStopTimes []*models.ScheduledStopTime

	ServiceRouteDays       []*models.ServiceRouteDay
	ServiceRouteExceptions []*models.ServiceRouteException
}

func NewLoader(dir string, routeFilter string) *Loader {
	l := Loader{
		dir:          dir,
		trips:        map[string]*models.Trip{},
		stopTrips:    map[string][]string{},
		tripRoute:    map[string]string{},
		uniqueStop:   map[string]*models.Stop{},
		tripService:  map[string]*models.Service{},
		serviceRoute: map[string]map[string]bool{},
	}

	if len(routeFilter) > 0 {
		l.routes = map[string]bool{}
		allRoutes := strings.Split(routeFilter, "|")
		for _, v := range allRoutes {
			l.routes[v] = true
		}
	}

	l.init()

	return &l
}

func (l *Loader) init() {
	l.loadTrips()
	l.loadStopTrips()
	l.loadTripRoute()
	l.loadUniqueStop()
	l.loadCalendars()

	l.Stops = make([]*models.Stop, len(l.uniqueStop))

	i := 0
	for _, v := range l.uniqueStop {
		l.Stops[i] = v
		i++
	}
}

func getcsv(dir, name string) *csv.Reader {

	f, err := os.Open(path.Join(dir, name))
	if err != nil {
		panic(err)
	}

	r := csv.NewReader(f)
	r.LazyQuotes = true
	return r
}

// find index of col in header

func find(header []string, col string) int {
	for i := 0; i < len(header); i++ {
		if header[i] == col {
			return i
		}
	}

	log.Fatalf("can't find header col %v", col)
	return -1
}

// skipRoute returns true if we should skip this route given our routeFilter
// config
func (l *Loader) skipRoute(routeID string) bool {
	if l.routes != nil && l.routes[routeID] == false {
		return true
	} else {
		return false
	}
}

func (l *Loader) loadTrips() {
	f := getcsv(l.dir, "trips.txt")

	header, err := f.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	tripIdx := find(header, "trip_id")
	dirIdx := find(header, "direction_id")
	headIdx := find(header, "trip_headsign")
	serviceIdx := find(header, "service_id")
	routeIdx := find(header, "route_id")

	for i := 0; ; i++ {
		rec, err := f.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of trips.txt", err, i)
		}
		direction, err := strconv.Atoi(rec[dirIdx])
		if err != nil {
			log.Fatalf("%v on line %v of trips.txt", err, i)
		}

		trip := &models.Trip{
			Id:          rec[tripIdx],
			DirectionId: direction,
			Headsign:    rec[headIdx],
		}

		service := rec[serviceIdx]
		route := rec[routeIdx]

		if l.skipRoute(route) {
			continue
		}

		l.trips[trip.Id] = trip

		serviceObj := &models.Service{
			Id:      service,
			RouteId: route,
		}

		l.tripService[trip.Id] = serviceObj

		if l.serviceRoute[service] == nil {
			l.serviceRoute[service] = map[string]bool{}
		}
		l.serviceRoute[service][route] = true
	}

}

func (l *Loader) loadStopTrips() {
	stopTimes := getcsv(l.dir, "stop_times.txt")

	header, err := stopTimes.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	stopIdx := find(header, "stop_id")
	tripIdx := find(header, "trip_id")
	timeIdx := find(header, "departure_time")
	for i := 0; ; i++ {
		rec, err := stopTimes.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of stop_times.txt", err, i)
		}

		stop := rec[stopIdx]
		trip := rec[tripIdx]
		timeStr := rec[timeIdx]

		l.stopTrips[stop] = append(l.stopTrips[stop], trip)

		service, exists := l.tripService[trip]
		if !exists {
			continue
		}

		sst, err := models.NewScheduledStopTime(service.RouteId, stop, service.Id, timeStr)
		if err != nil {
			log.Fatal("can't create sst", rec, err)
		}

		l.ScheduledStopTimes = append(l.ScheduledStopTimes, &sst)

	}
}

func (l *Loader) loadTripRoute() {

	trips := getcsv(l.dir, "trips.txt")

	header, err := trips.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	tripIdx := find(header, "trip_id")
	routeIdx := find(header, "route_id")

	trips.Read()
	for i := 0; ; i++ {
		rec, err := trips.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of trips.txt", err, i)
		}

		trip := rec[tripIdx]
		route := rec[routeIdx]

		if l.skipRoute(route) {
			continue
		}

		l.tripRoute[trip] = route
	}
}

func (l *Loader) loadUniqueStop() {
	stops := getcsv(l.dir, "stops.txt")

	header, err := stops.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	stopIdx := find(header, "stop_id")
	stopNameIdx := find(header, "stop_name")
	stopLatIdx := find(header, "stop_lat")
	stopLonIdx := find(header, "stop_lon")

	for i := 0; ; i++ {
		rec, err := stops.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of stops.txt", err, i)
		}

		stopLat, err := strconv.ParseFloat(
			strings.TrimSpace(rec[stopLatIdx]), 64,
		)
		if err != nil {
			log.Fatalf("%v on line %v of stops.txt", err, i)
		}

		stopLon, err := strconv.ParseFloat(
			strings.TrimSpace(rec[stopLonIdx]), 64,
		)
		if err != nil {
			log.Fatalf("%v on line %v of stops.txt", err, i)
		}

		trips, exists := l.stopTrips[rec[stopIdx]]
		if exists {
			for _, trip := range trips {
				if l.skipRoute(l.tripRoute[trip]) {
					continue
				}
				obj := models.Stop{
					Id:      rec[stopIdx],
					Name:    rec[stopNameIdx],
					Lat:     stopLat,
					Lon:     stopLon,
					RouteId: l.tripRoute[trip],

					DirectionId: l.trips[trip].DirectionId,
					Headsign:    l.trips[trip].Headsign,
				}
				l.uniqueStop[obj.Key()] = &obj
			}
		}
	}
}

func (l *Loader) loadCalendars() {
	stops := getcsv(l.dir, "calendar.txt")

	header, err := stops.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	idxs := map[string]int{}
	for _, day := range days {
		idxs[day] = find(header, day)
	}
	serviceIdx := find(header, "service_id")
	startDateIdx := find(header, "start_date")
	endDateIdx := find(header, "end_date")

	for i := 0; ; i++ {
		rec, err := stops.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of calendar.txt", err, i)
		}

		serviceId := rec[serviceIdx]

		startDate, err := time.Parse(datefmt, rec[startDateIdx])
		if err != nil {
			log.Fatalf("can't parse start date %v %v", err, rec[startDateIdx])
		}

		endDate, err := time.Parse(datefmt, rec[endDateIdx])
		if err != nil {
			log.Fatalf("can't parse end date %v %v", err, rec[endDateIdx])
		}

		for day, dayIdx := range idxs {
			dayVal := rec[dayIdx]
			if dayVal != "1" {
				continue
			}
			for route, _ := range l.serviceRoute[serviceId] {
				srd := models.ServiceRouteDay{
					ServiceId: serviceId,
					RouteId:   route,
					Day:       day,
					StartDate: startDate,
					EndDate:   endDate,
				}

				l.ServiceRouteDays = append(l.ServiceRouteDays, &srd)
			}
		}
	}
}

func doOne(dir string, stype string, routeFilter string, db *sqlx.DB) {

	l := NewLoader(dir, routeFilter)

	for i, s := range l.ServiceRouteDays {
		_, err := db.Exec(`
				INSERT INTO service_route_day
				(route_id, service_id, day, start_date, end_date)
				VALUES($1, $2, $3, $4, $5)
			`, s.RouteId, s.ServiceId, s.Day, s.StartDate, s.EndDate,
		)
		if err != nil && !strings.Contains(err.Error(), "violates unique constraint") {
			log.Println("ERROR SERVICE ROUTE DAYS: ", err, s)
		}

		if i%100 == 0 {
			log.Printf("loaded %v service route days", i)
		}

	}

	for i, s := range l.Stops {
		_, err := db.Exec(`
				INSERT INTO stop
				(stop_id, stop_name, direction_id, headsign, route_id,
				 location, stype)
				VALUES($1, $2, $3, $4, $5, ll_to_earth($6, $7), $8)
			`,
			s.Id, s.Name, s.DirectionId, s.Headsign, s.RouteId,
			s.Lat, s.Lon, stype,
		)

		if err != nil && !strings.Contains(err.Error(), "violates unique constraint") {
			log.Println("ERROR STOPS: ", err, s)
		}

		if i%100 == 0 {
			log.Printf("loaded %v stops", i)
		}
	}

	for i, s := range l.ScheduledStopTimes {
		_, err := db.Exec(`
				INSERT INTO scheduled_stop_time
				(route_id, stop_id, service_id, departure_sec)
				VALUES($1, $2, $3, $4)
			`, s.RouteId, s.StopId, s.ServiceId, s.DepartureSec,
		)
		if err != nil && !strings.Contains(err.Error(), "violates unique constraint") {
			log.Println("ERROR SCHEDULED STOP TIMES: ", err, s)
		}

		if i%100000 == 0 {
			log.Printf("loaded %v stop times", i)
		}
	}
}

func LoadOnce(routeFilter string, urls ...string) {
	for _, url := range urls {

		var stype string
		if strings.Contains(url, "subway") {
			stype = "subway"
		} else {
			stype = "bus"
		}

		// FIXME: do this in Go, need to make it integrated with loader
		dir, err := ioutil.TempDir(conf.TmpDir, "")
		if err != nil {
			panic(err)
		}
		cmd := exec.Command("wget", url, "-O", path.Join(dir, "file.zip"))
		err = cmd.Run()
		if err != nil {
			panic(err)
		}

		cmd = exec.Command("unzip", path.Join(dir, "file.zip"), "-d", dir)
		err = cmd.Run()
		if err != nil {
			panic(err)
		}

		func() {
			log.Println(url, dir, stype)
			defer os.RemoveAll(dir)
			t1 := time.Now()
			doOne(dir, stype, routeFilter, conf.DB)
			t2 := time.Now()
			log.Printf("took %v for %v\n", t2.Sub(t1), dir)
		}()

	}
}

func LoadForever(routeFilter string, urls ...string) {
	for {
		LoadOnce(routeFilter, urls...)
		log.Println("finished loading, sleeping for 24 hours")
		time.Sleep(time.Hour * 24)
	}
}
<commit_msg>don't log loading 0<commit_after>package loader

import (
	"encoding/csv"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/brnstz/bus/internal/conf"
	"github.com/brnstz/bus/models"
	"github.com/jmoiron/sqlx"
)

var days = []string{"monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"}

var datefmt = "20060102"

type Loader struct {
	// the dir from which we load google transit files
	dir string

	// mapping from trip id to a trip object
	trips map[string]*models.Trip

	// mapping from stop_id to a slice of trip_ids
	stopTrips map[string][]string

	// mapping trip_id to route_id
	tripRoute map[string]string

	// a map of "{stop_id}-{route_id}" to stop objects. Essentially a
	// list of unique stops by route.
	uniqueStop map[string]*models.Stop

	// mapping of trip_id to service object
	tripService map[string]*models.Service

	// mapping of service_id to map of unique route_id
	serviceRoute map[string]map[string]bool

	routes map[string]bool

	Stops []*models.Stop

	ScheduledStopTimes []*models.ScheduledStopTime

	ServiceRouteDays       []*models.ServiceRouteDay
	ServiceRouteExceptions []*models.ServiceRouteException
}

func NewLoader(dir string, routeFilter string) *Loader {
	l := Loader{
		dir:          dir,
		trips:        map[string]*models.Trip{},
		stopTrips:    map[string][]string{},
		tripRoute:    map[string]string{},
		uniqueStop:   map[string]*models.Stop{},
		tripService:  map[string]*models.Service{},
		serviceRoute: map[string]map[string]bool{},
	}

	if len(routeFilter) > 0 {
		l.routes = map[string]bool{}
		allRoutes := strings.Split(routeFilter, "|")
		for _, v := range allRoutes {
			l.routes[v] = true
		}
	}

	l.init()

	return &l
}

func (l *Loader) init() {
	l.loadTrips()
	l.loadStopTrips()
	l.loadTripRoute()
	l.loadUniqueStop()
	l.loadCalendars()

	l.Stops = make([]*models.Stop, len(l.uniqueStop))

	i := 0
	for _, v := range l.uniqueStop {
		l.Stops[i] = v
		i++
	}
}

func getcsv(dir, name string) *csv.Reader {

	f, err := os.Open(path.Join(dir, name))
	if err != nil {
		panic(err)
	}

	r := csv.NewReader(f)
	r.LazyQuotes = true
	return r
}

// find index of col in header

func find(header []string, col string) int {
	for i := 0; i < len(header); i++ {
		if header[i] == col {
			return i
		}
	}

	log.Fatalf("can't find header col %v", col)
	return -1
}

// skipRoute returns true if we should skip this route given our routeFilter
// config
func (l *Loader) skipRoute(routeID string) bool {
	if l.routes != nil && l.routes[routeID] == false {
		return true
	} else {
		return false
	}
}

func (l *Loader) loadTrips() {
	f := getcsv(l.dir, "trips.txt")

	header, err := f.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	tripIdx := find(header, "trip_id")
	dirIdx := find(header, "direction_id")
	headIdx := find(header, "trip_headsign")
	serviceIdx := find(header, "service_id")
	routeIdx := find(header, "route_id")

	for i := 0; ; i++ {
		rec, err := f.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of trips.txt", err, i)
		}
		direction, err := strconv.Atoi(rec[dirIdx])
		if err != nil {
			log.Fatalf("%v on line %v of trips.txt", err, i)
		}

		trip := &models.Trip{
			Id:          rec[tripIdx],
			DirectionId: direction,
			Headsign:    rec[headIdx],
		}

		service := rec[serviceIdx]
		route := rec[routeIdx]

		if l.skipRoute(route) {
			continue
		}

		l.trips[trip.Id] = trip

		serviceObj := &models.Service{
			Id:      service,
			RouteId: route,
		}

		l.tripService[trip.Id] = serviceObj

		if l.serviceRoute[service] == nil {
			l.serviceRoute[service] = map[string]bool{}
		}
		l.serviceRoute[service][route] = true
	}

}

func (l *Loader) loadStopTrips() {
	stopTimes := getcsv(l.dir, "stop_times.txt")

	header, err := stopTimes.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	stopIdx := find(header, "stop_id")
	tripIdx := find(header, "trip_id")
	timeIdx := find(header, "departure_time")
	for i := 0; ; i++ {
		rec, err := stopTimes.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of stop_times.txt", err, i)
		}

		stop := rec[stopIdx]
		trip := rec[tripIdx]
		timeStr := rec[timeIdx]

		l.stopTrips[stop] = append(l.stopTrips[stop], trip)

		service, exists := l.tripService[trip]
		if !exists {
			continue
		}

		sst, err := models.NewScheduledStopTime(service.RouteId, stop, service.Id, timeStr)
		if err != nil {
			log.Fatal("can't create sst", rec, err)
		}

		l.ScheduledStopTimes = append(l.ScheduledStopTimes, &sst)

	}
}

func (l *Loader) loadTripRoute() {

	trips := getcsv(l.dir, "trips.txt")

	header, err := trips.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	tripIdx := find(header, "trip_id")
	routeIdx := find(header, "route_id")

	trips.Read()
	for i := 0; ; i++ {
		rec, err := trips.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of trips.txt", err, i)
		}

		trip := rec[tripIdx]
		route := rec[routeIdx]

		if l.skipRoute(route) {
			continue
		}

		l.tripRoute[trip] = route
	}
}

func (l *Loader) loadUniqueStop() {
	stops := getcsv(l.dir, "stops.txt")

	header, err := stops.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	stopIdx := find(header, "stop_id")
	stopNameIdx := find(header, "stop_name")
	stopLatIdx := find(header, "stop_lat")
	stopLonIdx := find(header, "stop_lon")

	for i := 0; ; i++ {
		rec, err := stops.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of stops.txt", err, i)
		}

		stopLat, err := strconv.ParseFloat(
			strings.TrimSpace(rec[stopLatIdx]), 64,
		)
		if err != nil {
			log.Fatalf("%v on line %v of stops.txt", err, i)
		}

		stopLon, err := strconv.ParseFloat(
			strings.TrimSpace(rec[stopLonIdx]), 64,
		)
		if err != nil {
			log.Fatalf("%v on line %v of stops.txt", err, i)
		}

		trips, exists := l.stopTrips[rec[stopIdx]]
		if exists {
			for _, trip := range trips {
				if l.skipRoute(l.tripRoute[trip]) {
					continue
				}
				obj := models.Stop{
					Id:      rec[stopIdx],
					Name:    rec[stopNameIdx],
					Lat:     stopLat,
					Lon:     stopLon,
					RouteId: l.tripRoute[trip],

					DirectionId: l.trips[trip].DirectionId,
					Headsign:    l.trips[trip].Headsign,
				}
				l.uniqueStop[obj.Key()] = &obj
			}
		}
	}
}

func (l *Loader) loadCalendars() {
	stops := getcsv(l.dir, "calendar.txt")

	header, err := stops.Read()
	if err != nil {
		log.Fatalf("unable to read header: %v", err)
	}

	idxs := map[string]int{}
	for _, day := range days {
		idxs[day] = find(header, day)
	}
	serviceIdx := find(header, "service_id")
	startDateIdx := find(header, "start_date")
	endDateIdx := find(header, "end_date")

	for i := 0; ; i++ {
		rec, err := stops.Read()
		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("%v on line %v of calendar.txt", err, i)
		}

		serviceId := rec[serviceIdx]

		startDate, err := time.Parse(datefmt, rec[startDateIdx])
		if err != nil {
			log.Fatalf("can't parse start date %v %v", err, rec[startDateIdx])
		}

		endDate, err := time.Parse(datefmt, rec[endDateIdx])
		if err != nil {
			log.Fatalf("can't parse end date %v %v", err, rec[endDateIdx])
		}

		for day, dayIdx := range idxs {
			dayVal := rec[dayIdx]
			if dayVal != "1" {
				continue
			}
			for route, _ := range l.serviceRoute[serviceId] {
				srd := models.ServiceRouteDay{
					ServiceId: serviceId,
					RouteId:   route,
					Day:       day,
					StartDate: startDate,
					EndDate:   endDate,
				}

				l.ServiceRouteDays = append(l.ServiceRouteDays, &srd)
			}
		}
	}
}

func doOne(dir string, stype string, routeFilter string, db *sqlx.DB) {

	l := NewLoader(dir, routeFilter)

	for i, s := range l.ServiceRouteDays {
		_, err := db.Exec(`
				INSERT INTO service_route_day
				(route_id, service_id, day, start_date, end_date)
				VALUES($1, $2, $3, $4, $5)
			`, s.RouteId, s.ServiceId, s.Day, s.StartDate, s.EndDate,
		)
		if err != nil && !strings.Contains(err.Error(), "violates unique constraint") {
			log.Println("ERROR SERVICE ROUTE DAYS: ", err, s)
		}

		if i%100 == 0 && i > 0 {
			log.Printf("loaded %v service route days", i)
		}

	}

	for i, s := range l.Stops {
		_, err := db.Exec(`
				INSERT INTO stop
				(stop_id, stop_name, direction_id, headsign, route_id,
				 location, stype)
				VALUES($1, $2, $3, $4, $5, ll_to_earth($6, $7), $8)
			`,
			s.Id, s.Name, s.DirectionId, s.Headsign, s.RouteId,
			s.Lat, s.Lon, stype,
		)

		if err != nil && !strings.Contains(err.Error(), "violates unique constraint") {
			log.Println("ERROR STOPS: ", err, s)
		}

		if i%100 == 0 && i > 0 {
			log.Printf("loaded %v stops", i)
		}
	}

	for i, s := range l.ScheduledStopTimes {
		_, err := db.Exec(`
				INSERT INTO scheduled_stop_time
				(route_id, stop_id, service_id, departure_sec)
				VALUES($1, $2, $3, $4)
			`, s.RouteId, s.StopId, s.ServiceId, s.DepartureSec,
		)
		if err != nil && !strings.Contains(err.Error(), "violates unique constraint") {
			log.Println("ERROR SCHEDULED STOP TIMES: ", err, s)
		}

		if i%100000 == 0 && i > 0 {
			log.Printf("loaded %v stop times", i)
		}
	}
}

func LoadOnce(routeFilter string, urls ...string) {
	for _, url := range urls {

		var stype string
		if strings.Contains(url, "subway") {
			stype = "subway"
		} else {
			stype = "bus"
		}

		// FIXME: do this in Go, need to make it integrated with loader
		dir, err := ioutil.TempDir(conf.TmpDir, "")
		if err != nil {
			panic(err)
		}
		cmd := exec.Command("wget", url, "-O", path.Join(dir, "file.zip"))
		err = cmd.Run()
		if err != nil {
			panic(err)
		}

		cmd = exec.Command("unzip", path.Join(dir, "file.zip"), "-d", dir)
		err = cmd.Run()
		if err != nil {
			panic(err)
		}

		func() {
			log.Println(url, dir, stype)
			defer os.RemoveAll(dir)
			t1 := time.Now()
			doOne(dir, stype, routeFilter, conf.DB)
			t2 := time.Now()
			log.Printf("took %v for %v\n", t2.Sub(t1), dir)
		}()

	}
}

func LoadForever(routeFilter string, urls ...string) {
	for {
		LoadOnce(routeFilter, urls...)
		log.Println("finished loading, sleeping for 24 hours")
		time.Sleep(time.Hour * 24)
	}
}
<|endoftext|>