From 13123deebeecf839e44b0024f67f38d530ab75f2 Mon Sep 17 00:00:00 2001
From: RunasSudo <runassudo@yingtongli.me>
Date: Wed, 4 Jan 2023 18:37:05 +1100
Subject: [PATCH] Implement OFX 2.x import

---
 drcr/statements/importers/ofx1.py     |  6 +--
 drcr/statements/importers/ofx2.py     | 61 +++++++++++++++++++++++++++
 drcr/statements/views.py              |  3 ++
 drcr/templates/statements/import.html | 20 +++++++++
 4 files changed, 87 insertions(+), 3 deletions(-)
 create mode 100644 drcr/statements/importers/ofx2.py

diff --git a/drcr/statements/importers/ofx1.py b/drcr/statements/importers/ofx1.py
index 7c79d3a..b937020 100644
--- a/drcr/statements/importers/ofx1.py
+++ b/drcr/statements/importers/ofx1.py
@@ -24,11 +24,11 @@ from io import StringIO
 def import_ofx1(file):
 	raw_ofx = file.read().decode('utf-8')
 	
-	# Convert OFX header to XML and parse
+	# Convert OFX header to SGML and parse
 	raw_payload = raw_ofx[raw_ofx.index('<OFX>'):]
-	xml_input = StringIO(raw_payload.replace('&', '&amp;'))
+	sgml_input = StringIO(raw_payload.replace('&', '&amp;'))
 	try:
-		tree = ET.parse(xml_input, ET.HTMLParser())
+		tree = ET.parse(sgml_input, ET.HTMLParser())
 	except Exception as ex:
 		raise ex
 	root = tree.getroot()
diff --git a/drcr/statements/importers/ofx2.py b/drcr/statements/importers/ofx2.py
new file mode 100644
index 0000000..5dbc5fb
--- /dev/null
+++ b/drcr/statements/importers/ofx2.py
@@ -0,0 +1,61 @@
+#   DrCr: Web-based double-entry bookkeeping framework
+#   Copyright (C) 2022–2023  Lee Yingtong Li (RunasSudo)
+#
+#   This program 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.
+#
+#   This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from ..models import StatementLine
+
+from datetime import datetime
+from io import StringIO
+import xml.etree.ElementTree as ET
+
+def import_ofx2(file):
+	raw_ofx = file.read().decode('utf-8')
+	
+	# Convert OFX header to XML and parse
+	xml_header = '<?xml version="1.0" encoding="UTF-8" standalone="no"?>'
+	raw_payload = raw_ofx[raw_ofx.index('?>')+2:]
+	xml_input = StringIO(xml_header + raw_payload.replace('&', '&amp;'))
+	try:
+		tree = ET.parse(xml_input)
+	except Exception as ex:
+		raise ex
+	root = tree.getroot()
+	
+	# Read transactions
+	lines = []  # Do first pass to catch "extra description lines"
+	for transaction in root.find('BANKMSGSRSV1').find('STMTTRNRS').find('STMTRS').find('BANKTRANLIST').findall('STMTTRN'):
+		date = transaction.find('DTPOSTED').text
+		date = date[0:4] + '-' + date[4:6] + '-' + date[6:8]
+		description = transaction.find('NAME').text
+		amount = transaction.find('TRNAMT').text
+		
+		if amount == '0':
+			lines[-1][3].append(description)
+			continue
+		
+		lines.append([date, description, amount, []])
+	
+	imported_statement_lines = []
+	
+	# Import
+	for date, description, amount, notes in lines:
+		imported_statement_lines.append(StatementLine(
+			dt=datetime.strptime(date, '%Y-%m-%d'),
+			description=description,
+			quantity=round(float(amount)*100),
+			commodity='$'
+		))
+	
+	return imported_statement_lines
diff --git a/drcr/statements/views.py b/drcr/statements/views.py
index 22e6448..994a82d 100644
--- a/drcr/statements/views.py
+++ b/drcr/statements/views.py
@@ -102,6 +102,9 @@ def statement_lines_import():
 	if request.form['format'] == 'ofx1':
 		from .importers.ofx1 import import_ofx1
 		statement_lines = import_ofx1(request.files['file'])
+	elif request.form['format'] == 'ofx2':
+		from .importers.ofx2 import import_ofx2
+		statement_lines = import_ofx2(request.files['file'])
 	else:
 		abort(400)
 	
diff --git a/drcr/templates/statements/import.html b/drcr/templates/statements/import.html
index d7e231a..6ae6a98 100644
--- a/drcr/templates/statements/import.html
+++ b/drcr/templates/statements/import.html
@@ -40,4 +40,24 @@
 			</div>
 		</div>
 	</form>
+	
+	<h2 class="h3 mt-4">OFX 2.x</h2>
+	
+	<form method="POST" enctype="multipart/form-data">
+		<input type="hidden" name="format" value="ofx2">
+		<div class="d-flex">
+			<div class="flex-grow-1 me-2">
+				<input class="form-control" name="source-account" placeholder="Source account">
+			</div>
+			<div class="flex-grow-1">
+				<input class="form-control" type="file" name="file" accept=".ofx">
+			</div>
+			<div>
+				<button type="submit" name="action" value="preview" class="btn btn-secondary ms-2">Preview</button>
+			</div>
+			<div>
+				<button type="submit" name="action" value="import" class="btn btn-primary ms-2">Import</button>
+			</div>
+		</div>
+	</form>
 {% endblock %}