summaryrefslogtreecommitdiff
path: root/wallet.py
blob: a5126d5313b6928a848784f26820ef63e2c9ad71 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
#!/usr/bin/python
import pickle
import Crypto.Cipher.Blowfish
import os
import getpass
from threading import Lock


class Wallet(dict):
    def __init__(self, filename):
        self.lock = Lock()
        if os.path.isfile(filename):
            self.f = open(filename, "rb+")
            self.passwd = getpass.getpass()
            self.crypt = Crypto.Cipher.Blowfish.new(self.passwd)
            contents_encrypted = self.f.read()
            contents = self.crypt.decrypt(contents_encrypted +
                                          "\x00"*((-len(contents_encrypted))%8))
            if contents.startswith(self.passwd):
                self.update(dict(pickle.loads(contents[len(self.passwd):])))
            else:
                self.f.close()
                raise BaseException("Incorrect Password")
        else:
            self.f = open(filename, "wb+")
            passwd = self.passwd = None
            while passwd is None or passwd != self.passwd:
                passwd = getpass.getpass("Enter new password: ")
                self.passwd = getpass.getpass("Confirm new password: ")
                if passwd != self.passwd:
                    print "Passwords do not match!"
            self.crypt = Crypto.Cipher.Blowfish.new(self.passwd)
            self.flush()

    def flush(self):
        contents = self.passwd+pickle.dumps(self.items(), protocol=2)
        self.lock.acquire()
        try:
            self.f.seek(0)
            self.f.write(self.crypt.encrypt(
                contents+"\x00"*((-len(contents))%8)))
            self.f.truncate()
            self.f.flush()
        finally:
            self.lock.release()