#!/usr/bin/env python3

from pathlib import Path
import os
import re
import yaml

DIR = os.path.realpath(__file__ + '/../..')

def tagify(name):
  return 'kali-' + re.sub(r'[^a-z]+', '-', name.lower())

def xmlesc(txt):
  return txt.translate(str.maketrans({
    "&": "&amp;",
    "<": "&lt;",
    ">": "&gt;",
    '"': "&quot;",
    "'": "&apos;"
  }))

def set_category_to_desktop_file(tool_name, tag):
  desktop_file = f'{DIR}/desktop-files/kali-{tool_name}.desktop'
  if not os.path.isfile(desktop_file):
    desktop_file = f'{DIR}/desktop-files/{tool_name}.desktop'
  if not os.path.isfile(desktop_file):
    print(f"[!] Error: No desktop file found for {tag} : {tool_name}")
    return
  with open(desktop_file, 'r+') as file:
    content = re.sub(r'Categories=([^\n]*);?', f'Categories=\\1;{tag}', file.read())
    content = re.sub(r'=;', f'=', content)
    file.seek(0)
    file.write(content)

def generate_desktop_directory(category, tag):
  name_data = '\n'.join(f'{key}={value}' for key, value in category.items() if key.startswith('Name'))
  with open(f'{DIR}/desktop-directories/{tag}.directory', 'w') as file:
    file.write(f"""\
[Desktop Entry]
{name_data}
Type=Directory
Icon={tag}
""")

def generate_category_layout(category):
  category_name = category['Category']
  category_tag = tagify(category_name)
  generate_desktop_directory(category, category_tag)
  subcategories = '\n'.join(generate_subcategory_layout(sub) for sub in category.get('Include', []) if isinstance(sub, dict))
  return f"""\
  <Menu>
    <Name>{xmlesc(category_name)}</Name>
    <Directory>{category_tag}.directory</Directory>
    <Include>
      <And>
        <Category>{category_tag}</Category>
      </And>
    </Include>
{subcategories}
  </Menu>"""

def generate_subcategory_layout(subcategory):
  subcategory_name = subcategory['Category']
  subcategory_tag = tagify(subcategory_name)
  generate_desktop_directory(subcategory, subcategory_tag)
  return f"""\
    <Menu>
      <Name>{xmlesc(subcategory_name)}</Name>
      <Directory>{subcategory_tag}.directory</Directory>
      <Include>
        <And>
          <Category>{subcategory_tag}</Category>
        </And>
      </Include>
    </Menu>"""

def generate_menu_layout(data):
  menu_layout = '\n'.join(f"    <Menuname>{xmlesc(cat['Category'])}</Menuname>" for cat in data)
  menu_categories = '\n'.join(generate_category_layout(cat) for cat in data)
  return f"""\
<?xml version="1.0"?>
<!DOCTYPE Menu PUBLIC "-//freedesktop//DTD Menu 1.0//EN" "http://www.freedesktop.org/standards/menu-spec/1.0/menu.dtd">
<Menu>
  <Name>Applications</Name>
  <Layout>
{menu_layout}
    <Merge type="all"/>
  </Layout>

{menu_categories}

  <!-- Move default top-level menus one level deeper -->
  <Move>
    <Old>Accessories</Old>
    <New>Usual Applications/Accessories</New>
  </Move>
  <Move>
    <Old>Universal Access</Old>
    <New>Usual Applications/Universal Access</New>
  </Move>
  <Move>
    <Old>Electronics</Old>
    <New>Usual Applications/Electronics</New>
  </Move>
  <Move>
    <Old>Development</Old>
    <New>Usual Applications/Development</New>
  </Move>
  <Move>
    <Old>Education</Old>
    <New>Usual Applications/Education</New>
  </Move>
  <Move>
    <Old>Science</Old>
    <New>Usual Applications/Science</New>
  </Move>
  <Move>
    <Old>Games</Old>
    <New>Usual Applications/Games</New>
  </Move>
  <Move>
    <Old>Graphics</Old>
    <New>Usual Applications/Graphics</New>
  </Move>
  <Move>
    <Old>Internet</Old>
    <New>Usual Applications/Internet</New>
  </Move>
  <Move>
    <Old>Web Applications</Old>
    <New>Usual Applications/Web Applications</New>
  </Move>
  <Move>
    <Old>Multimedia</Old>
    <New>Usual Applications/Multimedia</New>
  </Move>
  <Move>
    <Old>Office</Old>
    <New>Usual Applications/Office</New>
  </Move>
  <Move>
    <Old>Settings</Old>
    <New>Usual Applications/Settings</New>
  </Move>
  <Move>
    <Old>Sundry</Old>
    <New>Usual Applications/Sundry</New>
  </Move>
  <Move>
    <Old>System</Old>
    <New>Usual Applications/System</New>
  </Move>
  <Move>
    <Old>System Settings</Old>
    <New>Usual Applications/System Settings</New>
  </Move>
  <Move>
    <Old>Utilities</Old>
    <New>Usual Applications/Utilities</New>
  </Move>
  <Move>
    <Old>Other</Old>
    <New>Usual Applications/Other</New>
  </Move>
  <Move>
    <Old>Hamradio</Old>
    <New>Usual Applications/Hamradio</New>
  </Move>
</Menu>"""



if __name__ == '__main__':
  desktop_directories_path = Path(f'{DIR}/desktop-directories/')
  for file in desktop_directories_path.glob('*'):
    file.unlink()

  with open(f'{DIR}/categories.yaml', 'r') as file:
    data = yaml.safe_load(file)
  with open(f'{DIR}/categories-purple.yaml', 'r') as file:
    data_purple = yaml.safe_load(file)

  for category in (data + data_purple):
    for child in category.get('Include', []):
      if type(child) is str: # Tool
        set_category_to_desktop_file(child, tagify(category['Category']))
      else: # Subcategory
        for tool in child.get('Include', []):
          set_category_to_desktop_file(tool, tagify(child['Category']))

  with open(f'{DIR}/menus/kali-applications.menu', 'w') as file:
    file.write(generate_menu_layout(data))
  with open(f'{DIR}/menus-purple/kali-applications.menu', 'w') as file:
    file.write(generate_menu_layout(data + data_purple))
