From 281f723f6a2c3cc66eecba83accfc92556412d78 Mon Sep 17 00:00:00 2001
From: Jason Lai <jasonlai256@gmail.com>
Date: Sat, 13 Oct 2012 00:12:13 +0800
Subject: [PATCH] commit codes and some bugs need to be fix.

---
 pypw/handler.py | 76 ++++++++++++++++++++++++++++++-------------------
 1 file changed, 46 insertions(+), 30 deletions(-)

diff --git a/pypw/handler.py b/pypw/handler.py
index e8da073..e6bb09b 100644
--- a/pypw/handler.py
+++ b/pypw/handler.py
@@ -9,9 +9,6 @@
     :license: BSD, see LICENSE for more details.
 """
 
-from Crypto.Cipher import AES
-from Crypto import Random
-
 import os
 import json
 import logging
@@ -19,6 +16,9 @@
 from datetime import datetime
 from collections import defaultdict
 
+from Crypto.Cipher import AES
+from Crypto import Random
+
 
 class AESHandler(object):
     """
@@ -34,9 +34,9 @@ def __init__(self):
         self.filepath = ''
         self.cipher = ''
         self.data = None
-        # self.records is a proxy structure mapping to the records of 
+        # self._records is a proxy structure mapping to the records of 
         # self.data and is use for better contrive records.
-        self.records = defaultdict(dict)
+        self._records = defaultdict(dict)
 
     def initialize(self, cipher, filepath='records.dat'):
         """
@@ -48,22 +48,24 @@ def initialize(self, cipher, filepath='records.dat'):
         self.cipher = self.figure_32Byte_key(cipher)
         
         # first initial the program
-        if not os.path.existx(self.filepath):
-            self.data = compose_structure(cipher)
+        if not os.path.exists(self.filepath) or \
+                               not os.path.getsize(self.filepath):
+            # initiali the self.data and self._records
+            self._init_data()
             self.initialized = True
             return self.initialized
 
         # validate the cipher and load the data
         with open(filepath, 'rb') as f:
             ciphertext = f.read()
-            jsondata = self.decrypt(ciphertext, key)
+            jsondata = self.decrypt(ciphertext, self.cipher)
 
         try:
             self.data = json.loads(jsondata)
             if self.data['digest'] == self.cipher:
                 self.initialized = True
-            # initiali the self.data and self.records
-            self._init_data()
+            # initiali the self._records
+            self._setup_structure()
         except ValueError:
             self.log.error('Error occur when load the JSON text.')
         except Exception, err:
@@ -73,34 +75,42 @@ def initialize(self, cipher, filepath='records.dat'):
 
     def write(self):
         """
-        encrypt the infomations and dumps into outside file.
+        encrypt the infomations and dump into outside file.
         """
         jsontext = json.dumps(self.data)
         with open(self.filepath, 'w') as f:
             ciphertext = self.encrypt(jsontext, self.cipher)
             f.write(ciphertext)
 
-    def add_record(group, item, value, note):
+    def add_record(self, group, item, value, note=None):
         try:
             record = self._compose_record(group, item, value, note)
             self.data['records'].append(record)
             self.data['currentID'] += 1
-            self.data['currentGID'] += 1
+            if not self._records['gid'].has_key(record['gid']):
+                self.data['currentGID'] += 1
             self.write()
             self._adjust_structure(record)
+            return True
         except Exception, err:
             self.log.error('Error occur in adding record - %s', err)
+            raise
             return False
-        return True
 
-    def update_record(group, item, value, note):
-        try:
-            # TODO: complete the method
-            self.write()
-        except Exception, err:
-            self.log.error('error occur in updating record - %s', err)
+    def update_record(self, group, item, value, note=None):
+        if self._records.has_key(group) and self._records[group].has_key(item):
+            try:
+                record = self._records[group][item]
+                record['value'] = value
+                if note:
+                    record['note'] = note
+                self.write()
+                return True
+            except Exception, err:
+                self.log.error('Error occur in updating record - %s', err)
+                return False
+        else:
             return False
-        return True
 
     @property
     def records(self):
@@ -121,19 +131,19 @@ def _setup_structure(self):
             self._adjust_structure(record)
 
     def _adjust_structure(self, record):
-        rid = record['id']
-        gid = record['gid']
+        rid, gid = record['id'], record['gid']
         group = record['group']
         item = record['itemname']
-        self.records['rid'][rid] = record
-        self.records['gid'].setdefault(gid, []).append(record)
-        self.records[group][item] = record
+        self._records['rid'][rid] = record
+        self._records['gid'].setdefault(gid, []).append(record)
+        self._records[group][item] = record
 
-    def _compose_record(group, item, value, note=None):
+    def _compose_record(self, group, item, value, note=None):
         created = datetime.today().isoformat('_')
+        # TODO: should judge the correct gid
         record = {
-            'id': self.data['currentID'] + 1,
-            'gid': self.data['currentGID'] + 1,
+            'id': self.data['currentID'],
+            'gid': self.data['currentGID'],
             'group': group,
             'itemname': item,
             'value': value,
@@ -147,7 +157,13 @@ def _compose_record(group, item, value, note=None):
     def figure_32Byte_key(cls, text):
         sha = hashlib.sha256()
         sha.update(text)
-        return sha.hexdigest()[::2]
+        digest = sha.hexdigest()[::2]
+
+        if len(text) > 32:
+            return digest
+        else:
+            keystr = text + digest
+            return keystr[:32]
 
     @classmethod
     def encrypt(cls, plaintext, key):