Locksly/vaultlogic.py at main · BoneheadBreaker/Locksly · GitHub

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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import os
import json
import base64
from pathlib import Path

from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from cryptography.hazmat.backends import default_backend
from cryptography.fernet import Fernet, InvalidToken
import pyotp # for generating totp codes

class VaultLockedError(Exception):
    pass


class Vault:
    def __init__(self, path):
        self.path = Path(path)
        self.key = None
        self.data = None

    def _derive_key(self, master_password, salt):
        kdf = Scrypt(
            salt=salt,
            length=32,
            n=2**14,
            r=8,
            p=1,
            backend=default_backend()
        )
        return base64.urlsafe_b64encode(kdf.derive(master_password.encode()))

    def _require_unlocked(self):
        if self.key is None or self.data is None:
            raise VaultLockedError("Vault is locked")

    def create(self, master_password, overwrite=False):
        self.path.parent.mkdir(parents=True, exist_ok=True)

        if self.path.exists() and not overwrite:
            raise FileExistsError("Vault already exists")

        salt = os.urandom(16)
        self.key = self._derive_key(master_password, salt)
        self.data = {"entries": []}

        encrypted = Fernet(self.key).encrypt(json.dumps(self.data).encode())

        with open(self.path, "w") as f:
            json.dump({
                "salt": base64.b64encode(salt).decode(),
                "vault": encrypted.decode()
            }, f, indent=2)

    def unlock(self, master_password):
        if not self.path.exists():
            raise FileNotFoundError("Vault not found")

        with open(self.path, "r") as f:
            vault_file = json.load(f)

        salt = base64.b64decode(vault_file["salt"])
        self.key = self._derive_key(master_password, salt)

        try:
            decrypted = Fernet(self.key).decrypt(vault_file["vault"].encode())
        except InvalidToken:
            self.key = None
            raise ValueError("Invalid master password")

        self.data = json.loads(decrypted)

    def lock(self):
        self.key = None
        self.data = None

    def get_entries(self):
        self._require_unlocked()
        return self.data["entries"]

    @staticmethod
    def mask_string(string):
        return '•' * len(string)

    @staticmethod
    def generate_totp(totp_token):
        totp = pyotp.TOTP(totp_token)
        current_otp = totp.now()

        return current_otp

    def add_entry(self, entry):
        self._require_unlocked()
        self.data["entries"].append(entry)
        self._save()

    def _save(self):
        self.path.parent.mkdir(parents=True, exist_ok=True)

        encrypted = Fernet(self.key).encrypt(json.dumps(self.data).encode())

        vault_file = {}
        if self.path.exists():
            with open(self.path, "r") as f:
                vault_file = json.load(f)

        vault_file["vault"] = encrypted.decode()

        with open(self.path, "w") as f:
            json.dump(vault_file, f, indent=2)