clix/avatar.lua

local b64 = require"prosody.util.encodings".base64.encode;

local unb64 = require"prosody.util.encodings".base64.decode;

local st = require "prosody.util.stanza";

local sha1 = require "prosody.util.hashes".sha1;

return function(opts, arg)

	if opts.short_help then

		print("Manage PEP avatars");

		return;

	end

	local subcommands = {};

	function subcommands.fetch(conn)

		local waiting = {[true] = true};

		for _, userjid in ipairs(arg) do

			waiting[userjid] = true;

			local userpep = conn.pubsub:service(userjid);

			userpep:node("urn:xmpp:avatar:metadata"):items(true, function(result)

				local metadata_tag = result:find("{http://jabber.org/protocol/pubsub}pubsub/items/item/{urn:xmpp:avatar:metadata}metadata");

				if not metadata_tag or not metadata_tag:get_child("info") then

					if result.attr.type == "error" then

						conn:error("Got error from %s: %s:%s", userjid, result:get_error());

					else

						conn:error("%s has no avatar", userjid)

					end

					waiting[userjid] = nil;

					if next(waiting) == nil then

						conn:close();

					end

					return;

				end

				for info_tag in metadata_tag:childtags("info") do

					conn:info("Has avatar with type %s", info_tag.attr.type or "?")

					local filename = (opts.output or (userjid .. "_" .. info_tag.attr.id)) .. "." .. (info_tag.attr.type or "/dat"):match("/([^./+]+)");

					local output = assert(io.open(filename, "w"));

					waiting[info_tag.attr.id] = true;

					conn:debug("Writing to %s", filename);

					userpep:node("urn:xmpp:avatar:data"):item(info_tag.attr.id, function(dataresult)

						local data = unb64(dataresult:find("{http://jabber.org/protocol/pubsub}pubsub/items/item/{urn:xmpp:avatar:data}data#"));

						if data then

							assert(output:write(data));

							assert(output:close());

							conn:info("Avatar of %s written to %s", userjid, filename);

						else

							conn:error("Got no data for %s id %s", userjid, info_tag.attr.id);

						end

						waiting[info_tag.attr.id] = nil;

						if next(waiting) == nil then

							conn:close();

						end

					end);

					if not opts.all then

						break

					end

				end

				waiting[userjid] = nil;

			end)

		end

		waiting[true] = nil;

	end

	function subcommands.publish(conn)

		local waiting = {meta=true};

		local userpep = conn.pubsub:service(nil);

		local metadata_tag = st.stanza("metadata", { xmlns = "urn:xmpp:avatar:metadata" });

		local metadata_node = userpep:node("urn:xmpp:avatar:metadata");

		local data_node = userpep:node("urn:xmpp:avatar:data");

		local first_h = nil;

		local sha1pat = string.rep("%x", #sha1("",true));

		for _, file in ipairs(arg) do

			local h, width, height, typ = file:match("_("..sha1pat..")_(%d+)x(%d+)%.(%w+)$");

			if not h then h, typ = file:match("_("..sha1pat..")%.(%w+)$"); end

			if not h then typ = file:match("%.(%w+)$"); end

			local f = assert(io.open(file));

			local data = f:read("*a");

			f:close();

			local bytes = string.format("%d", #data);

			if not h then h = sha1(data, true); end

			if not first_h then first_h = h; end

			if typ == "jpg" then typ = "jpeg"; end

			local data_tag = st.stanza("data", { xmlns = "urn:xmpp:avatar:data" }):text(b64(data));

			waiting[h] = true;

			data_node:publish(h, nil, data_tag, function(ok)

				waiting[h] = nil;

				if next(waiting) == nil then

					conn:close();

				end

			end);

			metadata_tag:tag("info", {id = h; type = "image/" .. typ; bytes = bytes; width = width; height = height}):up();

		end

		metadata_node:publish(first_h, nil, metadata_tag, function(ok)

			waiting.meta = nil;

			if next(waiting) == nil then

				conn:close();

			end

		end);

	end

	if ((#arg == 0) or opts.help) then

		print("Subcommands:");

		print("  publish file_HASH_WxH.png");

		print("  fetch user@example.com");

		return 0;

	end

	if opts.output and opts.all then

		print("Can't download multiple avatars to a single file")

		return 1;

	end

	local subcommand = table.remove(arg, 1);

	local on_connect = subcommands[subcommand];

	if not on_connect then

		print("No such command: " .. subcommand);

		return 1;

	end

	return clix_connect(opts, on_connect, {"pubsub"})

end