2 files changed, 59 insertions, 49 deletions

diff --git a/lib/oracle.go b/lib/oracle.go
index 005e392..8c7a52b 100644
--- a/lib/oracle.go
+++ b/lib/oracle.go
@@ -3,32 +3,29 @@
 
 package lib
 
-import (
-	"crypto/rand"
-	"math/big"
-)
-
 var oracleUnknown string = `Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkg
 aGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBq
 dXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUg
 YnkK`
 
-var OracleKey []byte = make([]byte, 16)
-var OracleIV []byte = make([]byte, 16)
-var oracleRandom []byte = make([]byte, RandomInt(1, 4096))
+var OracleKey []byte
+var OracleIV []byte
+var oracleRandom []byte
 
 func init() {
-	_, err := rand.Read(OracleKey)
+	var err error
+
+	OracleKey, err = RandomKey(16)
 	if err != nil {
 		panic(err)
 	}
 
-	_, err = rand.Read(OracleIV)
+	OracleIV, err = RandomKey(16)
 	if err != nil {
 		panic(err)
 	}
 
-	_, err = rand.Read(oracleRandom)
+	oracleRandom, err = RandomKey(int(RandomInt(1, 4096)))
 	if err != nil {
 		panic(err)
 	}
@@ -39,14 +36,12 @@ func init() {
 // cipher.
 func OracleAESRandomEncrypt(in []byte) []byte {
 	// Generate random key.
-	key := make([]byte, 16)
-	_, err := rand.Read(key)
+	key, err := RandomKey(16)
 	if err != nil {
 		panic(err)
 	}
 	// Generate random initialization vector; needed for AES CBC.
-	iv := make([]byte, 16)
-	_, err = rand.Read(iv)
+	iv, err := RandomKey(16)
 	if err != nil {
 		panic(err)
 	}
@@ -56,12 +51,9 @@ func OracleAESRandomEncrypt(in []byte) []byte {
 	in = append(in, randomBytes(5, 10)...)
 
 	// Randomly encrypt `in` with AES in ECB or CBC mode.
-	m, err := rand.Int(rand.Reader, big.NewInt(2))
-	if err != nil {
-		panic(err)
-	}
+	m := RandomInt(0, 1)
 	var out []byte
-	if m.Int64() == 0 {
+	if m == 0 {
 		// Encrypt with AES in ECB mode.
 		out = AESEncryptECB(in, key)
 	} else {
@@ -80,32 +72,3 @@ func OracleAESVarEncryptECB(in []byte) []byte {
 	in = append(in, Base64ToBytes(oracleUnknown)...)
 	return AESEncryptECB(in, OracleKey)
 }
-
-// Return a random number from range [min, max]
-func RandomInt(min, max int64) int64 {
-	if min >= max {
-		panic("RandomInt: min cannot be >= max!")
-	}
-
-	var rn *big.Int
-	var err error
-	for {
-		rn, err = rand.Int(rand.Reader, big.NewInt(max+1))
-		if err != nil {
-			panic(err)
-		}
-		if rn.Int64() >= min {
-			return rn.Int64()
-		}
-	}
-}
-
-// Randomly generates `min` to `max` bytes.
-func randomBytes(min, max int64) []byte {
-	bs := make([]byte, RandomInt(min, max))
-	_, err := rand.Read(bs)
-	if err != nil {
-		panic(err)
-	}
-	return bs
-}
diff --git a/lib/rand.go b/lib/rand.go
new file mode 100644
index 0000000..03fd651
--- /dev/null
+++ b/lib/rand.go
@@ -0,0 +1,47 @@
+// Copyright © 2020 rsiddharth <s@ricketyspace.net>
+// SPDX-License-Identifier: ISC
+
+package lib
+
+import (
+	"crypto/rand"
+	"math/big"
+)
+
+// Return a random number from range [min, max]
+func RandomInt(min, max int64) int64 {
+	if min >= max {
+		panic("RandomInt: min cannot be >= max!")
+	}
+
+	var rn *big.Int
+	var err error
+	for {
+		rn, err = rand.Int(rand.Reader, big.NewInt(max+1))
+		if err != nil {
+			panic(err)
+		}
+		if rn.Int64() >= min {
+			return rn.Int64()
+		}
+	}
+}
+
+func RandomKey(size int) ([]byte, error) {
+	k := make([]byte, size)
+	_, err := rand.Read(k)
+	if err != nil {
+		return []byte{}, err
+	}
+	return k, nil
+}
+
+// Randomly generates `min` to `max` bytes.
+func randomBytes(min, max int64) []byte {
+	bs := make([]byte, RandomInt(min, max))
+	_, err := rand.Read(bs)
+	if err != nil {
+		panic(err)
+	}
+	return bs
+}