spinel_extract

extract.go

package main

import (
	"log"

	"flag"
	"fmt"
	"strconv"
	"sync"

	"archive/zip"
	"os"

	"bytes"
	"image/png"

	"encoding/json"
	"sync/atomic"
	"time"

	"encoding/hex"

	"github.com/boltdb/bolt"
	"github.com/cr0sh/spinel"
	"labix.org/v2/mgo/bson"
)

var numberOfAvatars int64
var avatarsConverted int64 = 0

const singleImageSize = 96 * 96 * 3

func main() {
	var limit = flag.Int64("max", 0, "Maximum entry count for extract, 0=unlimited")
	flag.Parse()
	db, err := bolt.Open("database.db", 0600, &bolt.Options{ReadOnly: true})
	if err != nil {
		log.Fatal("Error while opening database: %s", err)
	}

	if err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("numberOfAvatars"))
		if b == nil {
			return fmt.Errorf("No such bucket: numberOfAvatars")
		}
		bn := b.Get([]byte("value"))
		if bn == nil {
			return fmt.Errorf("numberOfAvatars.value is empty")
		}
		if numberOfAvatars, err = strconv.ParseInt(string(bn), 10, 0); err != nil {
			return fmt.Errorf("Cannot parse numberOfAvatars.value: %v", bn)
		}
		return nil
	}); err != nil {
		log.Fatal(err)
	}

	log.Println("numberOfAvatars.value:", numberOfAvatars)
	if *limit == 0 {
		*limit = numberOfAvatars
		log.Printf("Extract limit: %d entries (unlimited)", *limit)
	} else if *limit > numberOfAvatars {
		*limit = numberOfAvatars
		log.Printf("Extract limit: %d entries", *limit)
		log.Printf("Warning: limit is greater than total count of database")
	} else {
		log.Printf("Extract limit: %d entries", *limit)
	}

	wg := new(sync.WaitGroup)

	zipfile, err := os.Create("data.zip")
	if err != nil {
		log.Fatal("Error while creating data.zip:", err)
	}
	rawfile, err := os.Create("data.dat")
	if err != nil {
		log.Fatal("Error while creating data.dat:", err)
	}
	zipwriter := zip.NewWriter(zipfile)
	defer func() {
		log.Println("Closing writers and files...")
		if err := zipwriter.Close(); err != nil {
			log.Fatal("Error while closing zipwriter: %s", err)
		}
		if err := zipfile.Close(); err != nil {
			log.Fatal("Error while closing data.zip: %s", err)
		}
		if err := rawfile.Close(); err != nil {
			log.Fatal("Error while closing data.dat: %s", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		if err = db.View(func(tx *bolt.Tx) error {
			imgb := tx.Bucket([]byte("AvatarImages"))
			if imgb == nil {
				return fmt.Errorf("No such bucket: AvatarImages")
			}
			metab := tx.Bucket([]byte("PlayerMetadata"))
			if metab == nil {
				return fmt.Errorf("No such bucket: PlayerMetadata")
			}

			for rank := int64(1); rank <= *limit; rank++ {
				fmtrank := []byte(strconv.FormatInt(rank, 10))
				pngbuf := imgb.Get(fmtrank)
				if pngbuf == nil {
					log.Printf("Avatar image does not exist on rank %d", rank)
					continue
					return fmt.Errorf("Avatar image does not exist on rank %d", rank)
				}
				metabs := metab.Get(fmtrank)
				if pngbuf == nil {
					return fmt.Errorf("Player metadata does not exist on rank %d", rank)
				}

				var le spinel.ListElement
				if err := bson.Unmarshal(metabs, &le); err != nil {
					return fmt.Errorf("BSON unmarshaling failed: %s", err)
				}
				jsonbs, err := json.MarshalIndent(le, "", "  ")
				if err != nil {
					return fmt.Errorf("Error while marshaling metadata to JSON: %s", err)
				}

				w, err := zipwriter.Create(fmt.Sprintf("img/%d.png", rank))
				if err != nil {
					return fmt.Errorf("Error while creating archived image file: %s, err")
				}
				if _, err := w.Write(pngbuf); err != nil {
					return fmt.Errorf("Error while writing PNG data: %s", err)
				}
				w, err = zipwriter.Create(fmt.Sprintf("meta/%d.json", rank))
				if err != nil {
					return fmt.Errorf("Error while creating archived metadata file: %s", err)
				}
				if _, err := w.Write(jsonbs); err != nil {
					return fmt.Errorf("Error while writing JSON metadata: %s", err)
				}

				p, err := png.Decode(bytes.NewReader(pngbuf))
				if err != nil {
					log.Println("Rank", rank)
					fmt.Println(hex.Dump(pngbuf))
					return fmt.Errorf("Error while decoding PNG data: %s", err)
				}
				rawbuf := make([]byte, singleImageSize)

				x1, y1, x2, y2 := p.Bounds().Min.X, p.Bounds().Min.Y, p.Bounds().Max.X, p.Bounds().Max.Y
				for y := y1; y < y2; y++ {
					for x := x1; x < x2; x++ {
						r, g, b, _ := p.At(x, y).RGBA()
						idx := ((x2-x1)*y + x) * 3
						rawbuf[idx] = byte(r >> 8)
						rawbuf[idx+1] = byte(g >> 8)
						rawbuf[idx+2] = byte(b >> 8)
					}
				}

				rawfile.WriteAt(rawbuf, singleImageSize*rank)
				atomic.AddInt64(&avatarsConverted, 1)
			}
			return nil
		}); err != nil {
			log.Fatal(err)
		}
	}()
	itr := make(chan struct{})
	go func() {
		t := time.Tick(time.Minute * 5)
		for {
			select {
			case <-itr:
				return
			case <-t:
				n := atomic.LoadInt64(&avatarsConverted)
				log.Printf("Number of avatars converted: %d (%05.2f%%)", n, float64(n)/float64(*limit)*100)
			}
		}
	}()
	wg.Wait()
	log.Println("All workers are done.")
}