1 files changed, 59 insertions, 3 deletions
diff --git a/Swiften/Avatars/AvatarFileStorage.cpp b/Swiften/Avatars/AvatarFileStorage.cpp
index 046ac16..44a6a7f 100644
--- a/Swiften/Avatars/AvatarFileStorage.cpp
+++ b/Swiften/Avatars/AvatarFileStorage.cpp
@@ -4,14 +4,39 @@
  * See Documentation/Licenses/GPLv3.txt for more information.
  */
 
-#include "Swiften/Avatars/AvatarFileStorage.h"
+#include <Swiften/Avatars/AvatarFileStorage.h>
 
 #include <iostream>
 #include <boost/filesystem/fstream.hpp>
+#include <boost/property_tree/ptree.hpp>
+#include <boost/property_tree/xml_parser.hpp>
+
+#include <Swiften/Base/foreach.h>
 
 namespace Swift {
 
-AvatarFileStorage::AvatarFileStorage(const boost::filesystem::path& path) : path_(path) {
+AvatarFileStorage::AvatarFileStorage(const boost::filesystem::path& avatarsDir, const boost::filesystem::path& avatarsFile) : avatarsDir(avatarsDir), avatarsFile(avatarsFile) {
+	if (boost::filesystem::exists(avatarsFile)) {
+		boost::property_tree::ptree tree;
+		try {
+			boost::property_tree::xml_parser::read_xml(avatarsFile.string(), tree);
+		}
+		catch (const boost::property_tree::xml_parser::xml_parser_error& e) {
+			std::cerr << "Error reading avatars file: " << e.filename() << ":" << e.line() << ": " << e.message() << std::endl;
+		}
+		foreach(const boost::property_tree::ptree::value_type &v, tree.get_child("avatars")) {
+			try {
+				JID jid(v.second.get<std::string>("jid"));
+				std::string hash(v.second.get<std::string>("hash"));
+				if (jid.isValid()) {
+					jidAvatars.insert(std::make_pair(jid, hash));
+				}
+			}
+			catch (const boost::property_tree::ptree_error& e) {
+				std::cerr << "Invalid avatar value: " << e.what() << std::endl;
+			}
+		}
+	}
 }
 
 bool AvatarFileStorage::hasAvatar(const std::string& hash) const { 
@@ -34,7 +59,7 @@ void AvatarFileStorage::addAvatar(const std::string& hash, const ByteArray& avat
 }
 
 boost::filesystem::path AvatarFileStorage::getAvatarPath(const std::string& hash) const {
-	return path_ / hash;
+	return avatarsDir / hash;
 }
 
 ByteArray AvatarFileStorage::getAvatar(const std::string& hash) const {
@@ -43,5 +68,36 @@ ByteArray AvatarFileStorage::getAvatar(const std::string& hash) const {
 	return data;
 }
 
+void AvatarFileStorage::setAvatarForJID(const JID& jid, const std::string& hash) {
+	std::pair<JIDAvatarMap::iterator, bool> r = jidAvatars.insert(std::make_pair(jid, hash));
+	if (r.second) {
+		saveJIDAvatars();
+	}
+	else if (r.first->second != hash) {
+		r.first->second = hash;
+		saveJIDAvatars();
+	}
+}
+
+std::string AvatarFileStorage::getAvatarForJID(const JID& jid) const {
+	JIDAvatarMap::const_iterator i = jidAvatars.find(jid);
+	return i == jidAvatars.end() ? "" : i->second;
+}
+
+void AvatarFileStorage::saveJIDAvatars() {
+	boost::property_tree::ptree tree;
+	for (JIDAvatarMap::const_iterator i = jidAvatars.begin(); i != jidAvatars.end(); ++i) {
+		boost::property_tree::ptree entry;
+		entry.put("jid", i->first.toString());
+		entry.put("hash", i->second);
+		tree.add_child("avatars.avatar", entry);
+	}
+	try {
+		boost::property_tree::xml_parser::write_xml(avatarsFile.string(), tree);
+	}
+	catch (const boost::property_tree::xml_parser::xml_parser_error& e) {
+		std::cerr << "Error writing avatars file: " << e.filename() << ": " << e.message() << std::endl;
+	}
+}
 
 }