2 files changed, 23 insertions, 2 deletions

diff --git a/README.md b/README.md
index 6fd582f..3932606 100644
--- a/README.md
+++ b/README.md
@@ -13,8 +13,16 @@ $ ./decode-qr-uri.py 'igmobileotp://?action=secactivate&enc=VRUq6IoLWQRCMRITZEHt
 # generate-otp.py
 Once you have the required information from a QR code, you can combine it with a "registration code" to derive the OTP secret. This registration code contains random bytes that were generated on the end-users device (their mobile phone), and are thus required to determine the OTP secret. An example way to obtain all of this information would be through email, if the user recieves a QR code in their email, then responds with their registration code.
 
+The OTP secret optionally includes the policy specification, which is provided as part of the QR code. *If you are having problems generating a valid OTP secret, try with or without the policy parameter*.
+
 Example:
 ```
+$ ./generate-otp.py 48244-13456 1745-7712-6942-8698 12211-49352 --policy '{"allowUnsecured":"false","trustedExecution":"NOT_ALLOWED"}'
+bb9b6d72ae99b006de5e106935ec96da
+
+To generate a code immediately, run:
+oathtool -v --totp=sha256 --digits=6 bb9b6d72ae99b006de5e106935ec96da
+
 $ ./generate-otp.py 48244-13456 1745-7712-6942-8698 12211-49352
 9a8eab5ecc9fc413758a92ac223dc6a0
 
diff --git a/generate-otp.py b/generate-otp.py
index a931a77..82925c0 100755
--- a/generate-otp.py
+++ b/generate-otp.py
@@ -5,10 +5,14 @@ import logging
 
 logging.basicConfig(level=logging.WARNING)
 
-parser = argparse.ArgumentParser(description='Generate an OTP secret for an Entrust IdentityGuard soft token')
+parser = argparse.ArgumentParser(
+	description='Generate an OTP secret for an Entrust IdentityGuard soft token',
+	epilog='If your token does not work, try without the Policy argument, as in some cases, this is not used to generate the OTP secret'
+)
 parser.add_argument('Serial', type=str, nargs=1, help='Given to the user (such as through a QR code). Example: 48244-13456')
 parser.add_argument('ActivationCode', type=str, nargs=1, help='Given to the user (such as through a QR code). Example: 1745-7712-6942-8698')
 parser.add_argument('RegistrationCode', type=str, nargs=1, help='The user provides this to the activation service. Example: 12211-49352')
+parser.add_argument('--policy', type=str, nargs=1, required=False, help='The policy associated with the identity. Example: {"allowUnsecured":"false","trustedExecution":"NOT_ALLOWED"}')
 args = parser.parse_args()
 
 # Remove dashes from input so we can work with the data
@@ -32,10 +36,19 @@ rngbytes = registrationbytes[-2:]
 
 logging.info("RNG Bytes: 0x%s", rngbytes.hex())
 
+password = activationbytes + rngbytes
+
+# The secret may or may not include the policy
+if args.policy is not None:
+	password += args.policy[0].encode('utf-8')
+	logging.info("Policy: %s", args.policy[0].encode('utf-8'))
+else:
+	logging.debug("Policy not provided")
+
 # Derive the secret key
 key = pbkdf2_hmac(
     hash_name='sha256',
-    password=activationbytes + rngbytes,
+    password=password,
     salt=serial.encode("utf-8"),
     iterations=8,
     dklen=16