From 11a578995276a89a9befd06fff71f7c064518e97 Mon Sep 17 00:00:00 2001
From: Anshal Dwivedi <anshaldwivedi@yahoo.com>
Date: Sun, 27 Mar 2022 03:23:40 +0530
Subject: [PATCH] add support for custom otp length

---
 src/TOTP.cpp | 41 +++++++++++++++++++++++++++++++++++------
 src/TOTP.h   |  6 ++++--
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/src/TOTP.cpp b/src/TOTP.cpp
index 0aa1e40..2f3da27 100644
--- a/src/TOTP.cpp
+++ b/src/TOTP.cpp
@@ -7,6 +7,8 @@
 #include "TOTP.h"
 #include "sha1.h"
 
+const int DEFAULT_CODE_LEN = 6;
+
 // Init the library with the private key, its length and the timeStep duration
 TOTP::TOTP(uint8_t* hmacKey, int keyLength, int timeStep) {
 
@@ -23,16 +25,42 @@ TOTP::TOTP(uint8_t* hmacKey, int keyLength) {
 	_timeStep = 30;
 };
 
+long pow10(int n) {
+    static long pow10[10] = {
+        1, 10, 100, 1000, 10000, 
+        100000, 1000000, 10000000, 100000000, 1000000000,
+    };
+
+    return pow10[n]; 
+} 
+
+char* getFormatString(int codeLen) {
+	static char format[10][6] = {
+        "%01ld", "%02ld", "%03ld", "%04ld", "%05ld", 
+         "%06ld", "%07ld",  "%08ld",  "%09ld",
+    };
+
+	return format[codeLen-1];
+}
+
 // Generate a code, using the timestamp provided
 char* TOTP::getCode(long timeStamp) {
-	
+	return getCode(timeStamp, DEFAULT_CODE_LEN);
+}
+
+// Generate a code of specified length, using the timestamp provided  
+char* TOTP::getCode(long timeStamp, int codeLen) {
 	long steps = timeStamp / _timeStep;
-	return getCodeFromSteps(steps);
+	return getCodeFromSteps(steps, codeLen);
 }
 
 // Generate a code, using the number of steps provided
 char* TOTP::getCodeFromSteps(long steps) {
-	
+	return getCodeFromSteps(steps, DEFAULT_CODE_LEN);
+
+}
+
+char *TOTP::getCodeFromSteps(long steps, int codeLen) {
 	// STEP 0, map the number of steps in a 8-bytes array (counter value)
 	_byteArray[0] = 0x00;
 	_byteArray[1] = 0x00;
@@ -58,9 +86,10 @@ char* TOTP::getCodeFromSteps(long steps) {
 
 	// STEP 3, compute the OTP value
 	_truncatedHash &= 0x7FFFFFFF;
-	_truncatedHash %= 1000000;
+	_truncatedHash %= pow10(codeLen);
 	
-	// convert the value in string, with heading zeroes
-	sprintf(_code, "%06ld", _truncatedHash);
+	sprintf(_code, getFormatString(codeLen), _truncatedHash);
+
 	return _code;
 }
+
diff --git a/src/TOTP.h b/src/TOTP.h
index a69e3df..a04fe1f 100644
--- a/src/TOTP.h
+++ b/src/TOTP.h
@@ -16,8 +16,10 @@ class TOTP {
 		TOTP(uint8_t* hmacKey, int keyLength);
 		TOTP(uint8_t* hmacKey, int keyLength, int timeStep);
 		char* getCode(long timeStamp);
+		char* getCode(long timeStamp, int codeLen);  
 		char* getCodeFromSteps(long steps);
-	
+		char* getCodeFromSteps(long steps, int codeLen);
+
 	private:
 
 		uint8_t* _hmacKey;
@@ -27,7 +29,7 @@ class TOTP {
 		uint8_t* _hash;
 		int _offset;
 		long _truncatedHash;
-		char _code[7];
+		char _code[10];
 };
 
 #endif
\ No newline at end of file