<commit_before>""" Stores functions for rendering menus with Urwid.
"""

import urwid
from collection import Collection

class MainMenu():
    def __init__(self, collections):
        self.collections = collections
        self.selected = None
        self.widget = urwid.Padding([], left = 2, right = 2)
        self.top = urwid.Overlay(self.widget,
            urwid.SolidFill(u'\N{MEDIUM SHADE}'),
            align = 'center', width = ('relative', 60),
            valign = 'middle', height = ('relative', 60),
            min_width = 20, min_height = 9)

        self.main()

    def run(self):
        """ Enter the menu system """
        urwid.MainLoop(self.top, palette=[('reversed', 'standout', '')]).run()

    def exit(self, button):
        """ Exit the menu system """
        raise urwid.ExitMainLoop()

    # ----------------------------------------
    # Helper functions
    # ----------------------------------------

    def _update_widget(self, new_widget):
        """ Recreates the display widget. This should be called at the end of
        menu display functions.
        """
        self.widget.original_widget = urwid.Padding(new_widget)

    def _button(self, label, callback, data=None):
        """ An easier way to create a good looking button. """
        button = urwid.Button(label, callback, data)
        return urwid.AttrMap(button, None, focus_map='reversed')

    def _edit(self, label):
        """ An easier way to create an edit widget. """
        return urwid.Edit(label)

    # ----------------------------------------
    # Functions for rendering menu screens
    # ----------------------------------------

    def main(self, button=None, data=None):
        menu = GenericMenu('Please select an option')
        menu.add_widget(self._button('New Collection', self.new_collection))
        menu.add_widget(self._button('Study Collection', self.study_collection))
        menu.add_widget(self._button('Edit Collection', self.edit_collection))
        menu.add_widget(self._button('Exit', self.exit))

        self._update_widget(menu)

    def study_collection(self, button, collection=None):
        """ Render the menu for studying a collection """

        def study_prompt(button, slide):
            """ Render the menu for the first half of the study menu: reading
                the prompt.
            """
            menu = GenericMenu("Studying " + collection.name + ".")
            menu.add_widget(urwid.Text(slide.prompt))
            menu.add_widget(urwid.GridFlow(
                [
                    self._button("Flip", study_answer, slide)
                ],
                40, 1, 1, 'center'))

            self._update_widget(menu)

        def study_answer(button, slide):
            """ Render the menu for the first half of the study menu: reading
                the answer and determining how difficult the question was.
            """
            menu = GenericMenu("Studying " + collection.name + ".")
            menu.add_widget(urwid.Text(slide.answer))
            menu.add_widget(urwid.GridFlow(
                [
                    self._button("Easy", study_process_difficulty, (slide, 1))
                    ,self._button("Medium", study_process_difficulty, (slide, 2))
                    ,self._button("Hard", study_process_difficulty, (slide, 3))
                    ,self._button("Back", self.main)
                ],
                20, 1, 1, 'center'))

            self._update_widget(menu)

        def study_process_difficulty(button, data):
            """ Determine when the slide should be reviewed next based on the
                subjective difficulty of the question.

                The "data" parameter should be a data structure with a slide in
                the first position (data[0]) and an integer representing
                difficulty (generally from 1 to 3) in the second position
                (data[1]).
            """
            # TODO: Find a better way to pass the slide and difficulty in.
            slide = data[0]
            difficulty = data[1]

            # How long should the user wait to review this slide again?
            if difficulty == 1:
                # If it's easy, double the time
                multiplier = 2
            elif difficulty == 3:
                # If it's hard, halve the time
                multiplier = 0.5
            else:
                # If it's medium or anything else, leave the time alone
                multiplier = 1

            slide.set_new_wait_time(slide.previous_wait_time * multiplier)

            self.study_collection(None, collection)
        # ----------------------------------------

        if collection == None:
            self.choose_collection(self.study_collection)
        else:
            slide = collection.get_next_slide()
            if slide == None:
                slide = collection.get_random_slides(1)[0]

            study_prompt(None, slide)

    def choose_slide(self, button, collection, forward=None, back=None):
        """ Render the menu for choosing a slide from a collection. """
        if forward == None:
            forward = self.edit_slide
        if back == None:
            back = self.main

        menu = GenericMenu('Choose a Slide')
        menu.add_widget(self._button("Back", back))
        for slide in collection.slides:
            menu.add_widget(self._button(slide.prompt, forward, slide))
        menu.add_widget(self._button("Back", back))

        self._update_widget(menu)

    def choose_collection(self, forward, back=None):
        """ Render the menu for choosing a collection from the list. """
        if back == None:
            back = self.main

        menu = GenericMenu('Choose a Collection')
        menu.add_widget(self._button("Back", back))
        for collection in self.collections:
            menu.add_widget(self._button(collection.name, forward, collection))
        menu.add_widget(self._button("Back", back))

        self._update_widget(menu)

    def new_collection(self, button, data=None):
        """ Render the menu for creating a new Collection """

        def do_new_collection(button, menu):
            """ Create a new collection with the data from the
                'new_collection' menu.
            """
            fields = menu.get_fields()
            collection = Collection(fields['Name'])
            self.collections.append(collection)

            self.main()
        # ----------------------------------------

        menu = DialogueMenu("Fill out the fields below to create a new collection.")
        menu.add_field("Name")
        menu.add_widget(self._button("Create Collection", do_new_collection, menu))

        self._update_widget(menu)

    def edit_collection(self, button, collection=None):
        """ Render the menu for editing a collection """

        def do_edit_collection(button, menu):
            """ Modify the collection with the data entered in the menu. """
            fields = menu.get_fields()

            if fields['Name']:
                collection.name = fields['Name']

            self.main()
        # ----------------------------------------

        if collection == None:
            self.choose_collection(self.edit_collection)
        else:
            menu = DialogueMenu("Edit " + collection.name + ". Leave blank any field which you do not want to change.")
            menu.add_field("Name")
            menu.add_widget(self._button("Edit Slides", self.choose_slide, collection))
            menu.add_widget(urwid.Divider('-'))
            menu.add_widget(self._button("Save", do_edit_collection, menu))
            menu.add_widget(self._button("Cancel", self.main))

            self._update_widget(menu)

    def edit_slide(self, button, slide):
        """ Render the menu for editing a slide in a collection """
        def do_edit_collection(button, menu):
            """ Modify the slide with the data entered in the menu. """
            fields = menu.get_fields()

            if fields['Prompt']:
                slide.prompt = fields['Prompt']
            if fields['Answer']:
                slide.answer = fields['Answer']

            self.main()
        # ----------------------------------------

        menu = DialogueMenu("Edit \"" + slide.prompt + "\". Leave blank any field which you do not want to change.")
        menu.add_field("Prompt")
        menu.add_field("Answer")
        menu.add_widget(urwid.Divider('-'))
        menu.add_widget(self._button("Save", do_edit_collection, menu))
        menu.add_widget(self._button("Cancel", self.main))

        self._update_widget(menu)

class GenericMenu(urwid.WidgetWrap):
    """ A bare-bones menu meant to encapsulate the most basic functions of a
    menu using Urwid. May be used to build more complex menus.
    """
    def __init__(self, title):
        self.head = [urwid.Text(title), urwid.Divider()]
        self.body = []
        self._update_widget()

    def add_widget(self, widget):
        """ Add a widget to the menu """
        self.body.append(widget)
        self._update_widget()

    def _update_widget(self):
        """ Recreates the display widget. This should be called whenenver
        anything is added to head or body.
        """
        self.widget = urwid.ListBox(urwid.SimpleFocusListWalker(
            self.head + self.body))
        super().__init__(self.widget)

class NavigationMenu(urwid.WidgetWrap):
    def __init__(self, title, back):
        if back == None:
            self.head = [urwid.Text(title), urwid.Divider()]
        else:
            back_button = self._button("Back", back, self)
            self.head = [urwid.Text(title), urwid.Divider(), back_button]
        self.body = []
        self._update_widget()

    def add_option(self, label, callback, data=None):
        """ Add a navigation option to the menu """
        self.body.append(self._button(label, callback, data))
        self._update_widget()

    def _update_widget(self):
        # Recreates the display widget. This should be called whenenver
        # anything is added to head or body.
        self.widget = urwid.ListBox(urwid.SimpleFocusListWalker(
            self.head + self.body))
        super().__init__(self.widget)

    def _button(self, label, callback, data=None):
        # An easier way to create a good looking button
        button = urwid.Button(label, callback, data)
        return urwid.AttrMap(button, None, focus_map='reversed')

class DialogueMenu(GenericMenu):
    """ A menu class which adds support for handing text fields. """
    def __init__(self, title):
        super().__init__(title)

    def add_field(self, label):
        """ Add an Edit widget to the body """
        self.add_widget(urwid.Edit(label))

    def get_fields(self):
        """ Returns the contents of all the dialogue boxes in a dictionary
            with the form, {"label": "contents"}
        """
        fields = {}
        for widget in self.body:
            # Try to get the label and contents of widgets. If it fails, the
            # widget probably isn't an edit widget, so there's no need to
            # worry about the exception.
            try:
                label = widget.caption
                contents = widget.get_edit_text()
                fields[label] = contents
            except:
                pass

        return fields
<commit_msg>Add new_slide menu.<commit_after>""" Stores functions for rendering menus with Urwid.
"""

import urwid
from collection import Collection
from slide import Slide

class MainMenu():
    def __init__(self, collections):
        self.collections = collections
        self.selected = None
        self.widget = urwid.Padding([], left = 2, right = 2)
        self.top = urwid.Overlay(self.widget,
            urwid.SolidFill(u'\N{MEDIUM SHADE}'),
            align = 'center', width = ('relative', 60),
            valign = 'middle', height = ('relative', 60),
            min_width = 20, min_height = 9)

        self.main()

    def run(self):
        """ Enter the menu system """
        urwid.MainLoop(self.top, palette=[('reversed', 'standout', '')]).run()

    def exit(self, button):
        """ Exit the menu system """
        raise urwid.ExitMainLoop()

    # ----------------------------------------
    # Helper functions
    # ----------------------------------------

    def _update_widget(self, new_widget):
        """ Recreates the display widget. This should be called at the end of
        menu display functions.
        """
        self.widget.original_widget = urwid.Padding(new_widget)

    def _button(self, label, callback, data=None):
        """ An easier way to create a good looking button. """
        button = urwid.Button(label, callback, data)
        return urwid.AttrMap(button, None, focus_map='reversed')

    def _edit(self, label):
        """ An easier way to create an edit widget. """
        return urwid.Edit(label)

    # ----------------------------------------
    # Functions for rendering menu screens
    # ----------------------------------------

    def main(self, button=None, data=None):
        menu = GenericMenu('Please select an option')
        menu.add_widget(self._button('New Collection', self.new_collection))
        menu.add_widget(self._button('Study Collection', self.study_collection))
        menu.add_widget(self._button('Edit Collection', self.edit_collection))
        menu.add_widget(self._button('Exit', self.exit))

        self._update_widget(menu)

    def study_collection(self, button, collection=None):
        """ Render the menu for studying a collection """

        def study_prompt(button, slide):
            """ Render the menu for the first half of the study menu: reading
                the prompt.
            """
            menu = GenericMenu("Studying " + collection.name + ".")
            menu.add_widget(urwid.Text(slide.prompt))
            menu.add_widget(urwid.GridFlow(
                [
                    self._button("Flip", study_answer, slide)
                ],
                40, 1, 1, 'center'))

            self._update_widget(menu)

        def study_answer(button, slide):
            """ Render the menu for the first half of the study menu: reading
                the answer and determining how difficult the question was.
            """
            menu = GenericMenu("Studying " + collection.name + ".")
            menu.add_widget(urwid.Text(slide.answer))
            menu.add_widget(urwid.GridFlow(
                [
                    self._button("Easy", study_process_difficulty, (slide, 1))
                    ,self._button("Medium", study_process_difficulty, (slide, 2))
                    ,self._button("Hard", study_process_difficulty, (slide, 3))
                    ,self._button("Back", self.main)
                ],
                20, 1, 1, 'center'))

            self._update_widget(menu)

        def study_process_difficulty(button, data):
            """ Determine when the slide should be reviewed next based on the
                subjective difficulty of the question.

                The "data" parameter should be a data structure with a slide in
                the first position (data[0]) and an integer representing
                difficulty (generally from 1 to 3) in the second position
                (data[1]).
            """
            # TODO: Find a better way to pass the slide and difficulty in.
            slide = data[0]
            difficulty = data[1]

            # How long should the user wait to review this slide again?
            if difficulty == 1:
                # If it's easy, double the time
                multiplier = 2
            elif difficulty == 3:
                # If it's hard, halve the time
                multiplier = 0.5
            else:
                # If it's medium or anything else, leave the time alone
                multiplier = 1

            slide.set_new_wait_time(slide.previous_wait_time * multiplier)

            self.study_collection(None, collection)
        # ----------------------------------------

        if collection == None:
            self.choose_collection(self.study_collection)
        else:
            slide = collection.get_next_slide()
            if slide == None:
                slide = collection.get_random_slides(1)[0]

            study_prompt(None, slide)

    def choose_slide(self, button, collection, forward=None, back=None):
        """ Render the menu for choosing a slide from a collection. """
        if forward == None:
            forward = self.edit_slide
        if back == None:
            back = self.main

        menu = GenericMenu('Choose a Slide')
        menu.add_widget(self._button("Back", back))
        for slide in collection.slides:
            menu.add_widget(self._button(slide.prompt, forward, slide))
        menu.add_widget(self._button("Back", back))

        self._update_widget(menu)

    def choose_collection(self, forward, back=None):
        """ Render the menu for choosing a collection from the list. """
        if back == None:
            back = self.main

        menu = GenericMenu('Choose a Collection')
        menu.add_widget(self._button("Back", back))
        for collection in self.collections:
            menu.add_widget(self._button(collection.name, forward, collection))
        menu.add_widget(self._button("Back", back))

        self._update_widget(menu)

    def new_collection(self, button, data=None):
        """ Render the menu for creating a new Collection """

        def do_new_collection(button, menu):
            """ Create a new collection with the data from the
                'new_collection' menu.
            """
            fields = menu.get_fields()
            collection = Collection(fields['Name'])
            self.collections.append(collection)

            self.main()
        # ----------------------------------------

        menu = DialogueMenu("Fill out the fields below to create a new collection.")
        menu.add_field("Name")
        menu.add_widget(self._button("Create Collection", do_new_collection, menu))

        self._update_widget(menu)

    def edit_collection(self, button, collection=None):
        """ Render the menu for editing a collection """

        def do_edit_collection(button, menu):
            """ Modify the collection with the data entered in the menu. """
            fields = menu.get_fields()

            if fields['Name']:
                collection.name = fields['Name']

            self.main()
        # ----------------------------------------

        if collection == None:
            self.choose_collection(self.edit_collection)
        else:
            menu = DialogueMenu("Edit " + collection.name + ". Leave blank any field which you do not want to change.")
            menu.add_field("Name")
            menu.add_widget(self._button("Create New Slide", self.new_slide, collection))
            menu.add_widget(self._button("Edit Slides", self.choose_slide, collection))
            menu.add_widget(urwid.Divider('-'))
            menu.add_widget(self._button("Save", do_edit_collection, menu))
            menu.add_widget(self._button("Cancel", self.main))

            self._update_widget(menu)

    def new_slide(self, button, collection):
        """ Render the menu for creating a new Slide """

        def do_new_slide(button, menu):
            """ Create a new slide with the data from the 'new_slide' menu. """
            fields = menu.get_fields()
            # TODO: Add a check to be sure the slide is valid.
            slide = Slide(fields['Prompt'], fields['Answer'])
            collection.add_slide(slide)

            self.edit_collection(None, collection)
        # ----------------------------------------

        # TODO: Make slide creation easier and faster.
        menu = DialogueMenu("Fill out the fields below to create a new slide.")
        menu.add_field("Prompt")
        menu.add_field("Answer")
        menu.add_widget(self._button("Create Slide", do_new_slide, menu))

        self._update_widget(menu)

    def edit_slide(self, button, slide):
        """ Render the menu for editing a slide in a collection """
        def do_edit_collection(button, menu):
            """ Modify the slide with the data entered in the menu. """
            fields = menu.get_fields()

            if fields['Prompt']:
                slide.prompt = fields['Prompt']
            if fields['Answer']:
                slide.answer = fields['Answer']

            self.main()
        # ----------------------------------------

        menu = DialogueMenu("Edit \"" + slide.prompt + "\". Leave blank any field which you do not want to change.")
        menu.add_field("Prompt")
        menu.add_field("Answer")
        menu.add_widget(urwid.Divider('-'))
        menu.add_widget(self._button("Save", do_edit_collection, menu))
        menu.add_widget(self._button("Cancel", self.main))

        self._update_widget(menu)

class GenericMenu(urwid.WidgetWrap):
    """ A bare-bones menu meant to encapsulate the most basic functions of a
    menu using Urwid. May be used to build more complex menus.
    """
    def __init__(self, title):
        self.head = [urwid.Text(title), urwid.Divider()]
        self.body = []
        self._update_widget()

    def add_widget(self, widget):
        """ Add a widget to the menu """
        self.body.append(widget)
        self._update_widget()

    def _update_widget(self):
        """ Recreates the display widget. This should be called whenenver
        anything is added to head or body.
        """
        self.widget = urwid.ListBox(urwid.SimpleFocusListWalker(
            self.head + self.body))
        super().__init__(self.widget)

class NavigationMenu(urwid.WidgetWrap):
    def __init__(self, title, back):
        if back == None:
            self.head = [urwid.Text(title), urwid.Divider()]
        else:
            back_button = self._button("Back", back, self)
            self.head = [urwid.Text(title), urwid.Divider(), back_button]
        self.body = []
        self._update_widget()

    def add_option(self, label, callback, data=None):
        """ Add a navigation option to the menu """
        self.body.append(self._button(label, callback, data))
        self._update_widget()

    def _update_widget(self):
        # Recreates the display widget. This should be called whenenver
        # anything is added to head or body.
        self.widget = urwid.ListBox(urwid.SimpleFocusListWalker(
            self.head + self.body))
        super().__init__(self.widget)

    def _button(self, label, callback, data=None):
        # An easier way to create a good looking button
        button = urwid.Button(label, callback, data)
        return urwid.AttrMap(button, None, focus_map='reversed')

class DialogueMenu(GenericMenu):
    """ A menu class which adds support for handing text fields. """
    def __init__(self, title):
        super().__init__(title)

    def add_field(self, label):
        """ Add an Edit widget to the body """
        self.add_widget(urwid.Edit(label))

    def get_fields(self):
        """ Returns the contents of all the dialogue boxes in a dictionary
            with the form, {"label": "contents"}
        """
        fields = {}
        for widget in self.body:
            # Try to get the label and contents of widgets. If it fails, the
            # widget probably isn't an edit widget, so there's no need to
            # worry about the exception.
            try:
                label = widget.caption
                contents = widget.get_edit_text()
                fields[label] = contents
            except:
                pass

        return fields
<|endoftext|>


<commit_before>COLS=20
ROWS=6

HUD_name="Crash the HUD"

import curses

def main(scr):
	scr.clear()
	scr.addstr(0, int((COLS-len(HUD_name)) / 2), HUD_name)
	scr.hline(1, 0, curses.ACS_HLINE, scr.getmaxyx()[1])
	scr.refresh()
	scr.getkey()

scr=curses.initscr()
curses.noecho()
curses.cbreak()
curses.curs_set(0)
scr.keypad(True)

main(scr)

curses.nocbreak()
scr.keypad(False)
curses.echo()
curses.endwin()
<commit_msg>Keep the program running<commit_after>COLS=20
ROWS=6

HUD_name="Crash the HUD"

import curses

def main(scr):
	menu_items=("Telemetry", "Navigation", "Temperature sync", "Hot tub")
	menu_display_max=4
	keypress=None

	while keypress!=curses.KEY_EXIT:
		scr.clear()
		scr.addstr(0, int((COLS - len(HUD_name)) / 2), HUD_name)
		scr.hline(1, 0, curses.ACS_HLINE, scr.getmaxyx()[1])
		for i in range(2, 6):
			menu_i=i - 2 + (menu_display_max - 4)
			scr.addstr(i, 2, " " * int((COLS - 4 - len(menu_items[menu_i])) / 2) + menu_items[menu_i])
		scr.refresh()
		keypress=scr.getkey()

scr=curses.initscr()
curses.noecho()
curses.cbreak()
curses.curs_set(0)
scr.keypad(True)

main(scr)

curses.nocbreak()
scr.keypad(False)
curses.echo()
curses.endwin()
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""This program inputs a MemberHub directory dump, and analyzes it.
"""

import directory_tools
import roster_tools
import hub_map_tools
import import_file_tools

def ContainsName(name_field, first_name, last_name):

    found_name = name_field.find(first_name) >= 0 and \
                 name_field.find(last_name) >= 0

    return found_name

def IsSameName(from_directory, from_roster):
    is_same = False

    if from_roster["family_relation"][:5].lower() == "child":

        is_same = from_directory["last_name"]  == from_roster["last_name"] and \
                  from_directory["first_name"] == from_roster["first_name"]

    else:

        is_same = ContainsName(from_roster["name_field"],
                               from_directory["first_name"],
                               from_directory["last_name"])

    return is_same

def FindMissingEmail(direct_d):
    """FindMissingEmail finds all the adults in the directory that do not have
an email address associated with their entry.
"""
    adult_count = no_email_count = 0
    for entry in direct_d:
        if direct_d[entry]["family_relation"][:5].lower() == "adult":
            adult_count += 1
            if direct_d[entry]["email"] == "":
                no_email_count += 1

    print "Found %d out of %d adults with missing email addresses" % \
          (no_email_count, adult_count)

def FindOrphans(direct_d):
    """FindOrphans finds all the children in the directory that do not have an
adult associated with them.
"""
    child_count = orphan_count = 0

    for entry in direct_d:
        if direct_d[entry]["family_relation"][:5].lower() == "child":
            child_count += 1
            if direct_d[entry]["parents"] == "":
                orphan_count += 1

    print "Found %d children without adults in the family out of %d children" % \
          (orphan_count, child_count)

def FindChildless(direct_d):
    print "This functionality has not been added yet"

def FindHubless(arg_list):
    direct_d = arg_list[0]
    map_d    = arg_list[1]

    child_count = adult_count = 0

    for entry in direct_d:
        if not hub_map_tools.IsAnyHubClassroomHub(map_d, direct_d[entry]["hubs"]):
            print "%s %s is not in a classroom hub (%s)" % (direct_d[entry]["first_name"],
                                                            direct_d[entry]["last_name"],
                                                            direct_d[entry]["hubs"])
            if direct_d[entry]["family_relation"][:5].lower() == "child":
                child_count += 1
            else:
                adult_count += 1

    print "Found %d children who are not in at least one hub" % child_count
    print "Found %d adults who are not in at least one hub" % adult_count

def MakeImportForHubless(arg_list):
    direct_d = arg_list[0]
    roster_d = arg_list[1]
    map_d    = arg_list[2]
    update_d = {}

    processed_count = hubed_count = hubless_count = 0

    for entry in direct_d:
        if direct_d[entry]["hubs"] == "":
            processed_count += 1
            for person in roster_d:
                if IsSameName(direct_d[entry], roster_d[person]):

                    temp = direct_d[entry]
                    hub_names = "|"
                    for teacher in (roster_d[person]["teacher"].split('|')):
                        hub_names += map_d[teacher] + "|"

                    temp.update({"hubs":hub_names})
                    update_d.update({processed_count:temp})
                    hubed_count += 1
                    break
            else:
                print "Did not find match for %s %s." % \
                          (direct_d[entry]["first_name"], direct_d[entry]["last_name"])
                hubless_count += 1

    print "Processed %d entries that were not in at least one hub." % processed_count
    print "Was able to find hubs for %d entries." % hubed_count
    print "Was not able to find hubs for %d entries." % hubless_count

    import_file_tools.CreateHublessImportFile(update_d)

def FindEntriless(arg_list):
    direct_d = arg_list[0]
    roster_d = arg_list[1]

    entriless_count = 0

    for person in roster_d:

        for entry in direct_d:
            if IsSameName(direct_d[entry], roster_d[person]):
                print "Found %s %s in directory." % \
                      (roster_d[person]["first_name"], roster_d[person]["last_name"])
                break

        else:
            print "Did not find match for %s %s." % \
                  (roster_d[person]["first_name"], roster_d[person]["last_name"])
            entriless_count += 1

    print "Found %d people on the roster who were not in the directory" % entriless_count

def MakeImportNotInDirectory(arg_list):
    direct_d = arg_list[0]
    roster_d = arg_list[1]
    map_d    = arg_list[2]
    update_d = {}

    entriless_count = 0

    for person in roster_d:
        for entry in direct_d:
            if IsSameName(direct_d[entry], roster_d[person]):
                break

        else:
            hub_names = "|"
            for teacher in (roster_d[person]["teacher"].split('|')):
                hub_names += map_d[teacher] + "|"
            update_d.update({entriless_count:{"last_name":roster_d[person]["last_name"],
                                              "first_name":roster_d[person]["first_name"],
                                              "family_relation":roster_d[person]["family_relation"],
                                              "person_id":" ",
                                              "hubs":hub_names}})
            entriless_count += 1

    print "Found %d people on the roster who were not in the directory" % entriless_count

    import_file_tools.CreateNewMemberImport(update_d)

def MakePrompt(choices):
    choice_list = sorted(choices)
    guts = '\n'.join(['(%s)%s' % (choice[0], choice[1:])
                      for choice in choice_list])
    return 'Choose:\n' + guts + '\nOr press <enter> to quit '

def RunMenu(direct_d, roster_d, map_d):
    """Runs the user interface for dictionary manipulation."""
    # The choices dictionary has function names for values.
    choices = {'1 - Find Missing Email':{'Function':FindMissingEmail,'Arg':direct_d},
               '2 - Find Orphans':{'Function':FindOrphans,'Arg':direct_d},
               '3 - Find Childless':{'Function':FindChildless,'Arg':direct_d},
               '4 - Find Not In Hub':{'Function':FindHubless,'Arg':[direct_d,map_d]},
               '5 - Make Import File for Not In Hub':{'Function':MakeImportForHubless,
                                                      'Arg':[direct_d,roster_d,map_d]},
               '6 - Find Not in Directory':{'Function':FindEntriless,
                                            'Arg':[direct_d,roster_d]},
               '7 - Make Import File for Not In Directory':
                                           {'Function':MakeImportNotInDirectory,
                                            'Arg':[direct_d,roster_d,map_d]}}

    prompt = MakePrompt(choices)

    while True:
        raw_choice = raw_input(prompt)
        if not raw_choice:
            break
        given_choice = raw_choice[0].lower()
        for maybe_choice in choices:
            if maybe_choice[0] == given_choice:
                # The appropriate function is called
                # using the dictionary value for the name
                # of the function.
                choices[maybe_choice]['Function'](choices[maybe_choice]['Arg'])
                break
        else:
            print '%s is not an acceptible choice.' % raw_choice

def main():
    direct_d = directory_tools.ReadDirectory()
    roster_d = roster_tools.ReadRoster()
    map_d    = hub_map_tools.ReadMap()
    RunMenu(direct_d, roster_d, map_d)

if __name__ == '__main__':
    main()
<commit_msg>extended new hubless logic to import file creator<commit_after>#!/usr/bin/env python
"""This program inputs a MemberHub directory dump, and analyzes it.
"""

import directory_tools
import roster_tools
import hub_map_tools
import import_file_tools

def ContainsName(name_field, first_name, last_name):

    found_name = name_field.find(first_name) >= 0 and \
                 name_field.find(last_name) >= 0

    return found_name

def IsSameName(from_directory, from_roster):
    is_same = False

    if from_roster["family_relation"][:5].lower() == "child":

        is_same = from_directory["last_name"]  == from_roster["last_name"] and \
                  from_directory["first_name"] == from_roster["first_name"]

    else:

        is_same = ContainsName(from_roster["name_field"],
                               from_directory["first_name"],
                               from_directory["last_name"])

    return is_same

def FindMissingEmail(direct_d):
    """FindMissingEmail finds all the adults in the directory that do not have
an email address associated with their entry.
"""
    adult_count = no_email_count = 0
    for entry in direct_d:
        if direct_d[entry]["family_relation"][:5].lower() == "adult":
            adult_count += 1
            if direct_d[entry]["email"] == "":
                no_email_count += 1

    print "Found %d out of %d adults with missing email addresses" % \
          (no_email_count, adult_count)

def FindOrphans(direct_d):
    """FindOrphans finds all the children in the directory that do not have an
adult associated with them.
"""
    child_count = orphan_count = 0

    for entry in direct_d:
        if direct_d[entry]["family_relation"][:5].lower() == "child":
            child_count += 1
            if direct_d[entry]["parents"] == "":
                orphan_count += 1

    print "Found %d children without adults in the family out of %d children" % \
          (orphan_count, child_count)

def FindChildless(direct_d):
    print "This functionality has not been added yet"

def FindHubless(arg_list):
    direct_d = arg_list[0]
    map_d    = arg_list[1]

    child_count = adult_count = 0

    for entry in direct_d:
        if not hub_map_tools.IsAnyHubClassroomHub(map_d, direct_d[entry]["hubs"]):
            print "%s %s is not in a classroom hub (%s)" % (direct_d[entry]["first_name"],
                                                            direct_d[entry]["last_name"],
                                                            direct_d[entry]["hubs"])
            if direct_d[entry]["family_relation"][:5].lower() == "child":
                child_count += 1
            else:
                adult_count += 1

    print "Found %d children who are not in at least one hub" % child_count
    print "Found %d adults who are not in at least one hub" % adult_count

def MakeImportForHubless(arg_list):
    direct_d = arg_list[0]
    roster_d = arg_list[1]
    map_d    = arg_list[2]
    update_d = {}

    processed_count = hubed_count = hubless_count = 0

    for entry in direct_d:
        if not hub_map_tools.IsAnyHubClassroomHub(map_d, direct_d[entry]["hubs"]):
            print "%s %s is not in a classroom hub (%s)" % (direct_d[entry]["first_name"],
                                                            direct_d[entry]["last_name"],
                                                            direct_d[entry]["hubs"])
            processed_count += 1
            for person in roster_d:
                if IsSameName(direct_d[entry], roster_d[person]):

                    temp = direct_d[entry]
                    hub_names = temp["hubs"] + "|"
                    for teacher in (roster_d[person]["teacher"].split('|')):
                        hub_names += map_d[teacher] + "|"

                    temp.update({"hubs":hub_names})
                    update_d.update({processed_count:temp})
                    hubed_count += 1
                    break
            else:
                print "Did not find match for %s %s in the roster." % \
                          (direct_d[entry]["first_name"], direct_d[entry]["last_name"])
                hubless_count += 1

    print "Processed %d entries that were not in any classroom hub." % processed_count
    print "Was able to find classroom hubs for %d entries." % hubed_count
    print "Was not able to find classroom hubs for %d entries." % hubless_count

    import_file_tools.CreateHublessImportFile(update_d)

def FindEntriless(arg_list):
    direct_d = arg_list[0]
    roster_d = arg_list[1]

    entriless_count = 0

    for person in roster_d:

        for entry in direct_d:
            if IsSameName(direct_d[entry], roster_d[person]):
                print "Found %s %s in directory." % \
                      (roster_d[person]["first_name"], roster_d[person]["last_name"])
                break

        else:
            print "Did not find match for %s %s." % \
                  (roster_d[person]["first_name"], roster_d[person]["last_name"])
            entriless_count += 1

    print "Found %d people on the roster who were not in the directory" % entriless_count

def MakeImportNotInDirectory(arg_list):
    direct_d = arg_list[0]
    roster_d = arg_list[1]
    map_d    = arg_list[2]
    update_d = {}

    entriless_count = 0

    for person in roster_d:
        for entry in direct_d:
            if IsSameName(direct_d[entry], roster_d[person]):
                break

        else:
            hub_names = "|"
            for teacher in (roster_d[person]["teacher"].split('|')):
                hub_names += map_d[teacher] + "|"
            update_d.update({entriless_count:{"last_name":roster_d[person]["last_name"],
                                              "first_name":roster_d[person]["first_name"],
                                              "family_relation":roster_d[person]["family_relation"],
                                              "person_id":" ",
                                              "hubs":hub_names}})
            entriless_count += 1

    print "Found %d people on the roster who were not in the directory" % entriless_count

    import_file_tools.CreateNewMemberImport(update_d)

def MakePrompt(choices):
    choice_list = sorted(choices)
    guts = '\n'.join(['(%s)%s' % (choice[0], choice[1:])
                      for choice in choice_list])
    return 'Choose:\n' + guts + '\nOr press <enter> to quit '

def RunMenu(direct_d, roster_d, map_d):
    """Runs the user interface for dictionary manipulation."""
    # The choices dictionary has function names for values.
    choices = {'1 - Find Missing Email':{'Function':FindMissingEmail,'Arg':direct_d},
               '2 - Find Orphans':{'Function':FindOrphans,'Arg':direct_d},
               '3 - Find Childless':{'Function':FindChildless,'Arg':direct_d},
               '4 - Find Not In Hub':{'Function':FindHubless,'Arg':[direct_d,map_d]},
               '5 - Make Import File for Not In Hub':{'Function':MakeImportForHubless,
                                                      'Arg':[direct_d,roster_d,map_d]},
               '6 - Find Not in Directory':{'Function':FindEntriless,
                                            'Arg':[direct_d,roster_d]},
               '7 - Make Import File for Not In Directory':
                                           {'Function':MakeImportNotInDirectory,
                                            'Arg':[direct_d,roster_d,map_d]}}

    prompt = MakePrompt(choices)

    while True:
        raw_choice = raw_input(prompt)
        if not raw_choice:
            break
        given_choice = raw_choice[0].lower()
        for maybe_choice in choices:
            if maybe_choice[0] == given_choice:
                # The appropriate function is called
                # using the dictionary value for the name
                # of the function.
                choices[maybe_choice]['Function'](choices[maybe_choice]['Arg'])
                break
        else:
            print '%s is not an acceptible choice.' % raw_choice

def main():
    direct_d = directory_tools.ReadDirectory()
    roster_d = roster_tools.ReadRoster()
    map_d    = hub_map_tools.ReadMap()
    RunMenu(direct_d, roster_d, map_d)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before><commit_msg>Draft: Revised the 5 Draft.scale unit tests. No longer using scale_vector_from_center.<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python
# coding: utf-8

"""
mpyq is a Python library for reading MPQ (MoPaQ) archives.
"""

import cStringIO
import os
import struct
import sys
import zlib
from collections import namedtuple
from optparse import OptionParser

__author__ = "Aku Kotkavuo"
__version__ = "0.1"

MPQ_FILE_IMPLODE        = 0x00000100
MPQ_FILE_COMPRESS       = 0x00000200
MPQ_FILE_ENCRYPTED      = 0x00010000
MPQ_FILE_FIX_KEY        = 0x00020000
MPQ_FILE_SINGLE_UNIT    = 0x01000000
MPQ_FILE_DELETE_MARKER  = 0x02000000
MPQ_FILE_SECTOR_CRC     = 0x04000000
MPQ_FILE_EXISTS         = 0x80000000

MPQFileHeader = namedtuple('MPQFileHeader',
    '''
    magic
    header_size
    arhive_size
    format_version
    sector_size_shift
    hash_table_offset
    block_table_offset
    hash_table_entries
    block_table_entries
    '''
)
MPQFileHeader.struct_format = '<4s2I2H4I'

MPQFileHeaderExt = namedtuple('MPQFileHeaderExt',
    '''
    extended_block_table_offset
    hash_table_offset_high
    block_table_offset_high
    '''
)
MPQFileHeaderExt.struct_format = 'q2h'

MPQUserDataHeader = namedtuple('MPQUserDataHeader',
    '''
    magic
    user_data_size
    mpq_header_offset
    user_data_header_size
    '''
)
MPQUserDataHeader.struct_format = '<4s3I'

MPQHashTableEntry = namedtuple('MPQHashTableEntry',
    '''
    hash_a
    hash_b
    locale
    platform
    block_table_index
    '''
)
MPQHashTableEntry.struct_format = '2I2HI'

MPQBlockTableEntry = namedtuple('MPQBlockTableEntry',
    '''
    offset
    archived_size
    size
    flags
    '''
)
MPQBlockTableEntry.struct_format = '4I'

SC2ReplayHeader = namedtuple('SC2ReplayHeader',
    '''
    identifier
    release_flag
    major_version
    minor_version
    maintenance_version
    build_number
    duration
    '''
)
SC2ReplayHeader.struct_format = '>x19s2x4bi6xhx'

class MPQArchive(object):

    def __init__(self, filename):
        self.file = open(filename, 'rb')
        self.header = self.read_header()
        self.hash_table =  self.read_table('hash')
        self.block_table = self.read_table('block')
        self.files = self.read_file('(listfile)').splitlines()

    def read_header(self):
        """Read the header of a MPQ archive."""

        def read_mpq_header(offset=None):
            if offset:
                self.file.seek(offset)
            data = self.file.read(32)
            header = MPQFileHeader._make(
                struct.unpack(MPQFileHeader.struct_format, data))
            header = header._asdict()
            if header['format_version'] == 1:
                data = self.file.read(12)
                extended_header = MPQFileHeaderExt._make(
                    struct.unpack(MPQFileHeaderExt.struct_format, data))
                header.update(extended_header._asdict())
            return header

        def read_mpq_user_data_header():
            data = self.file.read(16)
            header = MPQUserDataHeader._make(
                struct.unpack(MPQUserDataHeader.struct_format, data))
            header = header._asdict()
            header['content'] = self.file.read(header['user_data_header_size'])
            if header['content'].startswith('\x15StarCraft II replay'):
                header['starcraft2_replay_header'] = SC2ReplayHeader._make(
                    struct.unpack(SC2ReplayHeader.struct_format,
                                  header['content']))
            return header

        magic = self.file.read(4)
        self.file.seek(0)

        if magic == 'MPQ\x1a':
            header = read_mpq_header()
            header['offset'] = 0
        elif magic == 'MPQ\x1b':
            user_data_header = read_mpq_user_data_header()
            header = read_mpq_header(user_data_header['mpq_header_offset'])
            header['offset'] = user_data_header['mpq_header_offset']
            header['user_data_header'] = user_data_header

        return header

    def read_table(self, table_type):
        """Read either hash or block table of a MPQ archive."""

        table_offset = self.header['%s_table_offset' % table_type]
        table_size = self.header['%s_table_entries' % table_type] * 16
        key = self._hash('(%s table)' % table_type, 'TABLE')

        self.file.seek(table_offset + self.header['offset'])
        table = self.file.read(table_size)
        table = self._decrypt(table, key)
        return table

    def read_file(self, filename):
        """Read a file from the MPQ archive."""

        def decompress(data):
            """Read the compression type and decompress file data."""
            compression_type = ord(data[0])
            if compression_type == 0:
                return data
            if compression_type == 2:
                return zlib.decompress(data[1:], 15)

        # Iterate over all hash table entries, looking for the correct one.
        hash_a = self._hash(filename, 'HASH_A')
        hash_b = self._hash(filename, 'HASH_B')
        for i in range(self.header['hash_table_entries']):
            data = self.hash_table[i*16:i*16+16]
            hash_table_entry = MPQHashTableEntry._make(
                struct.unpack(MPQHashTableEntry.struct_format, data))
            if (hash_table_entry.hash_a == hash_a and
                hash_table_entry.hash_b == hash_b):
                # Correct hash table entry has been found.
                break

        # Get the corresponding block table entry.
        pos = hash_table_entry.block_table_index * 16
        data = self.block_table[pos:pos+16]
        block_table_entry = MPQBlockTableEntry._make(
            struct.unpack(MPQBlockTableEntry.struct_format, data))

        # Read the block.
        if block_table_entry.flags & MPQ_FILE_EXISTS:
            offset = block_table_entry.offset + self.header['offset']
            self.file.seek(offset)
            file_data = self.file.read(block_table_entry.archived_size)

            if not block_table_entry.flags & MPQ_FILE_SINGLE_UNIT:
                # File consist of many sectors. They all need to be
                # decompressed separately and united.
                sector_size = 512 << self.header['sector_size_shift']
                sectors = block_table_entry.size / sector_size + 1
                if block_table_entry.flags & MPQ_FILE_SECTOR_CRC:
                    crc = True
                    sectors += 1
                positions = struct.unpack('<%dI' % (sectors + 1),
                                          file_data[:4*(sectors+1)])
                result = cStringIO.StringIO()
                for i in range(len(positions) - (2 if crc else 1)):
                    sector = file_data[positions[i]:positions[i+1]]
                    sector = decompress(sector)
                    result.write(sector)
                file_data = result.getvalue()
            else:
                # Single unit files only need to be decompressed.
                if block_table_entry.flags & MPQ_FILE_COMPRESS:
                    file_data = decompress(file_data)

            return file_data

    def extract(self):
        """Extract all the files inside the MPQ archive in memory."""
        return dict((f, self.read_file(f)) for f in self.files)

    def extract_to_disk(self):
        """Extract all files and write them to disk."""
        archive_name, extension = os.path.splitext(self.file.name)
        if not os.path.isdir(os.path.join(os.getcwd(), archive_name)):
            os.mkdir(archive_name)
        os.chdir(archive_name)
        for filename, data in self.extract().items():
            f = open(filename, 'wb')
            f.write(data)
            f.close()

    def _hash(self, string, hash_type):
        """Hash a string using MPQ's hash function."""
        hash_types = {
            'TABLE_OFFSET': 0,
            'HASH_A': 1,
            'HASH_B': 2,
            'TABLE': 3
        }
        seed1 = 0x7FED7FED
        seed2 = 0xEEEEEEEE

        for ch in string:
            ch = ord(ch.upper())
            value = self.encryption_table[(hash_types[hash_type] << 8) + ch]
            seed1 = (value ^ (seed1 + seed2)) & 0xFFFFFFFF
            seed2 = ch + seed1 + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

        return seed1

    def _decrypt(self, data, key):
        """Decrypt hash or block table or a sector."""
        seed1 = key
        seed2 = 0xEEEEEEEE
        result = cStringIO.StringIO()

        for i in range(len(data) // 4):
            seed2 += self.encryption_table[0x400 + (seed1 & 0xFF)]
            seed2 &= 0xFFFFFFFF
            value = struct.unpack("<I", data[i*4:i*4+4])[0]
            value = (value ^ (seed1 + seed2)) & 0xFFFFFFFF

            seed1 = ((~seed1 << 0x15) + 0x11111111) | (seed1 >> 0x0B)
            seed1 &= 0xFFFFFFFF
            seed2 = value + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

            result.write(struct.pack("<I", value))

        return result.getvalue()

    def _prepare_encryption_table():
        """Prepare encryption table for MPQ hash function."""
        seed = 0x00100001
        crypt_table = {}

        for i in range(256):
            index = i
            for j in range(5):
                seed = (seed * 125 + 3) % 0x2AAAAB
                temp1 = (seed & 0xFFFF) << 0x10;

                seed = (seed * 125 + 3) % 0x2AAAAB
                temp2 = (seed & 0xFFFF)

                crypt_table[index] = (temp1 | temp2)

                index += 0x100

        return crypt_table

    encryption_table = _prepare_encryption_table()

def main(argv):
    parser = OptionParser(usage="%prog [options] -f FILE", version=__version__)
    parser.add_option("-I", "--headers", action="store_true", dest="headers",
                      help="print header information from archive")
    parser.add_option("-x", "--extract", action="store_true", dest="extract",
                      help="extract files from archive")
    parser.add_option("-f", "--file", action="store", dest="file",
                      help="path to archive")
    options, args = parser.parse_args()
    if options.file:
        archive = MPQArchive(options.file)
        if options.headers:
            print archive.header
        if options.extract:
            archive.extract_to_disk()

if __name__ == '__main__':
    main(sys.argv)
<commit_msg>Refactored hash and block table extraction.<commit_after>#!/usr/bin/env python
# coding: utf-8

"""
mpyq is a Python library for reading MPQ (MoPaQ) archives.
"""

import cStringIO
import os
import struct
import sys
import zlib
from collections import namedtuple
from optparse import OptionParser

__author__ = "Aku Kotkavuo"
__version__ = "0.1"

MPQ_FILE_IMPLODE        = 0x00000100
MPQ_FILE_COMPRESS       = 0x00000200
MPQ_FILE_ENCRYPTED      = 0x00010000
MPQ_FILE_FIX_KEY        = 0x00020000
MPQ_FILE_SINGLE_UNIT    = 0x01000000
MPQ_FILE_DELETE_MARKER  = 0x02000000
MPQ_FILE_SECTOR_CRC     = 0x04000000
MPQ_FILE_EXISTS         = 0x80000000

MPQFileHeader = namedtuple('MPQFileHeader',
    '''
    magic
    header_size
    arhive_size
    format_version
    sector_size_shift
    hash_table_offset
    block_table_offset
    hash_table_entries
    block_table_entries
    '''
)
MPQFileHeader.struct_format = '<4s2I2H4I'

MPQFileHeaderExt = namedtuple('MPQFileHeaderExt',
    '''
    extended_block_table_offset
    hash_table_offset_high
    block_table_offset_high
    '''
)
MPQFileHeaderExt.struct_format = 'q2h'

MPQUserDataHeader = namedtuple('MPQUserDataHeader',
    '''
    magic
    user_data_size
    mpq_header_offset
    user_data_header_size
    '''
)
MPQUserDataHeader.struct_format = '<4s3I'

MPQHashTableEntry = namedtuple('MPQHashTableEntry',
    '''
    hash_a
    hash_b
    locale
    platform
    block_table_index
    '''
)
MPQHashTableEntry.struct_format = '2I2HI'

MPQBlockTableEntry = namedtuple('MPQBlockTableEntry',
    '''
    offset
    archived_size
    size
    flags
    '''
)
MPQBlockTableEntry.struct_format = '4I'

SC2ReplayHeader = namedtuple('SC2ReplayHeader',
    '''
    identifier
    release_flag
    major_version
    minor_version
    maintenance_version
    build_number
    duration
    '''
)
SC2ReplayHeader.struct_format = '>x19s2x4bi6xhx'

class MPQArchive(object):

    def __init__(self, filename):
        self.file = open(filename, 'rb')
        self.header = self.read_header()
        self.hash_table =  self.read_table('hash')
        self.block_table = self.read_table('block')
        self.files = self.read_file('(listfile)').splitlines()

    def read_header(self):
        """Read the header of a MPQ archive."""

        def read_mpq_header(offset=None):
            if offset:
                self.file.seek(offset)
            data = self.file.read(32)
            header = MPQFileHeader._make(
                struct.unpack(MPQFileHeader.struct_format, data))
            header = header._asdict()
            if header['format_version'] == 1:
                data = self.file.read(12)
                extended_header = MPQFileHeaderExt._make(
                    struct.unpack(MPQFileHeaderExt.struct_format, data))
                header.update(extended_header._asdict())
            return header

        def read_mpq_user_data_header():
            data = self.file.read(16)
            header = MPQUserDataHeader._make(
                struct.unpack(MPQUserDataHeader.struct_format, data))
            header = header._asdict()
            header['content'] = self.file.read(header['user_data_header_size'])
            if header['content'].startswith('\x15StarCraft II replay'):
                header['starcraft2_replay_header'] = SC2ReplayHeader._make(
                    struct.unpack(SC2ReplayHeader.struct_format,
                                  header['content']))
            return header

        magic = self.file.read(4)
        self.file.seek(0)

        if magic == 'MPQ\x1a':
            header = read_mpq_header()
            header['offset'] = 0
        elif magic == 'MPQ\x1b':
            user_data_header = read_mpq_user_data_header()
            header = read_mpq_header(user_data_header['mpq_header_offset'])
            header['offset'] = user_data_header['mpq_header_offset']
            header['user_data_header'] = user_data_header

        return header

    def read_table(self, table_type):
        """Read either hash or block table of a MPQ archive."""

        if table_type == 'hash':
            entry_class = MPQHashTableEntry
        elif table_type == 'block':
            entry_class = MPQBlockTableEntry
        else:
            raise ValueError("Invalid table type.")

        table_offset = self.header['%s_table_offset' % table_type]
        table_entries = self.header['%s_table_entries' % table_type]
        key = self._hash('(%s table)' % table_type, 'TABLE')

        self.file.seek(table_offset + self.header['offset'])
        data = self.file.read(table_entries * 16)
        data = self._decrypt(data, key)

        table = []
        for i in range(table_entries):
            entry_data = data[i*16:i*16+16]
            entry = entry_class._make(
                struct.unpack(entry_class.struct_format, entry_data))
            table.append(entry)
        return table

    def get_hash_table_entry(self, filename):
        """Get the hash table entry corresponding to filename."""
        hash_a = self._hash(filename, 'HASH_A')
        hash_b = self._hash(filename, 'HASH_B')
        for entry in self.hash_table:
            if (entry.hash_a == hash_a and entry.hash_b == hash_b):
                return entry

    def read_file(self, filename):
        """Read a file from the MPQ archive."""

        def decompress(data):
            """Read the compression type and decompress file data."""
            compression_type = ord(data[0])
            if compression_type == 0:
                return data
            if compression_type == 2:
                return zlib.decompress(data[1:], 15)

        hash_entry = self.get_hash_table_entry(filename)
        block_entry = self.block_table[hash_entry.block_table_index]

        # Read the block.
        if block_entry.flags & MPQ_FILE_EXISTS:
            offset = block_entry.offset + self.header['offset']
            self.file.seek(offset)
            file_data = self.file.read(block_entry.archived_size)

            if not block_entry.flags & MPQ_FILE_SINGLE_UNIT:
                # File consist of many sectors. They all need to be
                # decompressed separately and united.
                sector_size = 512 << self.header['sector_size_shift']
                sectors = block_entry.size / sector_size + 1
                if block_entry.flags & MPQ_FILE_SECTOR_CRC:
                    crc = True
                    sectors += 1
                positions = struct.unpack('<%dI' % (sectors + 1),
                                          file_data[:4*(sectors+1)])
                result = cStringIO.StringIO()
                for i in range(len(positions) - (2 if crc else 1)):
                    sector = file_data[positions[i]:positions[i+1]]
                    sector = decompress(sector)
                    result.write(sector)
                file_data = result.getvalue()
            else:
                # Single unit files only need to be decompressed.
                if block_entry.flags & MPQ_FILE_COMPRESS:
                    file_data = decompress(file_data)

            return file_data

    def extract(self):
        """Extract all the files inside the MPQ archive in memory."""
        return dict((f, self.read_file(f)) for f in self.files)

    def extract_to_disk(self):
        """Extract all files and write them to disk."""
        archive_name, extension = os.path.splitext(self.file.name)
        if not os.path.isdir(os.path.join(os.getcwd(), archive_name)):
            os.mkdir(archive_name)
        os.chdir(archive_name)
        for filename, data in self.extract().items():
            f = open(filename, 'wb')
            f.write(data)
            f.close()

    def _hash(self, string, hash_type):
        """Hash a string using MPQ's hash function."""
        hash_types = {
            'TABLE_OFFSET': 0,
            'HASH_A': 1,
            'HASH_B': 2,
            'TABLE': 3
        }
        seed1 = 0x7FED7FED
        seed2 = 0xEEEEEEEE

        for ch in string:
            ch = ord(ch.upper())
            value = self.encryption_table[(hash_types[hash_type] << 8) + ch]
            seed1 = (value ^ (seed1 + seed2)) & 0xFFFFFFFF
            seed2 = ch + seed1 + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

        return seed1

    def _decrypt(self, data, key):
        """Decrypt hash or block table or a sector."""
        seed1 = key
        seed2 = 0xEEEEEEEE
        result = cStringIO.StringIO()

        for i in range(len(data) // 4):
            seed2 += self.encryption_table[0x400 + (seed1 & 0xFF)]
            seed2 &= 0xFFFFFFFF
            value = struct.unpack("<I", data[i*4:i*4+4])[0]
            value = (value ^ (seed1 + seed2)) & 0xFFFFFFFF

            seed1 = ((~seed1 << 0x15) + 0x11111111) | (seed1 >> 0x0B)
            seed1 &= 0xFFFFFFFF
            seed2 = value + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

            result.write(struct.pack("<I", value))

        return result.getvalue()

    def _prepare_encryption_table():
        """Prepare encryption table for MPQ hash function."""
        seed = 0x00100001
        crypt_table = {}

        for i in range(256):
            index = i
            for j in range(5):
                seed = (seed * 125 + 3) % 0x2AAAAB
                temp1 = (seed & 0xFFFF) << 0x10;

                seed = (seed * 125 + 3) % 0x2AAAAB
                temp2 = (seed & 0xFFFF)

                crypt_table[index] = (temp1 | temp2)

                index += 0x100

        return crypt_table

    encryption_table = _prepare_encryption_table()

def main(argv):
    parser = OptionParser(usage="%prog [options] -f FILE", version=__version__)
    parser.add_option("-I", "--headers", action="store_true", dest="headers",
                      help="print header information from archive")
    parser.add_option("-x", "--extract", action="store_true", dest="extract",
                      help="extract files from archive")
    parser.add_option("-f", "--file", action="store", dest="file",
                      help="path to archive")
    options, args = parser.parse_args()
    if options.file:
        archive = MPQArchive(options.file)
        if options.headers:
            print archive.header
        if options.extract:
            archive.extract_to_disk()

if __name__ == '__main__':
    main(sys.argv)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding: utf-8

"""
mpyq is a Python library for reading MPQ (MoPaQ) archives.
"""

import argparse
import bz2
import cStringIO
import os
import struct
import zlib
from collections import namedtuple

__author__ = "Aku Kotkavuo"
__version__ = "0.1.4"

MPQ_FILE_IMPLODE        = 0x00000100
MPQ_FILE_COMPRESS       = 0x00000200
MPQ_FILE_ENCRYPTED      = 0x00010000
MPQ_FILE_FIX_KEY        = 0x00020000
MPQ_FILE_SINGLE_UNIT    = 0x01000000
MPQ_FILE_DELETE_MARKER  = 0x02000000
MPQ_FILE_SECTOR_CRC     = 0x04000000
MPQ_FILE_EXISTS         = 0x80000000

MPQFileHeader = namedtuple('MPQFileHeader',
    '''
    magic
    header_size
    arhive_size
    format_version
    sector_size_shift
    hash_table_offset
    block_table_offset
    hash_table_entries
    block_table_entries
    '''
)
MPQFileHeader.struct_format = '<4s2I2H4I'

MPQFileHeaderExt = namedtuple('MPQFileHeaderExt',
    '''
    extended_block_table_offset
    hash_table_offset_high
    block_table_offset_high
    '''
)
MPQFileHeaderExt.struct_format = 'q2h'

MPQUserDataHeader = namedtuple('MPQUserDataHeader',
    '''
    magic
    user_data_size
    mpq_header_offset
    user_data_header_size
    '''
)
MPQUserDataHeader.struct_format = '<4s3I'

MPQHashTableEntry = namedtuple('MPQHashTableEntry',
    '''
    hash_a
    hash_b
    locale
    platform
    block_table_index
    '''
)
MPQHashTableEntry.struct_format = '2I2HI'

MPQBlockTableEntry = namedtuple('MPQBlockTableEntry',
    '''
    offset
    archived_size
    size
    flags
    '''
)
MPQBlockTableEntry.struct_format = '4I'

class MPQArchive(object):

    def __init__(self, filename, listfile=True):
        """Create an MPQArchive object.

        You can skip reading the listfile if you pass listfile=False
        to the constructor. The 'files' attribute will be unavailable
        if you do this.
        """
        self.file = open(filename, 'rb')
        self.header = self.read_header()
        self.hash_table = self.read_table('hash')
        self.block_table = self.read_table('block')
        if listfile:
            self.files = self.read_file('(listfile)').splitlines()
        else:
            self.files = None

    def read_header(self):
        """Read the header of a MPQ archive."""

        def read_mpq_header(offset=None):
            if offset:
                self.file.seek(offset)
            data = self.file.read(32)
            header = MPQFileHeader._make(
                struct.unpack(MPQFileHeader.struct_format, data))
            header = header._asdict()
            if header['format_version'] == 1:
                data = self.file.read(12)
                extended_header = MPQFileHeaderExt._make(
                    struct.unpack(MPQFileHeaderExt.struct_format, data))
                header.update(extended_header._asdict())
            return header

        def read_mpq_user_data_header():
            data = self.file.read(16)
            header = MPQUserDataHeader._make(
                struct.unpack(MPQUserDataHeader.struct_format, data))
            header = header._asdict()
            header['content'] = self.file.read(header['user_data_header_size'])
            return header

        magic = self.file.read(4)
        self.file.seek(0)

        if magic == 'MPQ\x1a':
            header = read_mpq_header()
            header['offset'] = 0
        elif magic == 'MPQ\x1b':
            user_data_header = read_mpq_user_data_header()
            header = read_mpq_header(user_data_header['mpq_header_offset'])
            header['offset'] = user_data_header['mpq_header_offset']
            header['user_data_header'] = user_data_header

        return header

    def read_table(self, table_type):
        """Read either hash or block table of a MPQ archive."""

        if table_type == 'hash':
            entry_class = MPQHashTableEntry
        elif table_type == 'block':
            entry_class = MPQBlockTableEntry
        else:
            raise ValueError("Invalid table type.")

        table_offset = self.header['%s_table_offset' % table_type]
        table_entries = self.header['%s_table_entries' % table_type]
        key = self._hash('(%s table)' % table_type, 'TABLE')

        self.file.seek(table_offset + self.header['offset'])
        data = self.file.read(table_entries * 16)
        data = self._decrypt(data, key)

        def unpack_entry(position):
            entry_data = data[position*16:position*16+16]
            return entry_class._make(
                struct.unpack(entry_class.struct_format, entry_data))

        return [unpack_entry(i) for i in range(table_entries)]

    def get_hash_table_entry(self, filename):
        """Get the hash table entry corresponding to filename."""
        hash_a = self._hash(filename, 'HASH_A')
        hash_b = self._hash(filename, 'HASH_B')
        for entry in self.hash_table:
            if (entry.hash_a == hash_a and entry.hash_b == hash_b):
                return entry

    def read_file(self, filename):
        """Read a file from the MPQ archive."""

        def decompress(data):
            """Read the compression type and decompress file data."""
            compression_type = ord(data[0])
            if compression_type == 0:
                return data
            elif compression_type == 2:
                return zlib.decompress(data[1:], 15)
            elif compression_type == 16:
                return bz2.decompress(data[1:])
            else:
                raise RuntimeError("Unsupported compression type.")

        hash_entry = self.get_hash_table_entry(filename)
        block_entry = self.block_table[hash_entry.block_table_index]

        # Read the block.
        if block_entry.flags & MPQ_FILE_EXISTS:
            offset = block_entry.offset + self.header['offset']
            self.file.seek(offset)
            file_data = self.file.read(block_entry.archived_size)

            if block_entry.flags & MPQ_FILE_ENCRYPTED:
                raise NotImplementedError("Encryption is not supported yet.")

            if not block_entry.flags & MPQ_FILE_SINGLE_UNIT:
                # File consist of many sectors. They all need to be
                # decompressed separately and united.
                sector_size = 512 << self.header['sector_size_shift']
                sectors = block_entry.size / sector_size + 1
                if block_entry.flags & MPQ_FILE_SECTOR_CRC:
                    crc = True
                    sectors += 1
                else:
                    crc = False
                positions = struct.unpack('<%dI' % (sectors + 1),
                                          file_data[:4*(sectors+1)])
                result = cStringIO.StringIO()
                for i in range(len(positions) - (2 if crc else 1)):
                    sector = file_data[positions[i]:positions[i+1]]
                    sector = decompress(sector)
                    result.write(sector)
                file_data = result.getvalue()
            else:
                # Single unit files only need to be decompressed, but
                # compression only happens when at least one byte is gained.
                if (block_entry.flags & MPQ_FILE_COMPRESS and
                    block_entry.size > block_entry.archived_size):
                    file_data = decompress(file_data)

            return file_data

    def extract(self):
        """Extract all the files inside the MPQ archive in memory."""
        if self.files:
            return dict((f, self.read_file(f)) for f in self.files)
        else:
            raise RuntimeError("Can't extract whole archive without listfile.")

    def extract_to_disk(self):
        """Extract all files and write them to disk."""
        archive_name, extension = os.path.splitext(self.file.name)
        if not os.path.isdir(os.path.join(os.getcwd(), archive_name)):
            os.mkdir(archive_name)
        os.chdir(archive_name)
        for filename, data in self.extract().items():
            f = open(filename, 'wb')
            f.write(data)
            f.close()

    def extract_files(self, *filenames):
        """Extract given files from the archive to disk."""
        for filename in filenames:
            data = self.read_file(filename)
            f = open(filename, 'wb')
            f.write(data)
            f.close()

    def print_headers(self):
        print "MPQ archive header"
        print "------------------"
        for key, value in self.header.iteritems():
            if key == "user_data_header":
                continue
            print "{0:30} {1!r}".format(key, value)
        if self.header['user_data_header']:
            print
            print "MPQ user data header"
            print "--------------------"
            for key, value in self.header['user_data_header'].iteritems():
                print "{0:30} {1!r}".format(key, value)

    def print_hash_table(self):
        print "MPQ archive hash table"
        print "----------------------"
        print " Hash A   Hash B  Locl Plat BlockIdx"
        for entry in self.hash_table:
            print '%08X %08X %04X %04X %08X' % entry
        print

    def print_block_table(self):
        print "MPQ archive block table"
        print "-----------------------"
        print " Offset  ArchSize RealSize  Flags"
        for entry in self.block_table:
            print '%08X %8d %8d %8X' % entry
        print

    def print_files(self):
        if self.files:
            for filename in self.files:
                hash_entry = self.get_hash_table_entry(filename)
                block_entry = self.block_table[hash_entry.block_table_index]
                print "{0:30} {1:>8} bytes".format(filename, block_entry.size)

    def _hash(self, string, hash_type):
        """Hash a string using MPQ's hash function."""
        hash_types = {
            'TABLE_OFFSET': 0,
            'HASH_A': 1,
            'HASH_B': 2,
            'TABLE': 3
        }
        seed1 = 0x7FED7FED
        seed2 = 0xEEEEEEEE

        for ch in string:
            ch = ord(ch.upper())
            value = self.encryption_table[(hash_types[hash_type] << 8) + ch]
            seed1 = (value ^ (seed1 + seed2)) & 0xFFFFFFFF
            seed2 = ch + seed1 + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

        return seed1

    def _decrypt(self, data, key):
        """Decrypt hash or block table or a sector."""
        seed1 = key
        seed2 = 0xEEEEEEEE
        result = cStringIO.StringIO()

        for i in range(len(data) // 4):
            seed2 += self.encryption_table[0x400 + (seed1 & 0xFF)]
            seed2 &= 0xFFFFFFFF
            value = struct.unpack("<I", data[i*4:i*4+4])[0]
            value = (value ^ (seed1 + seed2)) & 0xFFFFFFFF

            seed1 = ((~seed1 << 0x15) + 0x11111111) | (seed1 >> 0x0B)
            seed1 &= 0xFFFFFFFF
            seed2 = value + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

            result.write(struct.pack("<I", value))

        return result.getvalue()

    def _prepare_encryption_table():
        """Prepare encryption table for MPQ hash function."""
        seed = 0x00100001
        crypt_table = {}

        for i in range(256):
            index = i
            for j in range(5):
                seed = (seed * 125 + 3) % 0x2AAAAB
                temp1 = (seed & 0xFFFF) << 0x10;

                seed = (seed * 125 + 3) % 0x2AAAAB
                temp2 = (seed & 0xFFFF)

                crypt_table[index] = (temp1 | temp2)

                index += 0x100

        return crypt_table

    encryption_table = _prepare_encryption_table()

def main():
    description = "mpyq reads and extracts MPQ archives."
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("file", action="store", help="path to the archive")
    parser.add_argument("-I", "--headers", action="store_true", dest="headers",
                        help="print header information from the archive")
    parser.add_argument("-H", "--hash-table", action="store_true",
                        dest="hash_table", help="print hash table"),
    parser.add_argument("-b", "--block-table", action="store_true",
                        dest="block_table", help="print block table"),
    parser.add_argument("-s", "--skip-listfile", action="store_true",
                        dest="skip_listfile", help="skip reading (listfile)"),
    parser.add_argument("-t", "--list-files", action="store_true", dest="list",
                        help="list files inside the archive")
    parser.add_argument("-x", "--extract", action="store_true", dest="extract",
                        help="extract files from the archive")
    args = parser.parse_args()
    if args.file:
        if not args.skip_listfile:
            archive = MPQArchive(args.file)
        else:
            archive = MPQArchive(args.file, listfile=False)
        if args.headers:
            archive.print_headers()
        if args.hash_table:
            archive.print_hash_table()
        if args.block_table:
            archive.print_block_table()
        if args.list:
            archive.print_files()
        if args.extract:
            archive.extract_to_disk()

if __name__ == '__main__':
    main()
<commit_msg>Bumped version to 0.1.5.<commit_after>#!/usr/bin/env python
# coding: utf-8

"""
mpyq is a Python library for reading MPQ (MoPaQ) archives.
"""

import argparse
import bz2
import cStringIO
import os
import struct
import zlib
from collections import namedtuple

__author__ = "Aku Kotkavuo"
__version__ = "0.1.5"

MPQ_FILE_IMPLODE        = 0x00000100
MPQ_FILE_COMPRESS       = 0x00000200
MPQ_FILE_ENCRYPTED      = 0x00010000
MPQ_FILE_FIX_KEY        = 0x00020000
MPQ_FILE_SINGLE_UNIT    = 0x01000000
MPQ_FILE_DELETE_MARKER  = 0x02000000
MPQ_FILE_SECTOR_CRC     = 0x04000000
MPQ_FILE_EXISTS         = 0x80000000

MPQFileHeader = namedtuple('MPQFileHeader',
    '''
    magic
    header_size
    arhive_size
    format_version
    sector_size_shift
    hash_table_offset
    block_table_offset
    hash_table_entries
    block_table_entries
    '''
)
MPQFileHeader.struct_format = '<4s2I2H4I'

MPQFileHeaderExt = namedtuple('MPQFileHeaderExt',
    '''
    extended_block_table_offset
    hash_table_offset_high
    block_table_offset_high
    '''
)
MPQFileHeaderExt.struct_format = 'q2h'

MPQUserDataHeader = namedtuple('MPQUserDataHeader',
    '''
    magic
    user_data_size
    mpq_header_offset
    user_data_header_size
    '''
)
MPQUserDataHeader.struct_format = '<4s3I'

MPQHashTableEntry = namedtuple('MPQHashTableEntry',
    '''
    hash_a
    hash_b
    locale
    platform
    block_table_index
    '''
)
MPQHashTableEntry.struct_format = '2I2HI'

MPQBlockTableEntry = namedtuple('MPQBlockTableEntry',
    '''
    offset
    archived_size
    size
    flags
    '''
)
MPQBlockTableEntry.struct_format = '4I'

class MPQArchive(object):

    def __init__(self, filename, listfile=True):
        """Create an MPQArchive object.

        You can skip reading the listfile if you pass listfile=False
        to the constructor. The 'files' attribute will be unavailable
        if you do this.
        """
        self.file = open(filename, 'rb')
        self.header = self.read_header()
        self.hash_table = self.read_table('hash')
        self.block_table = self.read_table('block')
        if listfile:
            self.files = self.read_file('(listfile)').splitlines()
        else:
            self.files = None

    def read_header(self):
        """Read the header of a MPQ archive."""

        def read_mpq_header(offset=None):
            if offset:
                self.file.seek(offset)
            data = self.file.read(32)
            header = MPQFileHeader._make(
                struct.unpack(MPQFileHeader.struct_format, data))
            header = header._asdict()
            if header['format_version'] == 1:
                data = self.file.read(12)
                extended_header = MPQFileHeaderExt._make(
                    struct.unpack(MPQFileHeaderExt.struct_format, data))
                header.update(extended_header._asdict())
            return header

        def read_mpq_user_data_header():
            data = self.file.read(16)
            header = MPQUserDataHeader._make(
                struct.unpack(MPQUserDataHeader.struct_format, data))
            header = header._asdict()
            header['content'] = self.file.read(header['user_data_header_size'])
            return header

        magic = self.file.read(4)
        self.file.seek(0)

        if magic == 'MPQ\x1a':
            header = read_mpq_header()
            header['offset'] = 0
        elif magic == 'MPQ\x1b':
            user_data_header = read_mpq_user_data_header()
            header = read_mpq_header(user_data_header['mpq_header_offset'])
            header['offset'] = user_data_header['mpq_header_offset']
            header['user_data_header'] = user_data_header

        return header

    def read_table(self, table_type):
        """Read either hash or block table of a MPQ archive."""

        if table_type == 'hash':
            entry_class = MPQHashTableEntry
        elif table_type == 'block':
            entry_class = MPQBlockTableEntry
        else:
            raise ValueError("Invalid table type.")

        table_offset = self.header['%s_table_offset' % table_type]
        table_entries = self.header['%s_table_entries' % table_type]
        key = self._hash('(%s table)' % table_type, 'TABLE')

        self.file.seek(table_offset + self.header['offset'])
        data = self.file.read(table_entries * 16)
        data = self._decrypt(data, key)

        def unpack_entry(position):
            entry_data = data[position*16:position*16+16]
            return entry_class._make(
                struct.unpack(entry_class.struct_format, entry_data))

        return [unpack_entry(i) for i in range(table_entries)]

    def get_hash_table_entry(self, filename):
        """Get the hash table entry corresponding to filename."""
        hash_a = self._hash(filename, 'HASH_A')
        hash_b = self._hash(filename, 'HASH_B')
        for entry in self.hash_table:
            if (entry.hash_a == hash_a and entry.hash_b == hash_b):
                return entry

    def read_file(self, filename):
        """Read a file from the MPQ archive."""

        def decompress(data):
            """Read the compression type and decompress file data."""
            compression_type = ord(data[0])
            if compression_type == 0:
                return data
            elif compression_type == 2:
                return zlib.decompress(data[1:], 15)
            elif compression_type == 16:
                return bz2.decompress(data[1:])
            else:
                raise RuntimeError("Unsupported compression type.")

        hash_entry = self.get_hash_table_entry(filename)
        block_entry = self.block_table[hash_entry.block_table_index]

        # Read the block.
        if block_entry.flags & MPQ_FILE_EXISTS:
            offset = block_entry.offset + self.header['offset']
            self.file.seek(offset)
            file_data = self.file.read(block_entry.archived_size)

            if block_entry.flags & MPQ_FILE_ENCRYPTED:
                raise NotImplementedError("Encryption is not supported yet.")

            if not block_entry.flags & MPQ_FILE_SINGLE_UNIT:
                # File consist of many sectors. They all need to be
                # decompressed separately and united.
                sector_size = 512 << self.header['sector_size_shift']
                sectors = block_entry.size / sector_size + 1
                if block_entry.flags & MPQ_FILE_SECTOR_CRC:
                    crc = True
                    sectors += 1
                else:
                    crc = False
                positions = struct.unpack('<%dI' % (sectors + 1),
                                          file_data[:4*(sectors+1)])
                result = cStringIO.StringIO()
                for i in range(len(positions) - (2 if crc else 1)):
                    sector = file_data[positions[i]:positions[i+1]]
                    sector = decompress(sector)
                    result.write(sector)
                file_data = result.getvalue()
            else:
                # Single unit files only need to be decompressed, but
                # compression only happens when at least one byte is gained.
                if (block_entry.flags & MPQ_FILE_COMPRESS and
                    block_entry.size > block_entry.archived_size):
                    file_data = decompress(file_data)

            return file_data

    def extract(self):
        """Extract all the files inside the MPQ archive in memory."""
        if self.files:
            return dict((f, self.read_file(f)) for f in self.files)
        else:
            raise RuntimeError("Can't extract whole archive without listfile.")

    def extract_to_disk(self):
        """Extract all files and write them to disk."""
        archive_name, extension = os.path.splitext(self.file.name)
        if not os.path.isdir(os.path.join(os.getcwd(), archive_name)):
            os.mkdir(archive_name)
        os.chdir(archive_name)
        for filename, data in self.extract().items():
            f = open(filename, 'wb')
            f.write(data)
            f.close()

    def extract_files(self, *filenames):
        """Extract given files from the archive to disk."""
        for filename in filenames:
            data = self.read_file(filename)
            f = open(filename, 'wb')
            f.write(data)
            f.close()

    def print_headers(self):
        print "MPQ archive header"
        print "------------------"
        for key, value in self.header.iteritems():
            if key == "user_data_header":
                continue
            print "{0:30} {1!r}".format(key, value)
        if self.header['user_data_header']:
            print
            print "MPQ user data header"
            print "--------------------"
            for key, value in self.header['user_data_header'].iteritems():
                print "{0:30} {1!r}".format(key, value)

    def print_hash_table(self):
        print "MPQ archive hash table"
        print "----------------------"
        print " Hash A   Hash B  Locl Plat BlockIdx"
        for entry in self.hash_table:
            print '%08X %08X %04X %04X %08X' % entry
        print

    def print_block_table(self):
        print "MPQ archive block table"
        print "-----------------------"
        print " Offset  ArchSize RealSize  Flags"
        for entry in self.block_table:
            print '%08X %8d %8d %8X' % entry
        print

    def print_files(self):
        if self.files:
            for filename in self.files:
                hash_entry = self.get_hash_table_entry(filename)
                block_entry = self.block_table[hash_entry.block_table_index]
                print "{0:30} {1:>8} bytes".format(filename, block_entry.size)

    def _hash(self, string, hash_type):
        """Hash a string using MPQ's hash function."""
        hash_types = {
            'TABLE_OFFSET': 0,
            'HASH_A': 1,
            'HASH_B': 2,
            'TABLE': 3
        }
        seed1 = 0x7FED7FED
        seed2 = 0xEEEEEEEE

        for ch in string:
            ch = ord(ch.upper())
            value = self.encryption_table[(hash_types[hash_type] << 8) + ch]
            seed1 = (value ^ (seed1 + seed2)) & 0xFFFFFFFF
            seed2 = ch + seed1 + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

        return seed1

    def _decrypt(self, data, key):
        """Decrypt hash or block table or a sector."""
        seed1 = key
        seed2 = 0xEEEEEEEE
        result = cStringIO.StringIO()

        for i in range(len(data) // 4):
            seed2 += self.encryption_table[0x400 + (seed1 & 0xFF)]
            seed2 &= 0xFFFFFFFF
            value = struct.unpack("<I", data[i*4:i*4+4])[0]
            value = (value ^ (seed1 + seed2)) & 0xFFFFFFFF

            seed1 = ((~seed1 << 0x15) + 0x11111111) | (seed1 >> 0x0B)
            seed1 &= 0xFFFFFFFF
            seed2 = value + seed2 + (seed2 << 5) + 3 & 0xFFFFFFFF

            result.write(struct.pack("<I", value))

        return result.getvalue()

    def _prepare_encryption_table():
        """Prepare encryption table for MPQ hash function."""
        seed = 0x00100001
        crypt_table = {}

        for i in range(256):
            index = i
            for j in range(5):
                seed = (seed * 125 + 3) % 0x2AAAAB
                temp1 = (seed & 0xFFFF) << 0x10;

                seed = (seed * 125 + 3) % 0x2AAAAB
                temp2 = (seed & 0xFFFF)

                crypt_table[index] = (temp1 | temp2)

                index += 0x100

        return crypt_table

    encryption_table = _prepare_encryption_table()

def main():
    description = "mpyq reads and extracts MPQ archives."
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("file", action="store", help="path to the archive")
    parser.add_argument("-I", "--headers", action="store_true", dest="headers",
                        help="print header information from the archive")
    parser.add_argument("-H", "--hash-table", action="store_true",
                        dest="hash_table", help="print hash table"),
    parser.add_argument("-b", "--block-table", action="store_true",
                        dest="block_table", help="print block table"),
    parser.add_argument("-s", "--skip-listfile", action="store_true",
                        dest="skip_listfile", help="skip reading (listfile)"),
    parser.add_argument("-t", "--list-files", action="store_true", dest="list",
                        help="list files inside the archive")
    parser.add_argument("-x", "--extract", action="store_true", dest="extract",
                        help="extract files from the archive")
    args = parser.parse_args()
    if args.file:
        if not args.skip_listfile:
            archive = MPQArchive(args.file)
        else:
            archive = MPQArchive(args.file, listfile=False)
        if args.headers:
            archive.print_headers()
        if args.hash_table:
            archive.print_hash_table()
        if args.block_table:
            archive.print_block_table()
        if args.list:
            archive.print_files()
        if args.extract:
            archive.extract_to_disk()

if __name__ == '__main__':
    main()
<|endoftext|>