From 062cc5b840c0c9faf190c56e7e567dcdfac08719 Mon Sep 17 00:00:00 2001
From: Benjamin Dauvergne <bdauvergne@entrouvert.com>
Date: Tue, 8 Oct 2013 11:24:20 +0200
Subject: [PATCH] api: add a new procedure get_timestamp to extract the
 timestamp as an UTC naive datetime

---
 rfc3161/api.py | 30 ++++++++++++++++++++++++++----
 tests/api.py   | 13 +++++++------
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/rfc3161/api.py b/rfc3161/api.py
index 9d00245..2c9e543 100644
--- a/rfc3161/api.py
+++ b/rfc3161/api.py
@@ -1,3 +1,8 @@
+import datetime
+import hashlib
+import urllib2
+import base64
+
 from pyasn1.codec.der import encoder, decoder
 from pyasn1_modules import rfc2459
 from pyasn1.type import univ
@@ -6,11 +11,9 @@ import M2Crypto.X509 as X509
 import socket
 
 import rfc3161
-import hashlib
-import urllib2
-import base64
 
-__all__ = ('RemoteTimestamper','check_timestamp','get_hash_oid', 'TimestampingError')
+__all__ = ('RemoteTimestamper','check_timestamp','get_hash_oid',
+    'TimestampingError', 'get_timestamp')
 
 id_attribute_messageDigest = univ.ObjectIdentifier((1,2,840,113549,1,9,4,))
 
@@ -20,6 +23,25 @@ def get_hash_oid(hashname):
 class TimestampingError(RuntimeError):
     pass
 
+def get_timestamp(tst):
+    try:
+        if not isinstance(tst, rfc3161.TimeStampToken):
+            tst, substrate = decoder.decode(tst, asn1Spec=rfc3161.TimeStampToken())
+            if substrate:
+                raise ValueError("extra data after tst")
+
+        tstinfo = tst.getComponentByName('content').getComponentByPosition(2).getComponentByPosition(1)
+        tstinfo, substrate = decoder.decode(tstinfo, asn1Spec=univ.OctetString())
+        if substrate:
+            raise ValueError("extra data after tst")
+        tstinfo, substrate = decoder.decode(tstinfo, asn1Spec=rfc3161.TSTInfo())
+        if substrate:
+            raise ValueError("extra data after tst")
+        genTime = tstinfo.getComponentByName('genTime')
+        return datetime.datetime.strptime(str(genTime), '%Y%m%d%H%M%SZ')
+    except PyAsn1Error, e:
+        raise ValueError('not a valid TimeStampToken', e)
+
 def check_timestamp(tst, certificate, data=None, sha1=None, hashname=None):
     hashobj = hashlib.new(hashname or 'sha1')
     if not sha1:
diff --git a/tests/api.py b/tests/api.py
index 6e52c4d..3cc07e7 100644
--- a/tests/api.py
+++ b/tests/api.py
@@ -1,5 +1,8 @@
 import unittest
 import os.path
+import datetime
+
+from pyasn1.type import univ
 
 import rfc3161
 
@@ -9,11 +12,9 @@ class Rfc3161(unittest.TestCase):
             '../data/certum_certificate.crt')
 
     def test_timestamp(self):
-        try:
-            certificate = file(self.CERTIFICATE).read()
-            value, substrate = rfc3161.RemoteTimestamper(
-                    self.PUBLIC_TSA_SERVER, certificate=certificate)(data='xx')
-        except rfc3161.TimestampingError:
-            return
+        certificate = file(self.CERTIFICATE).read()
+        value, substrate = rfc3161.RemoteTimestamper(
+                self.PUBLIC_TSA_SERVER, certificate=certificate)(data='xx')
+        self.assertIsInstance(rfc3161.get_timestamp(value), datetime.datetime)
         self.assertNotEqual(value, None)
         self.assertEqual(substrate, '')