92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
import json
|
|
import pykeepass.icons
|
|
from urllib.parse import quote_plus, urlparse
|
|
from pykeepass import create_database
|
|
|
|
|
|
class KpWriter:
|
|
def __init__(self, filename, password="test"):
|
|
self.kp = create_database(filename, password)
|
|
self.current_entry = None
|
|
|
|
def add_entry(self, dest_group_name, title):
|
|
# Find group and create if not yet there
|
|
group = self.kp.find_groups(name=dest_group_name, first=True)
|
|
if not group:
|
|
# TODO: Handle nested groups?
|
|
group = self.kp.add_group(self.kp.root_group, dest_group_name)
|
|
|
|
suffix_ctr = 1
|
|
new_title = title
|
|
while self.kp.find_entries(title=new_title, group=group):
|
|
suffix_ctr += 1
|
|
new_title = "{} - {}".format(title, suffix_ctr)
|
|
|
|
self.current_entry = self.kp.add_entry(group, new_title, "", "")
|
|
return self.current_entry
|
|
|
|
def save(self):
|
|
return self.kp.save()
|
|
|
|
def set_icon(self, icon_id):
|
|
if icon_id in dir(pykeepass.icons):
|
|
kp_icon_id = getattr(pykeepass.icons, icon_id)
|
|
else:
|
|
# FIXME: Assume kp_icon is already ID, needed b/c icon 12 is missing from pykeepass.icons
|
|
kp_icon_id = icon_id
|
|
self.current_entry.icon = kp_icon_id
|
|
|
|
def set_tags(self, tag_list):
|
|
self.current_entry.tags = tag_list
|
|
|
|
def add_totp(self, init_string, otp_url=None, title=""):
|
|
if not otp_url:
|
|
otp_url = "otpauth://totp/Sample:username?secret={}&algorithm=SHA1&digits=6&period=30&issuer=Sample".format(quote_plus(init_string))
|
|
|
|
# It's possible to define multiple OTP-secrets in 1P7, so let's not lose one
|
|
suffix = ""
|
|
suffix_ctr = 1
|
|
while self.current_entry.get_custom_property("otp{}".format(suffix)):
|
|
suffix_ctr += 1
|
|
suffix = "_{}".format(suffix_ctr)
|
|
|
|
self.set_prop("TimeOtp-Secret-Base32{}".format(suffix), init_string, True)
|
|
self.set_prop("otp{}".format(suffix), otp_url, True)
|
|
if len(title) > 0:
|
|
self.set_prop("otp_title{}".format(suffix), title)
|
|
|
|
def add_url(self, url):
|
|
if not self.current_entry.url:
|
|
self.current_entry.url = url
|
|
else:
|
|
if url == self.current_entry.url:
|
|
return False
|
|
# https://github.com/keepassxreboot/keepassxc/pull/3558
|
|
suffix = ""
|
|
suffix_ctr = 0
|
|
while self.current_entry.get_custom_property("KP2A_URL{}".format(suffix)):
|
|
suffix_ctr += 1
|
|
suffix = "_{}".format(suffix_ctr)
|
|
self.set_prop("KP2A_URL{}".format(suffix), url)
|
|
|
|
# KeePassHttp
|
|
current_settings = self.current_entry.get_custom_property("KeePassHttp Settings")
|
|
if current_settings:
|
|
current_settings = json.loads(current_settings)
|
|
else:
|
|
current_settings = {
|
|
"Allow": [],
|
|
"Deny": [],
|
|
"Realm": "",
|
|
}
|
|
parsed_url = urlparse(url)
|
|
current_settings["Allow"].append(parsed_url.hostname)
|
|
current_settings["Allow"] = list(set(current_settings["Allow"]))
|
|
self.set_prop("KeePassHttp Settings", json.dumps(current_settings))
|
|
|
|
def set_prop(self, key, value, protected=False):
|
|
self.current_entry.set_custom_property(key, value)
|
|
if protected:
|
|
# https://github.com/libkeepass/pykeepass/issues/89
|
|
self.current_entry._element.xpath('String[Key[text()="{}"]]/Value'.format(key))[0].attrib["Protected"] = "True"
|