diff --git a/pykeepass/pykeepass.py b/pykeepass/pykeepass.py index 5a66d761..a5640524 100644 --- a/pykeepass/pykeepass.py +++ b/pykeepass/pykeepass.py @@ -10,7 +10,11 @@ import re import uuid import zlib +import socket from copy import deepcopy +from os.path import exists +from stat import S_ISSOCK +from io import BytesIO from construct import Container, ChecksumError from lxml import etree @@ -128,21 +132,31 @@ def reload(self): self.read(self.filename, self.password, self.keyfile) - def save(self, filename=None, transformed_key=None): - """Save current database object to disk. + def save(self, filename=None, transformed_key=None, force=False): + """Save current database object to disk, buffer or socket. Args: - filename (:obj:`str`, optional): path to database or stream object. + filename (:obj:`str`, optional): path to database, stream object + or socket. If None, the path given when the database was opened is used. PyKeePass.filename is unchanged. transformed_key (:obj:`bytes`, optional): precomputed transformed key. + force (:obj:`bool`, optional): force write if preconditions fail """ + seek_msg = ( + 'Using a non-seekable medium may cause KDBX corruption,' + ' use force=True to override' + ) + output = None if not filename: filename = self.filename if hasattr(filename, "write"): + if not getattr(filename, "seekable", lambda: 0)() and not force: + raise Exception(seek_msg) + output = KDBX.build_stream( self.kdbx, filename, @@ -150,7 +164,27 @@ def save(self, filename=None, transformed_key=None): keyfile=self.keyfile, transformed_key=transformed_key ) + + elif exists(filename) and S_ISSOCK(os.stat(filename).st_mode): + target = BytesIO() + output = KDBX.build_stream( + self.kdbx, + target, + password=self.password, + keyfile=self.keyfile, + transformed_key=transformed_key + ) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + target.seek(0) + sock.connect(filename) + sock.sendall(target.read()) + else: + if exists(filename): + with open(filename) as file: + if not file.seekable() and not force: + raise Exception(seek_msg) + output = KDBX.build_file( self.kdbx, filename,