Police in Tehran said they had arrested a suspect and Iranian authorities condemned the incident but played down talk of any political motive for the attack.
The incident comes amid increased tensions between Azerbaijan and Iran, which is home to a large ethnic Azeri minority after Baku appointed its first-ever ambassador to Israel this month.
"The attacker broke through the guard post, killing the head of security with a Kalashnikov assault rifle," Azerbaijan's foreign ministry said.
A video posted on social media appeared to show a gunman running past a guard post and forcing his way into the embassy building, firing through a door.
Iranian Foreign Minister Hossein Amirabdollahian, quoted on a government site, said that based on evidence and initial observations the gunman's motive was "completely personal."
"Necessary security measures have been taken to continue normal activities at the embassy and diplomats of the Republic of Azerbaijan in Tehran," he said.
However, the Azeri foreign ministry said it planned shortly to evacuate its embassy staff from Tehran. It also summoned the Iranian ambassador in Baku to express its anger and demand justice.
In a strongly-worded statement, the ministry said an "anti-Azerbaijani campaign" in Iran had contributed to the attack, without elaborating, and accused Tehran of long ignoring its appeals to bolster security at the embassy.
"Unfortunately, the latest bloody terrorist act demonstrates the serious consequences of the failure to give the necessary attention to our constant appeals in this regard," it said.
In a separate statement on Twitter, Azerbaijan's President Ilham Aliyev demanded swift punishment for those involved in the "act of terrorism."
Aliyev has complained in the past about Iran's treatment of its Azerbaijani minority, saying for example that Azeris have no schools where they can study in their own language.
The superintendent of criminal affairs in Tehran, Judge Mohammad Shahriari, quoted the assailant as saying his wife went to the Azerbaijan embassy and had never returned home, according to Iran's Tasnim news agency.
When the assailant contacted the embassy, he never got an answer and believed his wife was still inside.
Iran's police chief said on television that the man's wife was from Azerbaijan.
Azerbaijan, a secular Muslim former Soviet republic which borders Iran, has friendly ties with the United States and Israel and has had difficult diplomatic relations with the Islamic Republic.
Israel has had an embassy in Baku since the early 1990s and has been a significant military backer of Azerbaijan in recent years. It has also provided diplomatic support for Baku in its standoff with Armenia over the Nagorno-Karabakh region.
Iran's Revolutionary Guards carried out major military drills along the country’s borders with Armenia and Azerbaijan amid fears of renewed fighting between the two South Caucasus states last year.
Reuters