refactor

Author: enrichman

decode.go

@@ -3,32 +3,95 @@ package main
 import (
 	"crypto/aes"
 	"crypto/sha256"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
 	"gaen/export"
 	"io"
+	"io/ioutil"
 	"strconv"
 	"strings"
 	"time"
 
 	"golang.org/x/crypto/hkdf"
+	"google.golang.org/protobuf/proto"
 )
 
-// DecodeFromTEK ..
-func DecodeFromTEK(tek *export.TemporaryExposureKey) []RollingProximityIdentifier {
-	rpis := make([]RollingProximityIdentifier, 0)
+func DecodeFromFile(filename string) error {
+	export, err := UnmarshalExportFile(filename)
+	if err != nil {
+		return err
+	}
+	teks := DecodeExport(export)
+
+	// print something
+	fmt.Printf("\nTEK: [%s] - [%s]\n", base64.StdEncoding.EncodeToString(teks[0].ID), teks[0].IDBytes)
+	b, _ := json.MarshalIndent(teks[0].RPIs[0], "", "\t")
+	fmt.Printf("\nRPI:\n%v\n", string(b))
+
+	return nil
+}
+
+func UnmarshalExportFile(filename string) (*export.TemporaryExposureKeyExport, error) {
+	in, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return nil, err
+	}
+	in = in[16:]
+
+	export := &export.TemporaryExposureKeyExport{}
+	if err := proto.Unmarshal(in, export); err != nil {
+		return nil, err
+	}
+	return export, nil
+}
+
+func DecodeExport(export *export.TemporaryExposureKeyExport) []TemporaryExposureKey {
+	teks := make([]TemporaryExposureKey, 0)
 
-	hkdfReader := hkdf.New(sha256.New, tek.KeyData, nil, []byte("EN-RPIK"))
+	for _, tek := range export.Keys {
+		tek := NewTemporaryExposureKeyFromGaenTEK(tek)
+		tek.RPIs = DecodeTEK(tek)
+		teks = append(teks, tek)
+	}
+
+	return teks
+}
+
+// DecodeFromTEK ..
+func DecodeTEK(tek TemporaryExposureKey) []RollingProximityIdentifier {
+	hkdfReader := hkdf.New(sha256.New, tek.ID, nil, []byte("EN-RPIK"))
 	rpiKey := make([]byte, 16)
 	io.ReadFull(hkdfReader, rpiKey)
 
-	for rp := 0; rp < int(*tek.RollingPeriod); rp++ {
-		interval := rp + int(*tek.RollingStartIntervalNumber)
+	rpis := make([]RollingProximityIdentifier, 0)
+	for rp := 0; rp < tek.RollingPeriod; rp++ {
+		interval := rp + tek.RollingStartInterval
 		newRpi := NewRollingProximityIdentifier(rpiKey, interval)
 		rpis = append(rpis, newRpi)
 	}
 
 	return rpis
 }
 
+type TemporaryExposureKey struct {
+	ID                   []byte `json:"id"`
+	IDBytes              string `json:"id_bytes"`
+	RollingPeriod        int
+	RollingStartInterval int
+	RPIs                 []RollingProximityIdentifier `json:"rpis"`
+}
+
+func NewTemporaryExposureKeyFromGaenTEK(tek *export.TemporaryExposureKey) TemporaryExposureKey {
+	return TemporaryExposureKey{
+		ID:                   tek.KeyData,
+		IDBytes:              encodeToHexString(tek.KeyData),
+		RollingPeriod:        int(*tek.RollingPeriod),
+		RollingStartInterval: int(*tek.RollingStartIntervalNumber),
+		RPIs:                 make([]RollingProximityIdentifier, 0),
+	}
+}
+
 // RollingProximityIdentifier ..
 type RollingProximityIdentifier struct {
 	ID       []byte    `json:"id"`

main.go

@@ -1,14 +1,9 @@
 package main
 
 import (
-	"encoding/base64"
-	"encoding/json"
 	"fmt"
-	"gaen/export"
-	"io/ioutil"
 
 	"github.com/spf13/cobra"
-	"google.golang.org/protobuf/proto"
 )
 
 var version = "0.0.0-dev"
@@ -31,7 +26,7 @@ var decodeCmd = &cobra.Command{
 	Short: "Decode a TEK export binary file",
 	Args:  cobra.MinimumNArgs(1),
 	RunE: func(cmd *cobra.Command, args []string) error {
-		return decode(args[0])
+		return DecodeFromFile(args[0])
 	},
 }
 
@@ -50,30 +45,3 @@ func main() {
 	rootCmd.AddCommand(downloadCmd)
 	rootCmd.Execute()
 }
-
-func decode(filename string) error {
-	in, err := ioutil.ReadFile(filename)
-	if err != nil {
-		return err
-	}
-	in = in[16:]
-
-	teke := &export.TemporaryExposureKeyExport{}
-	if err := proto.Unmarshal(in, teke); err != nil {
-		return err
-	}
-
-	for _, tek := range teke.Keys {
-		fmt.Printf("\nTEK: [%s] - [%s]\n", base64.StdEncoding.EncodeToString(tek.KeyData), encodeToHexString(tek.KeyData))
-
-		rpis := DecodeFromTEK(tek)
-		for _, rpi := range rpis {
-			b, _ := json.MarshalIndent(rpi, "", "\t")
-			fmt.Printf("\nRPI:\n%v\n", string(b))
-			break
-		}
-		break
-	}
-
-	return nil
-}