"""
The latest version of this package is available at:
<http://github.com/jantman/biweeklybudget>
################################################################################
Copyright 2016 Jason Antman <jason@jasonantman.com> <http://www.jasonantman.com>
This file is part of biweeklybudget, also known as biweeklybudget.
biweeklybudget is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
biweeklybudget is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License
along with biweeklybudget. If not, see <http://www.gnu.org/licenses/>.
The Copyright and Authors attributions contained herein may not be removed or
otherwise altered, except to add the Author attribution of a contributor to
this work. (Additional Terms pursuant to Section 7b of the AGPL v3)
################################################################################
While not legally required, I sincerely request that anyone who finds
bugs please submit them at <https://github.com/jantman/biweeklybudget> or
to me via email, and that you send any contributions or improvements
either as a pull request on GitHub, or to me via email.
################################################################################
AUTHORS:
Jason Antman <jason@jasonantman.com> <http://www.jasonantman.com>
################################################################################
"""
from datetime import datetime
import os
import argparse
import logging
import atexit
from copy import deepcopy
from io import StringIO
import importlib
from ofxparse import OfxParser
from ofxclient.account import Account as OfxClientAccount
from biweeklybudget.utils import Vault
from biweeklybudget import settings
from biweeklybudget.models.account import Account
from biweeklybudget.db import init_db, db_session, cleanup_db
from biweeklybudget.ofxupdater import OFXUpdater
from biweeklybudget.cliutils import set_log_debug, set_log_info
logger = logging.getLogger(__name__)
# suppress requests logging
requests_log = logging.getLogger("requests")
requests_log.setLevel(logging.WARNING)
requests_log.propagate = True
[docs]class OfxGetter(object):
@staticmethod
[docs] def accounts():
"""
Return a sorted list of all Account objects that are for ofxgetter.
"""
return [a for a in db_session.query(
Account).filter(
Account.for_ofxgetter).order_by(Account.name).all()
]
def __init__(self, savedir='./'):
self.savedir = savedir
self._account_data = {}
for acct in self.accounts():
if acct.vault_creds_path is None and acct.ofxgetter_config == {}:
continue
self._account_data[acct.name] = {
'vault_path': acct.vault_creds_path,
'config': acct.ofxgetter_config,
'id': acct.id,
'cat_memo': acct.ofx_cat_memo_to_name
}
logger.info('Initialized with data for %d accounts',
len(self._account_data))
self._accounts = {}
self.vault = Vault()
for acct_name in self._account_data.keys():
data = self._account_data[acct_name]['config']
vault_path = self._account_data[acct_name]['vault_path']
logger.debug('Getting secrets for account %s', acct_name)
secrets = self.vault.read(vault_path)
data['institution']['password'] = secrets['password']
data['institution']['username'] = secrets['username']
if 'class_name' not in data:
self._accounts[acct_name] = OfxClientAccount.deserialize(data)
logger.info('Initialized %d accounts', len(self._accounts))
self.now_str = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
[docs] def get_ofx(self, account_name, write_to_file=True, days=30):
"""
Download OFX from the specified account. Return it as a string.
:param account_name: account name to download
:type account_name: str
:param write_to_file: if True, also write to a file named
"<account_name>_<date stamp>.ofx"
:type write_to_file: bool
:param days: number of days of data to download
:type days: int
:return: OFX string
:rtype: str
"""
fname = None
logger.info('Downloading OFX for account: %s', account_name)
if 'class_name' in self._account_data[account_name]['config']:
ofxdata = self._get_ofx_scraper(account_name, days=days)
else:
acct = self._accounts[account_name]
logger.info(
'Disabling logging for ofxclient, which has bad logging'
)
oldlvl = logging.getLogger().getEffectiveLevel()
logging.getLogger().setLevel(logging.WARNING)
ofxdata = acct.download(days=days).read()
logging.getLogger().setLevel(oldlvl)
logger.info('Re-enabling logging')
if write_to_file:
fname = self._write_ofx_file(account_name, ofxdata)
self._ofx_to_db(account_name, fname, ofxdata)
return ofxdata
[docs] def _ofx_to_db(self, account_name, fname, ofxdata):
"""
Put OFX Data to the DB
:param account_name: account name to download
:type account_name: str
:param ofxdata: raw OFX data
:type ofxdata: str
:param fname: filename OFX was written to
:type fname: str
"""
logger.debug('Parsing OFX')
ofx = OfxParser.parse(StringIO(ofxdata))
logger.debug('Instantiating OFXUpdater')
updater = OFXUpdater(
self._account_data[account_name]['id'],
account_name,
cat_memo=self._account_data[account_name]['cat_memo']
)
logger.debug('Updating OFX in DB')
updater.update(ofx, filename=fname)
db_session.commit()
logger.debug('Done updating OFX in DB')
[docs] def _get_ofx_scraper(self, account_name, days=30):
"""
Get OFX via a ScreenScraper subclass.
:param account_name: account name
:type account_name: str
:param days: number of days of data to download
:type days: int
:return: OFX string
:rtype: str
"""
data = self._account_data[account_name]['config']
clsname = data['class_name']
modname = data['module_name']
logger.debug('Scraper - getting class %s from module %s',
clsname, modname)
cls = getattr(
importlib.import_module(modname),
clsname
)
logger.debug('Getting secrets for account %s', account_name)
secrets = self.vault.read(
self._account_data[account_name]['vault_path']
)
if 'kwargs' in data:
kwargs = deepcopy(data['kwargs'])
else:
kwargs = {}
kwargs['username'] = secrets['username']
kwargs['password'] = secrets['password']
kwargs['savedir'] = os.path.join(self.savedir, account_name)
acct = cls(**kwargs)
ofxdata = acct.run()
return ofxdata
[docs] def _write_ofx_file(self, account_name, ofxdata):
"""
Write OFX data to a file.
:param account_name: account name
:type account_name: str
:param ofxdata: raw OFX data string
:type ofxdata: str
:returns: name of the file that was written
:rtype: str
"""
if not os.path.exists(os.path.join(self.savedir, account_name)):
os.makedirs(os.path.join(self.savedir, account_name))
fname = '%s_%s.ofx' % (account_name, self.now_str)
fpath = os.path.join(self.savedir, account_name, fname)
logger.debug('Writing %d bytes of OFX to: %s', len(ofxdata), fpath)
with open(fpath, 'w') as fh:
fh.write(ofxdata)
logger.info('Wrote OFX data to: %s', fpath)
return fname
[docs]def parse_args():
p = argparse.ArgumentParser(description='Download OFX transactions')
p.add_argument('-v', '--verbose', dest='verbose', action='count', default=0,
help='verbose output. specify twice for debug-level output.')
p.add_argument('-l', '--list-accts', dest='list', action='store_true',
help='list accounts and exit')
p.add_argument('ACCOUNT_NAME', type=str, action='store', default=None,
nargs='?',
help='Account name; omit to download all accounts')
args = p.parse_args()
return args
[docs]def main():
global logger
format = "[%(asctime)s %(levelname)s] %(message)s"
logging.basicConfig(level=logging.WARNING, format=format)
logger = logging.getLogger()
args = parse_args()
# set logging level
if args.verbose > 1:
set_log_debug(logger)
elif args.verbose == 1:
set_log_info(logger)
atexit.register(cleanup_db)
init_db()
if args.list:
for k in OfxGetter.accounts():
print(k.name)
raise SystemExit(0)
getter = OfxGetter(settings.STATEMENTS_SAVE_PATH)
if args.ACCOUNT_NAME is not None:
getter.get_ofx(args.ACCOUNT_NAME)
raise SystemExit(0)
# else all of them
total = 0
success = 0
for acct in OfxGetter.accounts():
try:
total += 1
getter.get_ofx(acct.name)
success += 1
except Exception:
logger.error(
'Failed to download account %s', acct.name, exc_info=True
)
if success != total:
logger.warning('Downloaded %d of %d accounts', success, total)
raise SystemExit(1)
raise SystemExit(0)
if __name__ == "__main__":
main()