rainbow table for phone number

rainbow.go

package main

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"crypto/rand"
	"encoding/binary"
	"encoding/hex"
	"flag"
	"fmt"
	bolt "go.etcd.io/bbolt"
	"math"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

type MobileRainbow struct {
	RBTable map[uint32][]uint32
	TableSize int
	ChainLength int
	Prefix string
}

const KeySpace = 1000000000
const ChainLength = 100000000
const TableSize = 64

func (rbt *MobileRainbow) ErrorRate() float64 {
	//1 - math.e ** (-100000000 * 100 / N)
	a := -float64(rbt.TableSize) * float64(rbt.ChainLength) / float64(KeySpace)
	return math.Exp(float64(a))
}

func (rbt *MobileRainbow) Lookup(hash []byte) string {
	for i := rbt.ChainLength - 1; i >= 0; i-- {
		end := rbt.gen_lookup_chain(hash, i)

		starting, found := rbt.RBTable[end]
		if found {
			for idx := range starting {
				start := starting[idx]
				number := rbt.find_in_chain(hash, start)
				if number != "" {
					return number
				}
			}
		}
	}
	return ""
}

func (rbt *MobileRainbow) LoadDB(file_path string) {
	f, err := os.Open(file_path)
	check(err)
	defer f.Close()
	rbt.RBTable = make(map[uint32][]uint32)
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		sp := strings.Split(line, ",")
		end_num := sp[0]
		begin_nums := strings.Split(sp[1][1:len(sp[1])-1], " ")

		end_int, _ := strconv.ParseUint(end_num, 10, 32)
		begin_ints := make([]uint32, len(begin_nums))
		for i := range begin_nums {
			bn, _ := strconv.ParseUint(begin_nums[i], 10, 32)
			begin_ints[i] = uint32(bn)
		}
		rbt.RBTable[uint32(end_int)] = begin_ints
	}
}

func (rbt *MobileRainbow) BuildDB(file_path string) {
	table := make(map[uint32]uint32)
	channel := make(chan uint32, 100)
	go key_generator(channel, uint32(KeySpace))
	var key uint32 = 0
	var counter int = 0
	for {
		key = <- channel
		_, exist := table[key]
		if !exist {
			end := rbt.generate_chain(key)
			table[key] = end
			counter++

			if counter % 100000 == 0 {
				fmt.Println(counter)
			}
		}
		if counter >= rbt.TableSize {
			break
		}
	}
	fmt.Println("Chain table size: ", len(table))

	// reverse map
	rbt.RBTable = make(map[uint32][]uint32)
	for k, v := range table {
		begin, exist := rbt.RBTable[v]
		if exist {
			rbt.RBTable[v] = append(begin, k)
		} else {
			rbt.RBTable[v] = []uint32{k}
		}
	}

	f, err := os.Create(file_path)
	check(err)
	defer f.Close()
	for k, v := range rbt.RBTable {
		_, err = fmt.Fprintf(f, "%d,%v\n", k, v)
		check(err)
	}
}

func (rbt *MobileRainbow) gen_lookup_chain(hash []byte, index int) uint32 {
	end := reduce(hash, index, uint32(KeySpace))
	for i := index + 1; i < rbt.ChainLength; i++ {
		h := md5_hash(end, rbt.Prefix)
		end = reduce(h, i, uint32(KeySpace))
	}
	return end
}

func (rbt *MobileRainbow) find_in_chain(needle []byte, start uint32) string {
	curr := start
	for i := 0; i < rbt.ChainLength; i++ {
		h := md5_hash(curr, rbt.Prefix)
		if bytes.Compare(h, needle) == 0 {
			number := fmt.Sprintf("%s%09v", rbt.Prefix, curr)
			return number
		}
		curr = reduce(h, i, uint32(KeySpace))
	}
	return ""
}

func (rbt *MobileRainbow) generate_chain(start uint32) uint32 {
	curr := start
	for i := 0; i < rbt.ChainLength; i++ {
		h := md5_hash(curr, rbt.Prefix)
		curr = reduce(h, i, uint32(KeySpace))
	}
	return curr
}

func (rbt *MobileRainbow) FillMissing(db *bolt.DB) {
	for i := 0; i < KeySpace; i++ {
		number := fmt.Sprintf("%s%09v", rbt.Prefix, i)
		hash := md5_hash(uint32(i), rbt.Prefix)
		found := rbt.Lookup(hash)
		if found == "" {
			db.Update(func(tx *bolt.Tx) error {
				b := tx.Bucket([]byte("MissingNumbers"))
				err := b.Put(hash, []byte(number))
				return err
			})
		}
		if i % 10000 == 0 {
			fmt.Println(i)
		}
	}
}

func md5_hash(key uint32, prefix string) []byte {
	number := fmt.Sprintf("%s%09v", prefix, key)
	h := md5.New()
	h.Write([]byte(number))
	return h.Sum(nil)
}

func reduce(hash []byte, index int, key_space uint32) uint32 {
	seg := binary.BigEndian.Uint32(hash[:4])
	return (seg + uint32(index)) % key_space
}

func key_generator(channel chan uint32, keyspace uint32) {
	buff := make([]byte, 4)
	var key uint32 = 0
	for {
		_, err := rand.Read(buff)
		if err != nil {
			return
		}
		key = binary.BigEndian.Uint32(buff)
		channel <- key % keyspace
	}
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func run_server(data_dir string, table_size int, chain_length int) {
	rbtables := []MobileRainbow {
		{TableSize: table_size, ChainLength: chain_length, Prefix: "13"},
		{TableSize: table_size, ChainLength: chain_length, Prefix: "18"},
		{TableSize: table_size, ChainLength: chain_length, Prefix: "15"},
		{TableSize: table_size, ChainLength: chain_length, Prefix: "16"},
		{TableSize: table_size, ChainLength: chain_length, Prefix: "17"},
		{TableSize: table_size, ChainLength: chain_length, Prefix: "14"},
		{TableSize: table_size, ChainLength: chain_length, Prefix: "19"},
	}
	start := time.Now()
	for i := range rbtables {
		fmt.Println("Loading ", rbtables[i].Prefix)
		rbtables[i].LoadDB(data_dir + "/rainbow." + rbtables[i].Prefix)
	}
	elapsed := time.Now().Sub(start)
	fmt.Printf("DB load time: %v \n", elapsed)

	db, err := bolt.Open(data_dir + "/exceptions.db", 0666, nil)
	if err != nil {
		panic(err)
	}
	defer db.Close()
	err1 := db.Update(func(tx *bolt.Tx) error {
		_, e := tx.CreateBucketIfNotExists([]byte("MissingNumbers"))
		return e
	})
	if err1 != nil {
		panic(err1)
	}

	http.HandleFunc("/check", func(w http.ResponseWriter, r *http.Request) {
		query := r.URL.Query()
		hash_str := query.Get("hash")
		hash, _ := hex.DecodeString(hash_str)
		mobile := query.Get("mobile")

		table_index := map[string]int {"13": 0, "18": 1, "15": 2, "16": 3, "17": 4, "14": 5, "19": 6}
		ti := table_index[mobile[:2]]
		table := rbtables[ti]

		number := table.Lookup(hash)
		if number == "" {
			db.Update(func(tx *bolt.Tx) error {
				b := tx.Bucket([]byte("MissingNumbers"))
				err := b.Put(hash, []byte(mobile))
				return err
			})
			fmt.Fprintf(w, "Missing, saved to db")
		} else {
			fmt.Fprintf(w, "Hit")
		}
	})

	http.HandleFunc("/query", func(w http.ResponseWriter, r *http.Request) {
		hash_str := r.URL.RawQuery
		hash_bytes, _ := hex.DecodeString(hash_str)

		var number []byte
		db.View(func(tx *bolt.Tx) error {
			b := tx.Bucket([]byte("MissingNumbers"))
			v := b.Get(hash_bytes)
			if v != nil {
				number = v
			}
			return nil
		})
		if number != nil {
			fmt.Fprintf(w, "Mobile Number: %s\n", string(number))
			return
		} else {
			for i := range rbtables {
				mnumber := rbtables[i].Lookup(hash_bytes)
				if mnumber != "" {
					fmt.Fprintf(w, "Mobile Number: %s\n", mnumber)
					return
				}
			}
			fmt.Fprintf(w, "Not Found")
		}
	})
	http.ListenAndServe(":8080", nil)
}


func fill_missing(data_dir string, prefix string, table_size int, chain_length int) {
	rbt := MobileRainbow{TableSize: table_size, ChainLength: chain_length, Prefix: prefix}
	rbt.LoadDB(data_dir + "/rainbow." + rbt.Prefix)
	fmt.Println("Database loaded")

	db, err := bolt.Open(data_dir + "/exceptions.db", 0666, nil)
	if err != nil {
		panic(err)
	}
	defer db.Close()
	err1 := db.Update(func(tx *bolt.Tx) error {
		_, e := tx.CreateBucketIfNotExists([]byte("MissingNumbers"))
		return e
	})
	if err1 != nil {
		panic(err1)
	}
	rbt.FillMissing(db)
}

func build_db(data_dir string, prefix string, table_size int, chain_length int) {
	rbt := MobileRainbow{TableSize: table_size, ChainLength: chain_length, Prefix: prefix}
	rbt.BuildDB(data_dir + "/rainbow." + rbt.Prefix)
}

func main() {
	op := flag.String("op", "", "Operation to run: fill_missing, build_db, server")
	data_dir := flag.String("data_dir", "", "Data directory")
	table_size := flag.Int("table_size", 100000000, "Rainbowtable lines")
	chain_length := flag.Int("chain_length", 64, "Hash chain length")
	prefix := flag.String("prefix", "", "Mobile number first two number")
	flag.Parse()

	if *op == "fill_missing" {
		fill_missing(*data_dir, *prefix, *table_size, *chain_length)
	} else if *op == "server" {
		run_server(*data_dir, *table_size, *chain_length)
	} else if *op == "build_db" {
		build_db(*data_dir, *prefix, *table_size, *chain_length)
	} else {
		fmt.Println(*op)
	}
}