<commit_before>#! /usr/bin/python3
# System imports
from datetime import datetime
import sys
import getopt

# Package imports
try:
    from bibtexparser.bparser import BibTexParser
except ImportError:
    print("We need the BibTeX parser from the bibtexparser package on PyPI. "
          "If pip is installed, `pip install bibtexparser` will install the "
          "package. Code is tested with version 0.4, may not work with other "
          "versions.")

# First is to define a function to format the names we get from BibTeX,
# since this task will be the same for every paper type. The current
# format is "F.M. Last, F.M. Last, and F.M. Last".
def reorder(names, faname='F.A. Author'):
    """Format the string of author names and return a string.

    Adapated from one of the `customization` functions in
    `bibtexparser`.
    INPUT:
    names -- string of names to be formatted. The names from BibTeX are
             formatted in the style "Last, First Middle and Last, First
             Middle and Last, First Middle" and this is the expected
             style here.
    faname -- string of the initialized name of the author to whom
              formatting will be applied
              default: 'F.A. Author'
    OUTPUT:
    nameout -- string of formatted names.

    """
    # Set the format tag for the website's owner, to highlight where on
    # the author list the website owner is. Default is **
    myNameFormatTag = '**'

    # Convert the input string to a list by splitting the string at the
    # " and " and strip out any remaining whitespace.
    nameslist = [i.strip() for i in names.replace('\n', ' ').split(" and ")]

    # Initialize a list to store the names after they've been tidied
    # up.
    tidynames = []

    # Loop through each name in the list.
    for namestring in nameslist:
        # Strip whitespace from the string
        namestring = namestring.strip()

        # If, for some reason, we've gotten a blank name, skip it
        if len(namestring) < 1:
            continue

        # Split the `namestring` at the comma, but only perform the
        # split once.
        namesplit = namestring.split(',', 1)

        # In the expected format, the first element of the split
        # namestring is the last name. Strip any whitespace.
        last = namesplit[0].strip()

        # There could be many first/middle names, so we collect them in
        # a list. All of the first/middle names are stored in the
        # second element of namesplit seperated by whitespace. Split
        # the first/middle names at the whitespace then strip out any
        # remaining whitespace and any periods (the periods will be
        # added in the proper place later).
        firsts = [i.strip().strip('.') for i in namesplit[1].split()]

        # Handle the case of juniors.
        if last in ['jnr', 'jr', 'junior']:
            last = firsts.pop()

        # Handle the case of "from" in the name
        for item in firsts:
            if item in ['ben', 'van', 'der', 'de', 'la', 'le']:
                last = firsts.pop() + ' ' + last

        # For the case of hyphenated first names, we need to split at
        # the hyphen as well. Possible bug: this only works if the
        # first first name is the hyphenated one, and this replaces all
        # of thefirst names with the names split at the hyphen. We'd
        # like to handle multiple hyphens or a hyphenated name with an
        # initial more intelligently.
        if '-' in firsts[0]:
            firsts = firsts[0].split('-')

        # Now that all the first name edge cases are sorted out, we
        # want to initialize all the first names. Set the variable
        # initials to an empty string to we can add to it. Then loop
        # through each of the items in the list of first names. Take
        # the first element of each item and append a period, but no
        # space.
        initials = ''
        for item in firsts:
            initials = initials + item[0] + '.'

        # Stick all of the parts of the name together in `tidynames`
        tidynames.append(initials + ' ' + last)

    # Find the case of the website author and set the format for that
    # name
    try:
        i = tidynames.index(faname)
        tidynames[i] = myNameFormatTag + tidynames[i] + myNameFormatTag
    except ValueError:
        print("Couldn't find ",faname,"in the names list. Sorry!")

    # Handle the various cases of number of authors and how they should
    # be joined. Convert the elements of `tidynames` to a string.
    if len(tidynames) > 2:
        tidynames[-1] = 'and ' + tidynames[-1]
        nameout = ', '.join(tidynames)
    elif len(tidynames) == 2:
        tidynames[-1] = 'and ' + tidynames[-1]
        nameout = ' '.join(tidynames)
    else:
        # If `tidynames` only has one name, we only need to convert it
        # to a string. The first way that came to mind was to join the
        # list to an empty string.
        nameout = ''.join(tidynames)

    # Return `nameout`, the string of formatted authors
    return nameout

def main(argv):
    bibFileName = 'refs.bib'
    outputFileName = 'pubs.md'
    faname = 'F.A. Author'
    help = ("Usage:\n"
    "-h, --help: Print this help dialog and exit\n"
    "-b filename, --bibfile=filename: Set the filename of the "
    "BibTeX reference file. Default: refs.bib\n"
    "-o filename, --output=filename: Set the filename of the "
    "kramdown output. Default: pubs.md\n"
    "-a 'author', --author='f.a. name': Set the name of the author "
    "to be highlighted. Default: 'F.A. Author'")
    try:
        opts, args = getopt.getopt(argv, "hb:o:a:",
                                   ["help", "bibfile=", "output=",
                                   "author="])
    except getopt.GetoptError as e:
        print("You did not enter an option properly. Please try again.")
        print(e)
        print(help)
        sys.exit(2)

    # If the user doesn't input any options, print the help.
    if not opts:
        print(help)
        sys.exit(0)

    for opt, arg in opts:
        if opt in {"-h", "--help"}:
            print(help)
            sys.exit()
        elif opt in {"-b", "--bibfile"}:
            bibFileName = arg
        elif opt in {"-o", "--output"}:
            outputFileName = arg
        elif opt in {"-a", "--author"}:
            faname = arg

    # Set the formatting identifiers. Since we're using kramdown, we
    # don't have to use the HTML tags.
    em = '_'
    st = '**'

    # Set some HTML tags to go around each part of the reference.
    openSpan = ''
    closeSpan = ''

    # Open and parse the BibTeX file in `bibFileName` using
    # `bibtexparser`

    with open(bibFileName,'r') as bibFile:
        bp = BibTexParser(bibFile)

    # Get a dictionary of dictionaries of key, value pairs from the
    # BibTeX file. The structure is
    # {ID:{authors:...},ID:{authors:...}}.
    refsdict = bp.get_entry_dict()

    # Create a list of all the types of documents found in the BibTeX
    # file, typically `article`, `inproceedings`, and `phdthesis`.
    # Dedupe the list.
    types = []
    for k,ref in refsdict.items():
        types.append(ref["type"])
    types  = set(types)

    # For each of the types of reference, we need to sort each by month
    # then year. We store the dictionary representing each reference in
    # a sorted list for each type of reference. Then we store each of
    # these sorted lists in a dictionary whose key is the type of
    # reference and value is the list of dictionaries.
    sortDict = {}
    for type in types:
        temp = sorted([val for key, val in refsdict.items()
                      if val["type"] == type], key=lambda l:
                      datetime.strptime(l["month"],'%b').month, reverse=True)
        sortDict[type] = sorted(temp, key=lambda k: k["year"], reverse=True)

    # Open the output file with utf-8 encoding, write mode, and Unix
    # newlines.
    with open(outputFileName, encoding='utf-8', mode='w',
              newline='') as outFile:

        # Start with journal articles. Print the header to the screen
        # and output file.
        print('Journal Articles\n---\n')
        outFile.write('Journal Articles\n---\n')

        # To get the year numbering correct, we have to set a dummy
        # value for pubyear (usage described below).
        pubyear = ''

        # Loop through all the references in the article type. The
        # logic in this loop (and the loops for the other reference
        # types) is not amenable to generalization due to different
        # information for each reference type. Therefore, its easiest
        # to write out the logic for each loop instead of writing the
        # logic into a function and calling that.
        for ref in sortDict["article"]:
            # Get the string of author names in the proper format from
            # the `reorder` function. Get some other information. Hack
            # the journal title to remove the '\' before '&' in
            # 'Energy & Fuels' because Mendeley inserts an extra '\'
            # into the BibTeX.
            authors = reorder(ref["author"],faname)
            title = ref["title"]
            journal = ref["journal"]
            if '\&' in journal:
                words = journal.strip().split('\&')
                journal = words[0] + '&' + words[1]
            doi = ref["doi"]

            # Get the publication year. If the year of the current
            # reference is not equal to the year of the previous
            # reference, we need to print the year out and set
            # `pubyear` equal to `year`.
            year = ref["year"]
            if year != pubyear:
                pubyear = year
                stri = '\n{:.year}\n###' + year + '\n'
                print(stri)
                outFile.write(stri)

            # Start building the string containing the formatted
            # reference. Each bit should be surrounded by a span. The
            # {:.XYZ} is the kramdown notation to add that class to the
            # HTML tag. Each line should be ended with two spaces
            # before the newline so that kramdown inserts an HTML

            # there.
            string = ('\n{:.paper}\n' +
                      openSpan + title + closeSpan + '{:.papertitle}  \n' +
                      openSpan + authors + closeSpan + '{:.authors}  \n' +
                      openSpan + em + journal + em + ', '
                      )

            # Not all journal articles will have vol., no., and pp.
            # because some may be "In Press".
            if "volume" in ref:
                string = string + 'vol. ' + ref["volume"] + ', '

            if "number" in ref:
                string = string + 'no. ' + ref["number"] + ', '

            if "pages" in ref:
                string = string + 'pp. ' + ref["pages"] + ', '

            string = (string + ref["month"].title() + '. ' +
                      year + closeSpan + '{:.journal}  \n')

            if "doi" in ref:
                string = (string + openSpan + st + 'DOI:' + st + ' [' +
                          ref["doi"] + '](http://dx.doi.org/' + ref["doi"]
                          + ')' + closeSpan + '{:.doi}  \n'
                          )

            # Extra comments, such as links to files, should be stored
            # as "Notes" for each reference in Mendeley. Mendeley will
            # export this field with the tag "annote" in BibTeX.
            if "annote" in ref:
                string = (string + openSpan + ref["annote"].replace('\\','') +
                           closeSpan + '{:.comment}  \n'
                         )
            print(string)
            outFile.write(string)

        # Next are conference papers and posters. Print the header to
        # the screen and the output file.
        print('Conference Publications and Posters\n---\n')
        outFile.write('\nConference Publications and Posters\n---\n')

        # Same trick for the pubyear as for the journal articles.
        pubyear = ''

        # Loop through the references in the `inproceedings` type.
        for ref in sortDict["inproceedings"]:
            authors = reorder(ref["author"],faname)
            title = ref["title"]
            year = ref["year"]
            if year != pubyear:
                pubyear = year
                stri = '\n{:.year}\n###' + year + '\n'
                print(stri)
                outFile.write(stri)

            # Start building the reference string.
            string = ('\n{:.paper}\n' +
                      openSpan + title + closeSpan + '{:.papertitle}  \n' +
                      openSpan + authors + closeSpan + '{:.authors}  \n' +
                      openSpan
                     )

            # Since Mendeley doesn't allow customization of BibTeX
            # output, we hack the "pages" field to contain the paper
            # number for the conference paper. Not all of this type of
            # reference will have this, so we check for it.
            if "pages" in ref:
                paperno = ref["pages"]
                string = string + paperno + ', '

            # Insert the conference title, stored in the "booktitle"
            # field.
            conf = ref["booktitle"]
            string = string + conf + ', '
            if "organization" in ref:
                string = string + ref["organization"] + ', '
            if "address" in ref:
                string = string + ref["address"] + ', '

            string = (string + ref["month"].title() + '. ' + year + closeSpan
                      + '{:.journal}  \n'
                     )

            if "doi" in ref:
                string = (string + openSpan + st + 'DOI:' + st + ' [' +
                          ref["doi"] + '](http://dx.doi.org/' + ref["doi"]
                          + ')' + closeSpan + '{:.doi}  \n'
                         )

            # Extra comments, such as links to files, should be stored
            # as "Notes" for each reference in Mendeley. Mendeley will
            # export this field with the tag "annote" in BibTeX.
            if "annote" in ref:
                string = (string + openSpan + ref["annote"].replace('\\','') +
                          closeSpan + '{:.comment}  \n'
                         )
            print(string)
            outFile.write(string)

        # Finally are the theses and dissertations. Same general logic
        # as for the other reference types.
        print("Master's Thesis\n---\n")
        outFile.write("\nMaster's Thesis\n---\n")
        pubyear = '2200'
        for ref in sortDict["phdthesis"]:
            authors = reorder(ref["author"],faname)
            title = ref["title"]
            year = ref["year"]
            if year != pubyear:
                pubyear = year
                stri = '\n{:.year}\n###' + year + '\n'
                print(stri)
                outFile.write(stri)

            string = ('\n{:.paper}\n' +
                      openSpan + title + closeSpan + '{:.papertitle}  \n' +
                      openSpan + authors + closeSpan + '{:.authors}  \n' +
                      openSpan )
            if "school" in ref:
                string = string + ref["school"] + ', '
            if "month" in ref:
                string = string + ref["month"].title() + '. '

            string = string + year + closeSpan + '{:.journal}  \n'

            if "annote" in ref:
                string = (string + openSpan + ref["annote"].replace('\\','') +
                          closeSpan + '{:.comment}  \n'
                         )
            print(string)
            outFile.write(string)

if __name__ == "__main__":
    main(sys.argv[1:])
<commit_msg>PEP-8 the variables<commit_after>#! /usr/bin/python3
# System imports
from datetime import datetime
import sys
import getopt

# Package imports
try:
    from bibtexparser.bparser import BibTexParser
except ImportError:
    print("We need the BibTeX parser from the bibtexparser package on PyPI. "
          "If pip is installed, `pip install bibtexparser` will install the "
          "package. Code is tested with version 0.4, may not work with other "
          "versions.")

# First is to define a function to format the names we get from BibTeX,
# since this task will be the same for every paper type. The current
# format is "F.M. Last, F.M. Last, and F.M. Last".
def reorder(names, faname='F.A. Author'):
    """Format the string of author names and return a string.

    Adapated from one of the `customization` functions in
    `bibtexparser`.
    INPUT:
    names -- string of names to be formatted. The names from BibTeX are
             formatted in the style "Last, First Middle and Last, First
             Middle and Last, First Middle" and this is the expected
             style here.
    faname -- string of the initialized name of the author to whom
              formatting will be applied
              default: 'F.A. Author'
    OUTPUT:
    nameout -- string of formatted names.

    """
    # Set the format tag for the website's owner, to highlight where on
    # the author list the website owner is. Default is **
    my_name_format_tag = '**'

    # Convert the input string to a list by splitting the string at the
    # " and " and strip out any remaining whitespace.
    nameslist = [i.strip() for i in names.replace('\n', ' ').split(" and ")]

    # Initialize a list to store the names after they've been tidied
    # up.
    tidynames = []

    # Loop through each name in the list.
    for namestring in nameslist:
        # Strip whitespace from the string
        namestring = namestring.strip()

        # If, for some reason, we've gotten a blank name, skip it
        if len(namestring) < 1:
            continue

        # Split the `namestring` at the comma, but only perform the
        # split once.
        namesplit = namestring.split(',', 1)

        # In the expected format, the first element of the split
        # namestring is the last name. Strip any whitespace.
        last = namesplit[0].strip()

        # There could be many first/middle names, so we collect them in
        # a list. All of the first/middle names are stored in the
        # second element of namesplit seperated by whitespace. Split
        # the first/middle names at the whitespace then strip out any
        # remaining whitespace and any periods (the periods will be
        # added in the proper place later).
        firsts = [i.strip().strip('.') for i in namesplit[1].split()]

        # Handle the case of juniors.
        if last in ['jnr', 'jr', 'junior']:
            last = firsts.pop()

        # Handle the case of "from" in the name
        for item in firsts:
            if item in ['ben', 'van', 'der', 'de', 'la', 'le']:
                last = firsts.pop() + ' ' + last

        # For the case of hyphenated first names, we need to split at
        # the hyphen as well. Possible bug: this only works if the
        # first first name is the hyphenated one, and this replaces all
        # of thefirst names with the names split at the hyphen. We'd
        # like to handle multiple hyphens or a hyphenated name with an
        # initial more intelligently.
        if '-' in firsts[0]:
            firsts = firsts[0].split('-')

        # Now that all the first name edge cases are sorted out, we
        # want to initialize all the first names. Set the variable
        # initials to an empty string to we can add to it. Then loop
        # through each of the items in the list of first names. Take
        # the first element of each item and append a period, but no
        # space.
        initials = ''
        for item in firsts:
            initials = initials + item[0] + '.'

        # Stick all of the parts of the name together in `tidynames`
        tidynames.append(initials + ' ' + last)

    # Find the case of the website author and set the format for that
    # name
    try:
        i = tidynames.index(faname)
        tidynames[i] = my_name_format_tag + tidynames[i] + my_name_format_tag
    except ValueError:
        print("Couldn't find ",faname,"in the names list. Sorry!")

    # Handle the various cases of number of authors and how they should
    # be joined. Convert the elements of `tidynames` to a string.
    if len(tidynames) > 2:
        tidynames[-1] = 'and ' + tidynames[-1]
        nameout = ', '.join(tidynames)
    elif len(tidynames) == 2:
        tidynames[-1] = 'and ' + tidynames[-1]
        nameout = ' '.join(tidynames)
    else:
        # If `tidynames` only has one name, we only need to convert it
        # to a string. The first way that came to mind was to join the
        # list to an empty string.
        nameout = ''.join(tidynames)

    # Return `nameout`, the string of formatted authors
    return nameout

def main(argv):
    bib_file_name = 'refs.bib'
    output_file_name = 'pubs.md'
    faname = 'F.A. Author'
    help = ("Usage:\n"
    "-h, --help: Print this help dialog and exit\n"
    "-b filename, --bibfile=filename: Set the filename of the "
    "BibTeX reference file. Default: refs.bib\n"
    "-o filename, --output=filename: Set the filename of the "
    "kramdown output. Default: pubs.md\n"
    "-a 'author', --author='f.a. name': Set the name of the author "
    "to be highlighted. Default: 'F.A. Author'")
    try:
        opts, args = getopt.getopt(argv, "hb:o:a:",
                                   ["help", "bibfile=", "output=",
                                   "author="])
    except getopt.GetoptError as e:
        print("You did not enter an option properly. Please try again.")
        print(e)
        print(help)
        sys.exit(1)

    # If the user doesn't input any options, print the help.
    if not opts:
        print(help)
        sys.exit(0)

    for opt, arg in opts:
        if opt in {"-h", "--help"}:
            print(help)
            sys.exit()
        elif opt in {"-b", "--bibfile"}:
            bib_file_name = arg
        elif opt in {"-o", "--output"}:
            output_file_name = arg
        elif opt in {"-a", "--author"}:
            faname = arg

    # Set the formatting identifiers. Since we're using kramdown, we
    # don't have to use the HTML tags.
    em = '_'
    st = '**'

    # Set some HTML tags to go around each part of the reference.
    openSpan = ''
    closeSpan = ''

    # Open and parse the BibTeX file in `bib_file_name` using
    # `bibtexparser`

    with open(bib_file_name,'r') as bib_file:
        bp = BibTexParser(bib_file)

    # Get a dictionary of dictionaries of key, value pairs from the
    # BibTeX file. The structure is
    # {ID:{authors:...},ID:{authors:...}}.
    refsdict = bp.get_entry_dict()

    # Create a list of all the types of documents found in the BibTeX
    # file, typically `article`, `inproceedings`, and `phdthesis`.
    # Dedupe the list.
    types = []
    for k,ref in refsdict.items():
        types.append(ref["type"])
    types  = set(types)

    # For each of the types of reference, we need to sort each by month
    # then year. We store the dictionary representing each reference in
    # a sorted list for each type of reference. Then we store each of
    # these sorted lists in a dictionary whose key is the type of
    # reference and value is the list of dictionaries.
    sort_dict = {}
    for type in types:
        temp = sorted([val for key, val in refsdict.items()
                      if val["type"] == type], key=lambda l:
                      datetime.strptime(l["month"],'%b').month, reverse=True)
        sort_dict[type] = sorted(temp, key=lambda k: k["year"], reverse=True)

    # Open the output file with utf-8 encoding, write mode, and Unix
    # newlines.
    with open(output_file_name, encoding='utf-8', mode='w',
              newline='') as out_file:

        # Start with journal articles. Print the header to the screen
        # and output file.
        print('Journal Articles\n---\n')
        out_file.write('Journal Articles\n---\n')

        # To get the year numbering correct, we have to set a dummy
        # value for pubyear (usage described below).
        pubyear = ''

        # Loop through all the references in the article type. The
        # logic in this loop (and the loops for the other reference
        # types) is not amenable to generalization due to different
        # information for each reference type. Therefore, its easiest
        # to write out the logic for each loop instead of writing the
        # logic into a function and calling that.
        for ref in sort_dict["article"]:
            # Get the string of author names in the proper format from
            # the `reorder` function. Get some other information. Hack
            # the journal title to remove the '\' before '&' in
            # 'Energy & Fuels' because Mendeley inserts an extra '\'
            # into the BibTeX.
            authors = reorder(ref["author"],faname)
            title = ref["title"]
            journal = ref["journal"]
            if '\&' in journal:
                words = journal.strip().split('\&')
                journal = words[0] + '&' + words[1]
            doi = ref["doi"]

            # Get the publication year. If the year of the current
            # reference is not equal to the year of the previous
            # reference, we need to print the year out and set
            # `pubyear` equal to `year`.
            year = ref["year"]
            if year != pubyear:
                pubyear = year
                stri = '\n{:.year}\n###' + year + '\n'
                print(stri)
                out_file.write(stri)

            # Start building the string containing the formatted
            # reference. Each bit should be surrounded by a span. The
            # {:.XYZ} is the kramdown notation to add that class to the
            # HTML tag. Each line should be ended with two spaces
            # before the newline so that kramdown inserts an HTML

            # there.
            string = ('\n{:.paper}\n' +
                      openSpan + title + closeSpan + '{:.papertitle}  \n' +
                      openSpan + authors + closeSpan + '{:.authors}  \n' +
                      openSpan + em + journal + em + ', '
                      )

            # Not all journal articles will have vol., no., and pp.
            # because some may be "In Press".
            if "volume" in ref:
                string = string + 'vol. ' + ref["volume"] + ', '

            if "number" in ref:
                string = string + 'no. ' + ref["number"] + ', '

            if "pages" in ref:
                string = string + 'pp. ' + ref["pages"] + ', '

            string = (string + ref["month"].title() + '. ' +
                      year + closeSpan + '{:.journal}  \n')

            if "doi" in ref:
                string = (string + openSpan + st + 'DOI:' + st + ' [' +
                          ref["doi"] + '](http://dx.doi.org/' + ref["doi"]
                          + ')' + closeSpan + '{:.doi}  \n'
                          )

            # Extra comments, such as links to files, should be stored
            # as "Notes" for each reference in Mendeley. Mendeley will
            # export this field with the tag "annote" in BibTeX.
            if "annote" in ref:
                string = (string + openSpan + ref["annote"].replace('\\','') +
                           closeSpan + '{:.comment}  \n'
                         )
            print(string)
            out_file.write(string)

        # Next are conference papers and posters. Print the header to
        # the screen and the output file.
        print('Conference Publications and Posters\n---\n')
        out_file.write('\nConference Publications and Posters\n---\n')

        # Same trick for the pubyear as for the journal articles.
        pubyear = ''

        # Loop through the references in the `inproceedings` type.
        for ref in sort_dict["inproceedings"]:
            authors = reorder(ref["author"],faname)
            title = ref["title"]
            year = ref["year"]
            if year != pubyear:
                pubyear = year
                stri = '\n{:.year}\n###' + year + '\n'
                print(stri)
                out_file.write(stri)

            # Start building the reference string.
            string = ('\n{:.paper}\n' +
                      openSpan + title + closeSpan + '{:.papertitle}  \n' +
                      openSpan + authors + closeSpan + '{:.authors}  \n' +
                      openSpan
                     )

            # Since Mendeley doesn't allow customization of BibTeX
            # output, we hack the "pages" field to contain the paper
            # number for the conference paper. Not all of this type of
            # reference will have this, so we check for it.
            if "pages" in ref:
                paperno = ref["pages"]
                string = string + paperno + ', '

            # Insert the conference title, stored in the "booktitle"
            # field.
            conf = ref["booktitle"]
            string = string + conf + ', '
            if "organization" in ref:
                string = string + ref["organization"] + ', '
            if "address" in ref:
                string = string + ref["address"] + ', '

            string = (string + ref["month"].title() + '. ' + year + closeSpan
                      + '{:.journal}  \n'
                     )

            if "doi" in ref:
                string = (string + openSpan + st + 'DOI:' + st + ' [' +
                          ref["doi"] + '](http://dx.doi.org/' + ref["doi"]
                          + ')' + closeSpan + '{:.doi}  \n'
                         )

            # Extra comments, such as links to files, should be stored
            # as "Notes" for each reference in Mendeley. Mendeley will
            # export this field with the tag "annote" in BibTeX.
            if "annote" in ref:
                string = (string + openSpan + ref["annote"].replace('\\','') +
                          closeSpan + '{:.comment}  \n'
                         )
            print(string)
            out_file.write(string)

        # Finally are the theses and dissertations. Same general logic
        # as for the other reference types.
        print("Master's Thesis\n---\n")
        out_file.write("\nMaster's Thesis\n---\n")
        pubyear = '2200'
        for ref in sort_dict["phdthesis"]:
            authors = reorder(ref["author"],faname)
            title = ref["title"]
            year = ref["year"]
            if year != pubyear:
                pubyear = year
                stri = '\n{:.year}\n###' + year + '\n'
                print(stri)
                out_file.write(stri)

            string = ('\n{:.paper}\n' +
                      openSpan + title + closeSpan + '{:.papertitle}  \n' +
                      openSpan + authors + closeSpan + '{:.authors}  \n' +
                      openSpan )
            if "school" in ref:
                string = string + ref["school"] + ', '
            if "month" in ref:
                string = string + ref["month"].title() + '. '

            string = string + year + closeSpan + '{:.journal}  \n'

            if "annote" in ref:
                string = (string + openSpan + ref["annote"].replace('\\','') +
                          closeSpan + '{:.comment}  \n'
                         )
            print(string)
            out_file.write(string)

if __name__ == "__main__":
    main(sys.argv[1:])
<|endoftext|>


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

import json
import click
import os
import config
import sqlite3
import time

from flask import Flask
from flask import request
from flask import g
from flask import redirect
from flask import abort
from flask import render_template
from flask import flash
from werkzeug.exceptions import HTTPException

app = Flask(__name__)

# Load default config and override config from an environment variable
app.config.update(dict(
    DATABASE=os.path.join(app.root_path, config.DATABASE_PATH),
    DEBUG=True,
))

class JSONException(HTTPException):
    def __init__(self, message, status_code=None, response=None):
        self.message = message
        self.response = response

        if status_code is not None:
            self.code = status_code
        else:
            self.code = 400

    def get_body(self, environ):
        return json.dumps({
            "error": self.message,
            "code": self.code
        })

    def get_headers(self, environ):
        return [('Content-Type', 'application/json')]

ALPHABET = "23456789abcdefghijkmnpqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"

def base62_encode(num, alphabet=ALPHABET):
    """Encode a number in Base X

    `num`: The number to encode
    `alphabet`: The alphabet to use for encoding
    """
    if (num == 0):
        return alphabet[0]
    arr = []
    base = len(alphabet)
    while num:
        rem = num % base
        num = num // base
        arr.append(alphabet[rem])
    arr.reverse()
    return ''.join(arr)

def base62_decode(string, alphabet=ALPHABET):
    """Decode a Base X encoded string into the number

    Arguments:
    - `string`: The encoded string
    - `alphabet`: The alphabet to use for encoding
    """
    base = len(alphabet)
    strlen = len(string)
    num = 0

    idx = 0
    for char in string:
        power = (strlen - (idx + 1))
        num += alphabet.index(char) * (base ** power)
        idx += 1

    return num

def connect_db():
    """Connects to the specific database."""
    rv = sqlite3.connect(app.config['DATABASE'])
    rv.row_factory = sqlite3.Row
    return rv

def init_db():
    """Initializes the database."""
    db = get_db()
    with app.open_resource('schema/v1.sql', mode='r') as f:
        db.cursor().executescript(f.read())
    db.commit()

def get_db():
    """Opens a new database connection if there is none yet for the
    current application context.
    """
    if not hasattr(g, 'sqlite_db'):
        g.sqlite_db = connect_db()
    return g.sqlite_db

def lookup_url(link_id):
    db = get_db()
    cur = db.execute('SELECT url FROM urls WHERE key = ?', (link_id,))
    link = cur.fetchall()
    if len(link) is not 1:
        return None
    else:
        return link

def save_url(url, wish=None):
    exists = None
    db = get_db()

    if wish is not None:
        exists = lookup_url(wish)
        if exists is not None:
            return wish
        else:
            db.execute('INSERT INTO urls (key, url) VALUES (?, ?)', (wish, url))
            db.commit()
            return wish
    else:
        try:
            cur = db.execute('SELECT key FROM urls WHERE url = ?', (url,))
            key_exists = cur.fetchone()[0]
            if key_exists is not None:
                return key_exists
        except TypeError:
            key_exists = None

    cur = db.execute('SELECT key FROM urls ORDER BY key LIMIT 1')
    last_key = cur.fetchone()

    if not last_key or last_key == '':
        key = base62_encode(8)
    else:
        key = base62_encode(base62_decode(last_key) + 1)

    db.execute('INSERT INTO urls (key, url) VALUES (?, ?)', (key, url))
    db.commit()
    return key

@app.cli.command('initdb')
def initdb_command():
    """Creates the database tables."""
    init_db()
    print('Initialized the database.')

@app.cli.command('addkey')
@click.option('--key', help='The API key to add')
@click.option('--limit', default=10000, help='Maximum requests per day')
def add_api_key(key, limit):
    """Adds an API key for authorization."""
    if key is None or len(key) is not 32:
        print('Keys must be exactly 32 characters long!')
        return

    if limit == 0:
        print('INFO: Limit set to 0 - setting no limit.')

    if limit < 0:
        print('INFO: Limit is less than 0, disabling account')

    db = get_db()
    db.execute('INSERT OR REPLACE INTO api (key, dlimit) VALUES (?, ?)', (key, limit))
    db.commit()

    print('Key "%s" added to the database.' % key)

@app.teardown_appcontext
def close_db(error):
    """Closes the database again at the end of the request."""
    if hasattr(g, 'sqlite_db'):
        g.sqlite_db.close()

@app.route('/')
def index():
    return render_template('index.html')

@app.errorhandler(404)
def page_not_found(error):
    return render_template('404.html'), 404

@app.route('/<link_id>')
def short_link(link_id):
    url = lookup_url(link_id)
    if url is None:
        abort(404)
    else:
        db = get_db()
        db.execute('INSERT INTO stats (link_id, time) VALUES(?, ?)',
                   (link_id, int(time.time()),))
        return redirect(url)

@app.route('/save', methods=['POST'])
def save_link():
    if not len(request.form['url']):
        flash('No URL supplied')
        return redirect('/')

    key = save_url(request.form['url'])

    if key is None:
        abort(500)

    return redirect('/' + key + '+')

@app.route('/<link_id>+')
def link_info(link_id):
    link_url = lookup_url(link_id)
    if link_url is None:
        abort(404)

    db = get_db()
    link_stats = db.execute('SELECT time FROM stats WHERE link_id = ?',
                            (link_id,))

    return json.dumps({'Link:': link_url, 'Stats:': link_stats})

@app.route('/api/v1/short', methods=['POST'])
def api_v1_short():
    if not request.json:
        raise JSONException(message="No data supplied", status_code=400)

    # TODO: api key validation
    if 'key' not in request.json:
        raise JSONException(message="No data supplied", status_code=401)

    # TODO: url validation
    if 'url' not in request.json:
        raise JSONException(message="No URL supplied", status_code=400)

    try:
        short_link = save_url(request.json['url'], request.json['wish'])
        return json.dumps({
            "url_long": request.json['url'],
            "url_short": short_link,
            "wish": request.json['wish'],
        })
    except Exception:
        return JSONException(message="Internal server error", status_code=500)

@app.route('/api/v1/long', methods=['POST'])
def api_v1_long():
    return json.dumps(request)

if __name__ == "__main__":
    app.run(host='localhost', port=9002)
<commit_msg>Added upgrade option stub to initdb CLI command<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import json
import click
import os
import config
import sqlite3
import time

from flask import Flask
from flask import request
from flask import g
from flask import redirect
from flask import abort
from flask import render_template
from flask import flash
from werkzeug.exceptions import HTTPException

app = Flask(__name__)

# Load default config and override config from an environment variable
app.config.update(dict(
    DATABASE=os.path.join(app.root_path, config.DATABASE_PATH),
    DEBUG=True,
))

class JSONException(HTTPException):
    def __init__(self, message, status_code=None, response=None):
        self.message = message
        self.response = response

        if status_code is not None:
            self.code = status_code
        else:
            self.code = 400

    def get_body(self, environ):
        return json.dumps({
            "error": self.message,
            "code": self.code
        })

    def get_headers(self, environ):
        return [('Content-Type', 'application/json')]

ALPHABET = "23456789abcdefghijkmnpqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"

def base62_encode(num, alphabet=ALPHABET):
    """Encode a number in Base X

    `num`: The number to encode
    `alphabet`: The alphabet to use for encoding
    """
    if (num == 0):
        return alphabet[0]
    arr = []
    base = len(alphabet)
    while num:
        rem = num % base
        num = num // base
        arr.append(alphabet[rem])
    arr.reverse()
    return ''.join(arr)

def base62_decode(string, alphabet=ALPHABET):
    """Decode a Base X encoded string into the number

    Arguments:
    - `string`: The encoded string
    - `alphabet`: The alphabet to use for encoding
    """
    base = len(alphabet)
    strlen = len(string)
    num = 0

    idx = 0
    for char in string:
        power = (strlen - (idx + 1))
        num += alphabet.index(char) * (base ** power)
        idx += 1

    return num

def connect_db():
    """Connects to the specific database."""
    rv = sqlite3.connect(app.config['DATABASE'])
    rv.row_factory = sqlite3.Row
    return rv

def init_db():
    """Initializes the database."""
    db = get_db()
    with app.open_resource('schema/v1.sql', mode='r') as f:
        db.cursor().executescript(f.read())
    db.commit()

def get_db():
    """Opens a new database connection if there is none yet for the
    current application context.
    """
    if not hasattr(g, 'sqlite_db'):
        g.sqlite_db = connect_db()
    return g.sqlite_db

def lookup_url(link_id):
    db = get_db()
    cur = db.execute('SELECT url FROM urls WHERE key = ?', (link_id,))
    link = cur.fetchall()
    if len(link) is not 1:
        return None
    else:
        return link

def save_url(url, wish=None):
    exists = None
    db = get_db()

    if wish is not None:
        exists = lookup_url(wish)
        if exists is not None:
            return wish
        else:
            db.execute('INSERT INTO urls (key, url) VALUES (?, ?)', (wish, url))
            db.commit()
            return wish
    else:
        try:
            cur = db.execute('SELECT key FROM urls WHERE url = ?', (url,))
            key_exists = cur.fetchone()[0]
            if key_exists is not None:
                return key_exists
        except TypeError:
            key_exists = None

    cur = db.execute('SELECT key FROM urls ORDER BY key LIMIT 1')
    last_key = cur.fetchone()

    if not last_key or last_key == '':
        key = base62_encode(8)
    else:
        key = base62_encode(base62_decode(last_key) + 1)

    db.execute('INSERT INTO urls (key, url) VALUES (?, ?)', (key, url))
    db.commit()
    return key

@app.cli.command('initdb')
@click.option('--upgrade', default='no', 'Only upgrade to a newer version')
def initdb_command(upgrade):
    if upgrade == 'yes':
        print('Not implemented')
        return
    else:
        """Creates the database tables."""
        init_db()
        print('Initialized the database.')

@app.cli.command('addkey')
@click.option('--key', help='The API key to add')
@click.option('--limit', default=10000, help='Maximum requests per day')
def add_api_key(key, limit):
    """Adds an API key for authorization."""
    if key is None or len(key) is not 32:
        print('Keys must be exactly 32 characters long!')
        return

    if limit == 0:
        print('INFO: Limit set to 0 - setting no limit.')

    if limit < 0:
        print('INFO: Limit is less than 0, disabling account')

    db = get_db()
    db.execute('INSERT OR REPLACE INTO api (key, dlimit) VALUES (?, ?)', (key, limit))
    db.commit()

    print('Key "%s" added to the database.' % key)

@app.teardown_appcontext
def close_db(error):
    """Closes the database again at the end of the request."""
    if hasattr(g, 'sqlite_db'):
        g.sqlite_db.close()

@app.route('/')
def index():
    return render_template('index.html')

@app.errorhandler(404)
def page_not_found(error):
    return render_template('404.html'), 404

@app.route('/<link_id>')
def short_link(link_id):
    url = lookup_url(link_id)
    if url is None:
        abort(404)
    else:
        db = get_db()
        db.execute('INSERT INTO stats (link_id, time) VALUES(?, ?)',
                   (link_id, int(time.time()),))
        return redirect(url)

@app.route('/save', methods=['POST'])
def save_link():
    if not len(request.form['url']):
        flash('No URL supplied')
        return redirect('/')

    key = save_url(request.form['url'])

    if key is None:
        abort(500)

    return redirect('/' + key + '+')

@app.route('/<link_id>+')
def link_info(link_id):
    link_url = lookup_url(link_id)
    if link_url is None:
        abort(404)

    db = get_db()
    link_stats = db.execute('SELECT time FROM stats WHERE link_id = ?',
                            (link_id,))

    return json.dumps({'Link:': link_url, 'Stats:': link_stats})

@app.route('/api/v1/short', methods=['POST'])
def api_v1_short():
    if not request.json:
        raise JSONException(message="No data supplied", status_code=400)

    # TODO: api key validation
    if 'key' not in request.json:
        raise JSONException(message="No data supplied", status_code=401)

    # TODO: url validation
    if 'url' not in request.json:
        raise JSONException(message="No URL supplied", status_code=400)

    try:
        short_link = save_url(request.json['url'], request.json['wish'])
        return json.dumps({
            "url_long": request.json['url'],
            "url_short": short_link,
            "wish": request.json['wish'],
        })
    except Exception:
        return JSONException(message="Internal server error", status_code=500)

@app.route('/api/v1/long', methods=['POST'])
def api_v1_long():
    return json.dumps(request)

if __name__ == "__main__":
    app.run(host='localhost', port=9002)
<|endoftext|>


<commit_before>#!/usr/bin/env python
import os
import sys
import paramiko
import subprocess

# definitions
bli_path = "/".join(sys.argv[0].split("/")[0:-1])
policies = ["phys"]

# defaults
sensor_file = "sensor.csv"
ssh_user    = "broadmin"
prefix      = "/opt/bro"
spooltmp    = "/data/bro/spool/tmp"
policy_type = "phys"

def print_usage():
	print "Usage: bli.py [OPTION]"
	print "Options:"
	print "  {0:20s} print downstream health".format("status")
	print "  {0:20s} print downstream config".format("config")
	print "  {0:20s} clear crash logs".format("clear_logs")
	print "  {0:20s} check policy".format("check_policy")
	print "  {0:20s} give this help list".format("-?, --help")

def menu():
	print "\n", "-"*30
	print "|{0:^28s}|".format("Brommand Line Interface")
	print "-"*30
	print "|{0:28s}|".format(" (1) Get status")
	print "|{0:28s}|".format(" (2) Print status")
	print "|{0:28s}|".format(" (3) Clear crash logs")
	print "|{0:28s}|".format(" (4) Check policy")
	print "|{0:28s}|".format(" (9) Print configuration")
	print "|{0:28s}|".format(" (0) Quit")
	print "-"*30

def get_sensors(sensors):
	sensor_list = open(os.path.join(bli_path, sensor_file), "r")

	for sensor in sensor_list:
		populate_sensors(sensor, sensors)

def populate_sensors(sensor, sensors):
	if sensor:
		line = sensor.split(",")

		if line[0][0] != "#":
			line[-1] = line[-1].strip()
			sensors[line[0]] = {}
			sensors[line[0]]['crash_logs']   = 0
			sensors[line[0]]['policy_file'] = {}
			sensors[line[0]]['hostname']    = line[1]
			sensors[line[0]]['ssh_user']    = line[2] if line[2] else ssh_user
			sensors[line[0]]['prefix']      = line[3] if line[3] else prefix
			sensors[line[0]]['spooltmp']    = line[4] if line[4] else spooltmp
			sensors[line[0]]['policy_type'] = line[5] if line[5] else policy_type

def print_config(sensors):
	print "\n{0:15s} : {1:20s} : {2:10s} : {3:20s} : {4:20s} : {5}".format("IP Address", "Hostname", "User", "Prefix", "SpoolTmp", "Policy")
	print "-"*120

	for sensor in sorted(sensors):
		print "{0:15s} : {1:20s} : {2:10s} : {3:20s} : {4:20s} : {5}".format(sensor, sensors[sensor]['hostname'], sensors[sensor]['ssh_user'], sensors[sensor]['prefix'], sensors[sensor]['spooltmp'], sensors[sensor]['policy_type'])

def get_status(sensors):
	for sensor in sensors:
		ssh = paramiko.SSHClient()
		ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

		try:
			ssh.connect(
				sensor,
				username = sensors[sensor]['ssh_user'],
				key_filename = os.path.expanduser(os.path.join("~", ".ssh", "id_rsa.pub")),
				timeout = 10
			)
			lines = running = stopped = crashed = warnings = fnf_prefix = fnf_spool = 0

			(stdin, stdout, stderr) = ssh.exec_command("ls " + sensors[sensor]['spooltmp'] + "|grep crash |wc -l")
			for line in stderr.readlines():
				if "No such file" in line:
					fnf_spool = 1
			if not fnf_spool:
				sensors[sensor]['crash_logs'] = int(stdout.readline().strip())

			(stdin, stdout, stderr) = ssh.exec_command(os.path.join(sensors[sensor]['prefix'], "bin", "broctl") + " status 2>&1")
			for line in stdout.readlines():
				if "warning" in line:
					warnings += 1
				if "No such file" in line:
					fnf_prefix = 1
				if "manager" in line or "proxy" in line or "worker" in line or "standalone" in line:
					lines += 1
					if "running" in line:
						running += 1
					elif "stopped" in line:
						stopped += 1
					elif "crashed" in line:
						crashed += 1

			(stdin, stdout, stderr) = ssh.exec_command("find " +  os.path.join(sensors[sensor]['prefix'], "share", "bro", "site", "*") + " -exec md5sum '{}' \;")
			for line in stdout.readlines():
				line = line.strip().split()
				line[1] = line[1].split(os.path.join(sensors[sensor]['prefix'], "share", "bro", "site/"))
				sensors[sensor]['policy_file'][line[1][-1]] = line[0]

			if not fnf_prefix and not fnf_spool:
				if running == lines:
					sensors[sensor]['status'] = "OK (" + str(warnings) + " warnings, " + str(sensors[sensor]['crash_logs']) + " crash logs)"
				else:
					sensors[sensor]['status'] = "Unhealthy (" + str(running) + " running, " + str(stopped) + " stopped, " + str(crashed) + " crashed, " + str(warnings) + " warnings, " + str(sensors[sensor]['crash_logs']) + " crash logs)"
			elif fnf_prefix and not fnf_spool:
				sensors[sensor]['status'] = "Error (broctl not found; validate prefix setting)"
			elif fnf_spool and not fnf_prefix:
				sensors[sensor]['status'] = "Error (Bro spool not found; validate spooltmp setting)"
			elif fnf_prefix and fnf_spool:
				sensors[sensor]['status'] = "Error (broctl and spool not found; validate path settings)"
			ssh.close()
		except Exception as e:
			sensors[sensor]['status'] = "Error (" + str(e) + ")"

	print "\nStatus loaded..."
	return 1

def print_status(sensors):
	print "\n{0:15s} : {1:20s} : {2}".format("IP Address", "Hostname", "Status")
	print "-"*120

	for sensor in sorted(sensors):
		print "{0:15s} : {1:20s} : {2}".format(sensor, sensors[sensor]['hostname'], sensors[sensor]['status'])

def clear_logs(sensors):
	cleared = 0

	for sensor in sensors:
		if sensors[sensor]['crash_logs'] > 0:
			ssh = paramiko.SSHClient()
			ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

			try:
				ssh.connect(
					sensor,
					username = sensors[sensor]['ssh_user'],
					key_filename = os.path.expanduser(os.path.join("~", ".ssh", "id_rsa.pub")),
					timeout = 10
				)
				(stdin, stdout, stderr) = ssh.exec_command("rm -rf " + sensors[sensor]['spooltmp'] + "/*crash")
				ssh.close()
				cleared += 1
				print "\nLogs cleared from", sensor + "..."
			except Exception as e:
				sensors[sensor]['status'] = "Error (" + str(e) + ")"

	if cleared == 0:
		print "\nNo log(s) cleared..."

def check_policy(sensors):
	policy = {}
	for pol in policies:
		policy[pol] = {}
		stdout = subprocess.Popen(
			["find", os.path.join(bli_path, "deploy", pol, "site"), "-exec", "md5sum", "{}", ";"],
			stdout = subprocess.PIPE,
			stderr = subprocess.STDOUT,
			universal_newlines = True
			).stdout

		for line in stdout.readlines():
			line = line.strip().split()
			if len(line) == 2:
				line[1] = line[1].split(os.path.join(bli_path, "deploy", pol, "site/"))[1]
				policy[pol][line[1]] = line[0]

		print "\n{0:15s} : {1:20s} : {2:6s} : {3:8s} : {4}".format("IP Address", "Hostname", "Policy", "Issue", "File")
		print "-"*120

		for sensor in sorted(sensors):
			if "Error" not in sensors[sensor]['status']:
				print "{0:15s} : {1:20s} : {2:6s} : {3:8s} :".format(sensor, sensors[sensor]['hostname'], sensors[sensor]['policy_type'], "")
				for policy_file in policy[sensors[sensor]['policy_type']]:
					if policy_file in sensors[sensor]['policy_file']:
						if policy[sensors[sensor]['policy_type']][policy_file] != sensors[sensor]['policy_file'][policy_file]:
							print "{0:15s} : {1:20s} : {2:7s}: {3:8s} : {4} ".format("", "", "", "modified", policy_file)
					else:
						print "{0:15s} : {1:20s} : {2:7s} : {3:8s} : {4} ".format("", "", "", "missing", policy_file)

def main():
	loaded = decision = None
	sensors = {}
	get_sensors(sensors)

	if len(sys.argv) == 1:
		while decision != 0:
			menu()

			try:
				decision = input("\naction> ")
			except Exception as e:
				decision = None

			if decision == 1:
				loaded = get_status(sensors)
				raw_input("\n<Press Enter to continue>")
			elif decision == 9:
				print_config(sensors)
				raw_input("\n<Press Enter to continue>")
			elif decision == 0:
				print "\nExiting..."
			else:
				if loaded:
					if decision == 2:
						print_status(sensors)
						raw_input("\n<Press Enter to continue>")
					elif decision == 3:
						clear_logs(sensors)
						raw_input("\n<Press Enter to continue>")
					elif decision == 4:
						check_policy(sensors)
						raw_input("\n<Press Enter to continue>")
				else:
					print "\nStatus not yet loaded (get status)"

	elif len(sys.argv) == 2:
		if sys.argv[1] == "status":
			get_status(sensors)
			print_status(sensors)
		elif sys.argv[1] == "config":
			print_config(sensors)
		elif sys.argv[1] == "clear_logs":
			get_status(sensors)
			clear_logs(sensors)
		elif sys.argv[1] == "check_policy":
			get_status(sensors)
			check_policy(sensors)
		elif sys.argv[1] == "-?" or sys.argv[1] == "--help":
			print_usage()
		else:
			print sys.argv[0] + ": invalid option:", sys.argv[1]
			print "Try '" + sys.argv[0] + " --help' for more information."

	else:
		print sys.argv[0] + ": too many arguments"
		print "Try '" + sys.argv[0] + " --help' for more information."

######
main()
<commit_msg>Initial check_policy directory structure validation<commit_after>#!/usr/bin/env python
import os
import sys
import paramiko
import subprocess

# definitions
bli_path = "/".join(sys.argv[0].split("/")[0:-1])
policies = ["phys"]

# defaults
sensor_file = "sensor.csv"
ssh_user    = "broadmin"
prefix      = "/opt/bro"
spooltmp    = "/data/bro/spool/tmp"
policy_type = "phys"

def print_usage():
	print "Usage: bli.py [OPTION]"
	print "Options:"
	print "  {0:20s} print downstream health".format("status")
	print "  {0:20s} print downstream config".format("config")
	print "  {0:20s} clear crash logs".format("clear_logs")
	print "  {0:20s} check policy".format("check_policy")
	print "  {0:20s} give this help list".format("-?, --help")

def menu():
	print "\n", "-"*30
	print "|{0:^28s}|".format("Brommand Line Interface")
	print "-"*30
	print "|{0:28s}|".format(" (1) Get status")
	print "|{0:28s}|".format(" (2) Print status")
	print "|{0:28s}|".format(" (3) Clear crash logs")
	print "|{0:28s}|".format(" (4) Check policy")
	print "|{0:28s}|".format(" (9) Print configuration")
	print "|{0:28s}|".format(" (0) Quit")
	print "-"*30

def get_sensors(sensors):
	sensor_list = open(os.path.join(bli_path, sensor_file), "r")

	for sensor in sensor_list:
		populate_sensors(sensor, sensors)

def populate_sensors(sensor, sensors):
	if sensor:
		line = sensor.split(",")

		if line[0][0] != "#":
			line[-1] = line[-1].strip()
			sensors[line[0]] = {}
			sensors[line[0]]['crash_logs']   = 0
			sensors[line[0]]['policy_file'] = {}
			sensors[line[0]]['hostname']    = line[1]
			sensors[line[0]]['ssh_user']    = line[2] if line[2] else ssh_user
			sensors[line[0]]['prefix']      = line[3] if line[3] else prefix
			sensors[line[0]]['spooltmp']    = line[4] if line[4] else spooltmp
			sensors[line[0]]['policy_type'] = line[5] if line[5] else policy_type

def print_config(sensors):
	print "\n{0:15s} : {1:20s} : {2:10s} : {3:20s} : {4:20s} : {5}".format("IP Address", "Hostname", "User", "Prefix", "SpoolTmp", "Policy")
	print "-"*120

	for sensor in sorted(sensors):
		print "{0:15s} : {1:20s} : {2:10s} : {3:20s} : {4:20s} : {5}".format(sensor, sensors[sensor]['hostname'], sensors[sensor]['ssh_user'], sensors[sensor]['prefix'], sensors[sensor]['spooltmp'], sensors[sensor]['policy_type'])

def get_status(sensors):
	for sensor in sensors:
		ssh = paramiko.SSHClient()
		ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

		try:
			ssh.connect(
				sensor,
				username = sensors[sensor]['ssh_user'],
				key_filename = os.path.expanduser(os.path.join("~", ".ssh", "id_rsa.pub")),
				timeout = 10
			)
			lines = running = stopped = crashed = warnings = fnf_prefix = fnf_spool = 0

			(stdin, stdout, stderr) = ssh.exec_command("ls " + sensors[sensor]['spooltmp'] + "|grep crash |wc -l")
			for line in stderr.readlines():
				if "No such file" in line:
					fnf_spool = 1
			if not fnf_spool:
				sensors[sensor]['crash_logs'] = int(stdout.readline().strip())

			(stdin, stdout, stderr) = ssh.exec_command(os.path.join(sensors[sensor]['prefix'], "bin", "broctl") + " status 2>&1")
			for line in stdout.readlines():
				if "warning" in line:
					warnings += 1
				if "No such file" in line:
					fnf_prefix = 1
				if "manager" in line or "proxy" in line or "worker" in line or "standalone" in line:
					lines += 1
					if "running" in line:
						running += 1
					elif "stopped" in line:
						stopped += 1
					elif "crashed" in line:
						crashed += 1

			(stdin, stdout, stderr) = ssh.exec_command("find " +  os.path.join(sensors[sensor]['prefix'], "share", "bro", "site", "*") + " -exec md5sum '{}' \;")
			for line in stdout.readlines():
				line = line.strip().split()
				line[1] = line[1].split(os.path.join(sensors[sensor]['prefix'], "share", "bro", "site/"))
				sensors[sensor]['policy_file'][line[1][-1]] = line[0]

			if not fnf_prefix and not fnf_spool:
				if running == lines:
					sensors[sensor]['status'] = "OK (" + str(warnings) + " warnings, " + str(sensors[sensor]['crash_logs']) + " crash logs)"
				else:
					sensors[sensor]['status'] = "Unhealthy (" + str(running) + " running, " + str(stopped) + " stopped, " + str(crashed) + " crashed, " + str(warnings) + " warnings, " + str(sensors[sensor]['crash_logs']) + " crash logs)"
			elif fnf_prefix and not fnf_spool:
				sensors[sensor]['status'] = "Error (broctl not found; validate prefix setting)"
			elif fnf_spool and not fnf_prefix:
				sensors[sensor]['status'] = "Error (Bro spool not found; validate spooltmp setting)"
			elif fnf_prefix and fnf_spool:
				sensors[sensor]['status'] = "Error (broctl and spool not found; validate path settings)"
			ssh.close()
		except Exception as e:
			sensors[sensor]['status'] = "Error (" + str(e) + ")"

	print "\nStatus loaded..."
	return 1

def print_status(sensors):
	print "\n{0:15s} : {1:20s} : {2}".format("IP Address", "Hostname", "Status")
	print "-"*120

	for sensor in sorted(sensors):
		print "{0:15s} : {1:20s} : {2}".format(sensor, sensors[sensor]['hostname'], sensors[sensor]['status'])

def clear_logs(sensors):
	cleared = 0

	for sensor in sensors:
		if sensors[sensor]['crash_logs'] > 0:
			ssh = paramiko.SSHClient()
			ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

			try:
				ssh.connect(
					sensor,
					username = sensors[sensor]['ssh_user'],
					key_filename = os.path.expanduser(os.path.join("~", ".ssh", "id_rsa.pub")),
					timeout = 10
				)
				(stdin, stdout, stderr) = ssh.exec_command("rm -rf " + sensors[sensor]['spooltmp'] + "/*crash")
				ssh.close()
				cleared += 1
				print "\nLogs cleared from", sensor + "..."
			except Exception as e:
				sensors[sensor]['status'] = "Error (" + str(e) + ")"

	if cleared == 0:
		print "\nNo log(s) cleared..."

def check_policy(sensors):
	stdout = subprocess.Popen(
		["find", os.path.join(bli_path, "deploy")],
		stdout = subprocess.PIPE,
		stderr = subprocess.STDOUT,
		universal_newlines = True
		).stdout

	dnf = False
	for line in stdout.readlines():
		if "No such file" in line:
			dnf = True

	if dnf == False:
		policy = {}
		for pol in policies:
			policy[pol] = {}
			stdout = subprocess.Popen(
				["find", os.path.join(bli_path, "deploy", pol, "site"), "-exec", 	"md5sum", "{}", ";"],
				stdout = subprocess.PIPE,
				stderr = subprocess.STDOUT,
				universal_newlines = True
				).stdout

			for line in stdout.readlines():
				line = line.strip().split()
				if len(line) == 2:
					line[1] = line[1].split(os.path.join(bli_path, "deploy", 	pol, "site/"))[1]
					policy[pol][line[1]] = line[0]

				print "\n{0:15s} : {1:20s} : {2:6s} : {3:8s} : {4}".format("IP Address", "Hostname", "Policy", "Issue", "File")
				print "-"*120

				for sensor in sorted(sensors):
					if "Error" not in sensors[sensor]['status']:
						print "{0:15s} : {1:20s} : {2:6s} : {3:8s} 	:".format(sensor, sensors[sensor]['hostname'], 	sensors[sensor]['policy_type'], "")
						for policy_file in 	policy[sensors[sensor]['policy_type']]:
							if policy_file in 	sensors[sensor]['policy_file']:
								if policy[sensors[sensor]['policy_type']][policy_file] != sensors[sensor]['policy_file'][policy_file]:
									print "{0:15s} : {1:20s} : {2:7s}: {3:8s} : {4} ".format("", "", "", "modified", policy_file)
								else:
									print "{0:15s} : {1:20s} : {2:7s} : {3:8s} : {4} ".format("", "", "", "missing", policy_file)
	else:
		print "\nPolicy validation unavailable due to lack of deployment data"

def main():
	loaded = decision = None
	sensors = {}
	get_sensors(sensors)

	if len(sys.argv) == 1:
		while decision != 0:
			menu()

			try:
				decision = input("\naction> ")
			except Exception as e:
				decision = None

			if decision == 1:
				loaded = get_status(sensors)
				raw_input("\n<Press Enter to continue>")
			elif decision == 9:
				print_config(sensors)
				raw_input("\n<Press Enter to continue>")
			elif decision == 0:
				print "\nExiting..."
			else:
				if loaded:
					if decision == 2:
						print_status(sensors)
						raw_input("\n<Press Enter to continue>")
					elif decision == 3:
						clear_logs(sensors)
						raw_input("\n<Press Enter to continue>")
					elif decision == 4:
						check_policy(sensors)
						raw_input("\n<Press Enter to continue>")
				else:
					print "\nStatus not yet loaded (get status)"

	elif len(sys.argv) == 2:
		if sys.argv[1] == "status":
			get_status(sensors)
			print_status(sensors)
		elif sys.argv[1] == "config":
			print_config(sensors)
		elif sys.argv[1] == "clear_logs":
			get_status(sensors)
			clear_logs(sensors)
		elif sys.argv[1] == "check_policy":
			get_status(sensors)
			check_policy(sensors)
		elif sys.argv[1] == "-?" or sys.argv[1] == "--help":
			print_usage()
		else:
			print sys.argv[0] + ": invalid option:", sys.argv[1]
			print "Try '" + sys.argv[0] + " --help' for more information."

	else:
		print sys.argv[0] + ": too many arguments"
		print "Try '" + sys.argv[0] + " --help' for more information."

######
main()
<|endoftext|>


<commit_before>import time
from enum import IntEnum

from PIL import Image
from pylab import *

num_of_colors = 2
colormap = plt.get_cmap('inferno')
colors = [colormap(k) for k in np.linspace(0.0, 0.8, num_of_colors)]
sec_colors = [colormap(k + 0.1) for k in np.linspace(0.0, 0.8, num_of_colors)]

class Cell(IntEnum):
    EMPTY = 0
    RIGHT = 1
    DOWN = 2
    BOTH = 3

class BML:
    def __init__(self, width, height, p, model):
        self.width = width
        self.height = height
        self.cells = np.zeros((height, width))
        self.model = model
        self.step = 0
        self.velocity = []

        num_of_cars = int(width * height * (p / 2))
        # Get a list of indices
        inds = list(np.ndindex(self.cells.shape))
        # Shuffle the indices in-place
        np.random.shuffle(inds)
        # Access array elements using the indices to do cool stuff
        for i in inds[:num_of_cars]:
            self.cells[i] = Cell.RIGHT

        for i in inds[num_of_cars:2 * num_of_cars]:
            self.cells[i] = Cell.DOWN

        self.total_number = 2 * num_of_cars

    def make_step(self):
        if self.model == 1:
            vel = self.step_right()
            vel += self.step_down()
            self.velocity.append(vel / self.total_number)
        elif self.model == 2:
            self.velocity.append(self.step_all() / self.total_number)
        elif self.model == 3:
            self.step_all_together()  # TODO
        self.step += 1

    def step_down(self):
        d = np.diff(np.r_[self.cells, [self.cells[0, :]]], axis=0)
        down_indexes = np.where(np.logical_and(self.cells == Cell.DOWN, d == Cell.EMPTY - Cell.DOWN))
        empty_indexes = self.down_set(down_indexes)
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.DOWN
        return len(down_indexes[0])

    def step_right(self):
        d = np.diff(np.c_[self.cells, self.cells[:, 0]])
        right_indexes = np.where(np.logical_and(self.cells == Cell.RIGHT, d == Cell.EMPTY - Cell.RIGHT))
        empty_indexes = self.right_set(right_indexes)
        self.cells[right_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.RIGHT
        return len(right_indexes[0])

    def step_all(self):
        dr = np.c_[self.cells[:, -1], self.cells[:, :-1]] - self.cells
        dd = np.r_[[self.cells[-1, :]], self.cells[:-1, :]] - self.cells
        empty_indexes_both = np.transpose(np.where(np.logical_and(self.cells == Cell.EMPTY,
                                                                  np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                                dr == Cell.RIGHT - Cell.EMPTY))))

        if empty_indexes_both.shape[0] > 0:
            np.random.shuffle(empty_indexes_both)
            right_nums = int(empty_indexes_both.shape[0] / 2)
            # @formatter:off
            empty_indexes_right = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                                   np.logical_and(self.cells == Cell.EMPTY,
                                                                   np.logical_and(dr == Cell.RIGHT - Cell.EMPTY,
                                                                                dd != Cell.DOWN - Cell.EMPTY)))),
                                                      empty_indexes_both[:right_nums]]))
            empty_indexes_down = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                              np.logical_and(self.cells == Cell.EMPTY,
                                                              np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                             dr != Cell.RIGHT - Cell.EMPTY)))),
                                                      empty_indexes_both[right_nums:]]))
            # @formatter:on
        else:
            empty_indexes_right = np.where(np.logical_and(self.cells == Cell.EMPTY, dr == Cell.RIGHT - Cell.EMPTY))
            empty_indexes_down = np.where(np.logical_and(self.cells == Cell.EMPTY, dd == Cell.DOWN - Cell.EMPTY))

        right_indexes = self.left_set(empty_indexes_right)
        down_indexes = self.up_set(empty_indexes_down)

        self.cells[right_indexes] = Cell.EMPTY
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes_right] = Cell.RIGHT
        self.cells[empty_indexes_down] = Cell.DOWN

        # if self.step % 100 == 1:
        #    self.save()
        return len(right_indexes[0]) + len(down_indexes[0])

    def step_all_together(self):
        return -1

    def upper(self, i, j):
        return self.height - 1 if i == 0 else i - 1, j

    def left(self, i, j):
        return i, self.width - 1 if j == 0 else j - 1

    def upper_t(self, ind):
        return self.height - 1 if ind[0] == 0 else ind[0] - 1, ind[1]

    def left_t(self, ind):
        return ind[0], self.width - 1 if ind[1] == 0 else ind[1] - 1

    def down_t(self, ind):
        return 0 if ind[0] == self.height - 1 else ind[0] + 1, ind[1]

    def right_t(self, ind):
        return ind[0], 0 if ind[1] == self.width - 1 else ind[1] + 1

    def down_set(self, inds):
        return tuple([[(x + 1) % self.height for x in inds[0]], inds[1]])

    def right_set(self, inds):
        return tuple([inds[0], [(x + 1) % self.width for x in inds[1]]])

    def up_set(self, inds):
        return tuple([[(x - 1) % self.height for x in inds[0]], inds[1]])

    def left_set(self, inds):
        return tuple([inds[0], [(x - 1) % self.width for x in inds[1]]])

    def run(self, steps):
        start = time.perf_counter()
        for i in range(steps):
            self.make_step()
        total_time = time.perf_counter() - start
        print('{} steps for {}x{} map run in {}s. Average of {}s per step'.format(steps, self.height, self.width,
                                                                                total_time, total_time / steps))

    def save(self, pixel_size=10):
        """
        Saves current state as png image
        """
        visual = np.zeros((self.height, self.width, 3)).astype('uint8')
        for i in range(self.height):
            for j in range(self.width):
                visual[i, j, 0] = 0 if self.cells[(i, j)] == Cell.DOWN else 255  # down is blue
                visual[i, j, 1] = 0 if self.cells[(i, j)] != Cell.EMPTY else 255
                visual[i, j, 2] = 0 if self.cells[(i, j)] == Cell.RIGHT else 255  # right are red
        im = Image.fromarray(visual, mode='RGB')
        im = im.resize((self.height * pixel_size, self.width * pixel_size))
        im.save("visual" + str(self.step) + ".png")

    def plot_velocity(self):
        """
        Builds plot of velocity as a function of step number until now
        """
        params = {'axes.labelsize': 8, 'font.size': 8, 'legend.fontsize': 10, 'xtick.labelsize': 10,
                  'ytick.labelsize': 10, 'text.usetex': False, 'figure.figsize': [4.5, 4.5]}
        rcParams.update(params)
        fig = figure()  # no frame
        ax = fig.add_subplot(111)
        ax.plot(self.velocity, linewidth=1, color=colors[1])
        ax.set_ylim(0, 1)
        ax.set_yticks(np.arange(0, 1, 0.2))

        ax.get_xaxis().tick_bottom()
        ax.get_yaxis().tick_left()

        ax.set_xlabel('Step')
        ax.set_ylabel('Velocity')

        ax.grid(axis='y', color="0.9", linestyle='-', linewidth=1)
        ax.set_axisbelow(True)

        plt.tight_layout()

        fig.savefig('velocity.png', dpi=600)
        # print(self.velocity)

automat = BML(521, 523, 0.1, 2)
automat.save()
automat.run(3500)
automat.save()
automat.plot_velocity()
<commit_msg>Add animation output<commit_after>import time
from enum import IntEnum

from PIL import Image
from pylab import *

num_of_colors = 2
colormap = plt.get_cmap('inferno')
colors = [colormap(k) for k in np.linspace(0.0, 0.8, num_of_colors)]
sec_colors = [colormap(k + 0.1) for k in np.linspace(0.0, 0.8, num_of_colors)]

class Cell(IntEnum):
    EMPTY = 0
    RIGHT = 1
    DOWN = 2
    BOTH = 3

class BML:
    def __init__(self, width, height, p, model, animation_ratio=10):
        self.width = width
        self.height = height
        self.cells = np.zeros((height, width))
        self.model = model  # multidimentional, probabilistic, klein bottle, chain
        self.step = 0
        self.velocity = []
        self.animation_ration = animation_ratio

        num_of_cars = int(width * height * (p / 2))
        # Get a list of indices
        inds = list(np.ndindex(self.cells.shape))
        # Shuffle the indices in-place
        np.random.shuffle(inds)
        # Access array elements using the indices to do cool stuff
        for i in inds[:num_of_cars]:
            self.cells[i] = Cell.RIGHT

        for i in inds[num_of_cars:2 * num_of_cars]:
            self.cells[i] = Cell.DOWN

        self.total_number = 2 * num_of_cars

        self.animation = [self.get_image()]

    def make_step(self):
        if self.model == 1:
            vel = self.step_right()
            vel += self.step_down()
            self.velocity.append(vel / self.total_number)
        elif self.model == 2:
            self.velocity.append(self.step_all() / self.total_number)
        elif self.model == 3:
            self.velocity.append(self.step_all_together() / self.total_number)  # TODO
        self.step += 1
        if self.step % self.animation_ration == 0:
            self.animation.append(self.get_image())

    def step_down(self):
        d = np.diff(np.r_[self.cells, [self.cells[0, :]]], axis=0)
        down_indexes = np.where(np.logical_and(self.cells == Cell.DOWN, d == Cell.EMPTY - Cell.DOWN))
        empty_indexes = self.down_set(down_indexes)
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.DOWN
        return len(down_indexes[0])

    def step_right(self):
        d = np.diff(np.c_[self.cells, self.cells[:, 0]])
        right_indexes = np.where(np.logical_and(self.cells == Cell.RIGHT, d == Cell.EMPTY - Cell.RIGHT))
        empty_indexes = self.right_set(right_indexes)
        self.cells[right_indexes] = Cell.EMPTY
        self.cells[empty_indexes] = Cell.RIGHT
        return len(right_indexes[0])

    def step_all(self):
        dr = np.c_[self.cells[:, -1], self.cells[:, :-1]] - self.cells
        dd = np.r_[[self.cells[-1, :]], self.cells[:-1, :]] - self.cells
        empty_indexes_both = np.transpose(np.where(np.logical_and(self.cells == Cell.EMPTY,
                                                                  np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                                dr == Cell.RIGHT - Cell.EMPTY))))

        if empty_indexes_both.shape[0] > 0:
            np.random.shuffle(empty_indexes_both)
            right_nums = int(empty_indexes_both.shape[0] / 2)
            # @formatter:off
            empty_indexes_right = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                                   np.logical_and(self.cells == Cell.EMPTY,
                                                                   np.logical_and(dr == Cell.RIGHT - Cell.EMPTY,
                                                                                dd != Cell.DOWN - Cell.EMPTY)))),
                                                      empty_indexes_both[:right_nums]]))
            empty_indexes_down = tuple(np.transpose(np.r_[
                                                      np.transpose(np.where(
                                                              np.logical_and(self.cells == Cell.EMPTY,
                                                              np.logical_and(dd == Cell.DOWN - Cell.EMPTY,
                                                                             dr != Cell.RIGHT - Cell.EMPTY)))),
                                                      empty_indexes_both[right_nums:]]))
            # @formatter:on
        else:
            empty_indexes_right = np.where(np.logical_and(self.cells == Cell.EMPTY, dr == Cell.RIGHT - Cell.EMPTY))
            empty_indexes_down = np.where(np.logical_and(self.cells == Cell.EMPTY, dd == Cell.DOWN - Cell.EMPTY))

        right_indexes = self.left_set(empty_indexes_right)
        down_indexes = self.up_set(empty_indexes_down)

        self.cells[right_indexes] = Cell.EMPTY
        self.cells[down_indexes] = Cell.EMPTY
        self.cells[empty_indexes_right] = Cell.RIGHT
        self.cells[empty_indexes_down] = Cell.DOWN

        # if self.step % 100 == 1:
        #    self.save()
        return len(right_indexes[0]) + len(down_indexes[0])

    def step_all_together(self):
        return -1

    def upper(self, i, j):
        return self.height - 1 if i == 0 else i - 1, j

    def left(self, i, j):
        return i, self.width - 1 if j == 0 else j - 1

    def upper_t(self, ind):
        return self.height - 1 if ind[0] == 0 else ind[0] - 1, ind[1]

    def left_t(self, ind):
        return ind[0], self.width - 1 if ind[1] == 0 else ind[1] - 1

    def down_t(self, ind):
        return 0 if ind[0] == self.height - 1 else ind[0] + 1, ind[1]

    def right_t(self, ind):
        return ind[0], 0 if ind[1] == self.width - 1 else ind[1] + 1

    def down_set(self, inds):
        return tuple([[(x + 1) % self.height for x in inds[0]], inds[1]])

    def right_set(self, inds):
        return tuple([inds[0], [(x + 1) % self.width for x in inds[1]]])

    def up_set(self, inds):
        return tuple([[(x - 1) % self.height for x in inds[0]], inds[1]])

    def left_set(self, inds):
        return tuple([inds[0], [(x - 1) % self.width for x in inds[1]]])

    def run(self, steps):
        start = time.perf_counter()
        for i in range(steps):
            self.make_step()
        total_time = time.perf_counter() - start
        print('{} steps for {}x{} map run in {}s. Average of {}s per step'.format(steps, self.height, self.width,
                                                                                total_time, total_time / steps))

    def get_image(self, pixel_size=1):
        """
        Returns image
        """
        visual = np.zeros((self.height, self.width, 3)).astype('uint8')
        down_indexes = np.where(self.cells == Cell.DOWN)
        right_indexes = np.where(self.cells == Cell.RIGHT)
        visual[:, :, :] = 255
        visual[:, :, 0][down_indexes] = 0  # down is blue

        visual[:, :, 1][down_indexes] = 0
        visual[:, :, 1][right_indexes] = 0

        visual[:, :, 2][right_indexes] = 0  # right are red
        im = Image.fromarray(visual, mode='RGB')
        im = im.resize((self.height * pixel_size, self.width * pixel_size))
        return im

    def save(self, pixel_size=10):
        """
        Saves current state as png image
        """
        visual = np.zeros((self.height, self.width, 3)).astype('uint8')
        for i in range(self.height):
            for j in range(self.width):
                visual[i, j, 0] = 0 if self.cells[(i, j)] == Cell.DOWN else 255  # down is blue
                visual[i, j, 1] = 0 if self.cells[(i, j)] != Cell.EMPTY else 255
                visual[i, j, 2] = 0 if self.cells[(i, j)] == Cell.RIGHT else 255  # right are red
        im = Image.fromarray(visual, mode='RGB')
        im = im.resize((self.height * pixel_size, self.width * pixel_size))
        im.save("visual" + str(self.step) + ".png")

    def save_animation(self):
        self.animation[0].save("progress.gif", append_images=self.animation[1:], save_all=True, optimize=True,
                               duration=10)

    def plot_velocity(self):
        """
        Builds plot of velocity as a function of step number until now
        """
        params = {'axes.labelsize': 8, 'font.size': 8, 'legend.fontsize': 10, 'xtick.labelsize': 10,
                  'ytick.labelsize': 10, 'text.usetex': False, 'figure.figsize': [4.5, 4.5]}
        rcParams.update(params)
        fig = figure()  # no frame
        ax = fig.add_subplot(111)
        ax.plot(self.velocity, linewidth=1, color=colors[1])
        ax.set_ylim(0, 1)
        ax.set_yticks(np.arange(0, 1, 0.2))

        ax.get_xaxis().tick_bottom()
        ax.get_yaxis().tick_left()

        ax.set_xlabel('Step')
        ax.set_ylabel('Velocity')

        ax.grid(axis='y', color="0.9", linestyle='-', linewidth=1)
        ax.set_axisbelow(True)

        plt.tight_layout()

        fig.savefig('velocity.png', dpi=600)
        # print(self.velocity)

automat = BML(256, 256, 0.36, 1)
automat.save()
automat.run(5000)
automat.save()
automat.save_animation()
automat.plot_velocity()
<|endoftext|>


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

from wxpy import *
from config import *
import re
from wxpy.utils import start_new_thread
import time
import os

'''
使用 cache 来缓存登陆信息，同时使用控制台登陆
'''
bot = Bot('bot.pkl', console_qr=True)
bot.messages.max_history = 0

'''
开启 PUID 用于后续的控制
'''
bot.enable_puid('wxpy_puid.pkl')

'''
邀请信息处理
'''
rp_new_member_name = (
    re.compile(r'^"(.+)"通过'),
    re.compile(r'邀请"(.+)"加入'),
)

# 格式化 Group
try:
    groups = list(map(lambda x: bot.groups().search(puid=x)[0], group_puids))
except:
    print("查找管理群出错！请检查管理群 puid 是否输入正确")
    quit()

# 格式化 Admin
try:
    admins = list(map(lambda x: bot.friends().search(puid=x)[0], admin_puids))
except:
    print("查找管理员出错！请检查管理员 puid 是否输入正确")
    quit()

# 远程踢人命令: 移出 @<需要被移出的人>
rp_kick = re.compile(r'^(?:移出|移除|踢出|拉黑)\s*@(.+?)(?:\u2005?\s*$)')

# 下方为函数定义

def get_time():
    return str(time.strftime("%Y-%m-%d %H:%M:%S"))

'''
机器人消息提醒设置
'''
if alert_group:
    try:
        alert_receiver = ensure_one(bot.groups().search(alert_group))
    except:
        print("警报群设置有误，请检查群名是否存在且唯一")
        alert_receiver = bot.file_helper
else:
    alert_receiver = bot.file_helper
logger = get_wechat_logger(alert_receiver)
logger.error(str("机器人登陆成功！"+ get_time()))

'''
重启机器人
'''
def _restart():
    os.execv(sys.executable, [sys.executable] + sys.argv)

'''
状态汇报
'''
def status():
    status_text = get_time() + " 机器人目前在线,共有好友 【" + str(len(bot.friends())) + "】 群 【 " + str(len(bot.groups())) + "】"
    return status_text

'''
定时报告进程状态
'''
def heartbeat():
    while bot.alive:
        time.sleep(3600)
        # noinspection PyBroadException
        try:
            logger.error(status())
        except ResponseError as e:
            if 1100 <= e.err_code <= 1102:
                logger.critical('LCBot offline: {}'.format(e))
                _restart()

start_new_thread(heartbeat)

'''
条件邀请
'''
def condition_invite(user):
    if user.sex == 2:
        female_groups = bot.groups().search(female_group)[0]
        try:
            female_groups.add_members(user, use_invitation=True)
            pass
        except:
            pass
    if (user.province in city_group.keys() or user.city in city_group.keys()):
        try:
            target_city_group = bot.groups().search(city_group[user.province])[0]
            pass
        except:
            target_city_group = bot.groups().search(city_group[user.city])[0]
            pass
        try:
            if user not in target_city_group:
                target_city_group.add_members(user, use_invitation=True)
        except:
            pass

'''
判断消息发送者是否在管理员列表
'''
def from_admin(msg):
    """
    判断 msg 中的发送用户是否为管理员
    :param msg:
    :return:
    """
    if not isinstance(msg, Message):
        raise TypeError('expected Message, got {}'.format(type(msg)))
    from_user = msg.member if isinstance(msg.chat, Group) else msg.sender
    return from_user in admins

'''
远程踢人命令
'''
def remote_kick(msg):
    if msg.type is TEXT:
        match = rp_kick.search(msg.text)
        if match:
            name_to_kick = match.group(1)

            if not from_admin(msg) and not silence_mode:
                return '感觉有点不对劲… @{}'.format(msg.member.name)

            member_to_kick = ensure_one(list(filter(
                lambda x: x.name == name_to_kick, msg.chat)))
            if member_to_kick  == bot.self:
                return '无法移出 @{}'.format(member_to_kick.name)
            if member_to_kick in admins:
                return '无法移出 @{}'.format(member_to_kick.name)

            logger.error(get_time() + str(" 【"+member_to_kick.name + "】 被 【"+msg.member.name+"】 移出 【" + msg.sender.name+"】"))
            member_to_kick.set_remark_name("[黑名单]-"+get_time())
            if member_to_kick in msg.sender:
                member_to_kick.remove()
                kick_info = '成功移出 @{}'.format(member_to_kick.name)
            else:
                kick_info = '@{} 已不在群中'.format(member_to_kick.name)

            for ready_to_kick_group in  groups:
                if member_to_kick in ready_to_kick_group:
                    ready_to_kick_group.remove_members(member_to_kick)
                    logger.error(get_time()+ str("【"+member_to_kick.name + "】 被系统自动移出 " +  ready_to_kick_group.name))

            return kick_info

'''
邀请消息处理
'''
def get_new_member_name(msg):
    # itchat 1.2.32 版本未格式化群中的 Note 消息
    from itchat.utils import msg_formatter
    msg_formatter(msg.raw, 'Text')

    for rp in rp_new_member_name:
        match = rp.search(msg.text)
        if match:
            return match.group(1)

'''
定义邀请用户的方法。
按关键字搜索相应的群，如果存在相应的群，就向用户发起邀请。
'''
def invite(user, keyword):
    from random import randrange
    group = bot.groups().search(keyword_of_group[keyword])
    if len(group) > 0:
        for i in range(0, len(group)):
            if user in group[i]:
                content = "您已经加入了 {} [微笑]".format(group[i].nick_name)
                user.send(content)
                return
        if len(group) == 1:
            target_group = group[0]
        else:
            index = randrange(len(group))
            target_group = group[index]
        try:
            target_group.add_members(user, use_invitation=True)
        except:
            user.send("邀请错误！机器人邀请好友进群已达当日限制。请您明日再试")
    else:
        user.send("该群状态有误，您换个关键词试试？")

# 下方为消息处理

'''
处理加好友请求信息。
如果验证信息文本是字典的键值之一，则尝试拉群。
'''
@bot.register(msg_types=FRIENDS)
def new_friends(msg):
    user = msg.card.accept()
    if msg.text.lower() in keyword_of_group.keys():
        invite(user, msg.text.lower())
    else:
        user.send(invite_text)

@bot.register(Friend, msg_types=TEXT)
def exist_friends(msg):
    if msg.sender.name.find("黑名单") != -1:
        return "您已被拉黑！"
    else:
        if msg.text.lower() in keyword_of_group.keys():
            invite(msg.sender, msg.text.lower())
        else:
            return invite_text

# 管理群内的消息处理
@bot.register(groups, except_self=False)
def wxpy_group(msg):
    ret_msg = remote_kick(msg)
    if ret_msg:
        return ret_msg
    elif msg.is_at and not silence_mode:
        if turing_key :
            tuling = Tuling(api_key=turing_key)
            tuling.do_reply(msg)
        else:
            return "忙着呢，别烦我！";
            pass

@bot.register(groups, NOTE)
def welcome(msg):
    name = get_new_member_name(msg)
    if name and not silence_mode:
        return welcome_text.format(name)

@bot.register(alert_receiver, except_self=False)
def alert_command(msg):
    if from_admin(msg):
        if msg.text == "状态":
            return status()
        elif msg.text == "重启":
            _restart()

embed()
<commit_msg>增加监控群消息提醒等级<commit_after>#!/usr/bin/env python3
# coding: utf-8

from wxpy import *
from config import *
import re
from wxpy.utils import start_new_thread
import time
import os

'''
使用 cache 来缓存登陆信息，同时使用控制台登陆
'''
bot = Bot('bot.pkl', console_qr=True)
bot.messages.max_history = 0

'''
开启 PUID 用于后续的控制
'''
bot.enable_puid('wxpy_puid.pkl')

'''
邀请信息处理
'''
rp_new_member_name = (
    re.compile(r'^"(.+)"通过'),
    re.compile(r'邀请"(.+)"加入'),
)

# 格式化 Group
try:
    groups = list(map(lambda x: bot.groups().search(puid=x)[0], group_puids))
except:
    print("查找管理群出错！请检查管理群 puid 是否输入正确")
    quit()

# 格式化 Admin
try:
    admins = list(map(lambda x: bot.friends().search(puid=x)[0], admin_puids))
except:
    print("查找管理员出错！请检查管理员 puid 是否输入正确")
    quit()

# 远程踢人命令: 移出 @<需要被移出的人>
rp_kick = re.compile(r'^(?:移出|移除|踢出|拉黑)\s*@(.+?)(?:\u2005?\s*$)')

# 下方为函数定义

def get_time():
    return str(time.strftime("%Y-%m-%d %H:%M:%S"))

'''
机器人消息提醒设置
'''
alert_level = 30 # DEBUG: 10, INFO: 20, WARNING: 30, ERROR: 40, FATAL: 50
if alert_group:
    try:
        alert_receiver = ensure_one(bot.groups().search(alert_group))
    except:
        print("警报群设置有误，请检查群名是否存在且唯一")
        alert_receiver = bot.file_helper
else:
    alert_receiver = bot.file_helper
logger = get_wechat_logger(alert_receiver, alert_level)
logger.error(str("机器人登陆成功！"+ get_time()))

'''
重启机器人
'''
def _restart():
    os.execv(sys.executable, [sys.executable] + sys.argv)

'''
状态汇报
'''
def status():
    status_text = get_time() + " 机器人目前在线,共有好友 【" + str(len(bot.friends())) + "】 群 【 " + str(len(bot.groups())) + "】"
    return status_text

'''
定时报告进程状态
'''
def heartbeat():
    while bot.alive:
        time.sleep(3600)
        # noinspection PyBroadException
        try:
            logger.error(status())
        except ResponseError as e:
            if 1100 <= e.err_code <= 1102:
                logger.critical('LCBot offline: {}'.format(e))
                _restart()

start_new_thread(heartbeat)

'''
条件邀请
'''
def condition_invite(user):
    if user.sex == 2:
        female_groups = bot.groups().search(female_group)[0]
        try:
            female_groups.add_members(user, use_invitation=True)
            pass
        except:
            pass
    if (user.province in city_group.keys() or user.city in city_group.keys()):
        try:
            target_city_group = bot.groups().search(city_group[user.province])[0]
            pass
        except:
            target_city_group = bot.groups().search(city_group[user.city])[0]
            pass
        try:
            if user not in target_city_group:
                target_city_group.add_members(user, use_invitation=True)
        except:
            pass

'''
判断消息发送者是否在管理员列表
'''
def from_admin(msg):
    """
    判断 msg 中的发送用户是否为管理员
    :param msg:
    :return:
    """
    if not isinstance(msg, Message):
        raise TypeError('expected Message, got {}'.format(type(msg)))
    from_user = msg.member if isinstance(msg.chat, Group) else msg.sender
    return from_user in admins

'''
远程踢人命令
'''
def remote_kick(msg):
    if msg.type is TEXT:
        match = rp_kick.search(msg.text)
        if match:
            name_to_kick = match.group(1)

            if not from_admin(msg) and not silence_mode:
                return '感觉有点不对劲… @{}'.format(msg.member.name)

            member_to_kick = ensure_one(list(filter(
                lambda x: x.name == name_to_kick, msg.chat)))
            if member_to_kick  == bot.self:
                return '无法移出 @{}'.format(member_to_kick.name)
            if member_to_kick in admins:
                return '无法移出 @{}'.format(member_to_kick.name)

            logger.error(get_time() + str(" 【"+member_to_kick.name + "】 被 【"+msg.member.name+"】 移出 【" + msg.sender.name+"】"))
            member_to_kick.set_remark_name("[黑名单]-"+get_time())
            if member_to_kick in msg.sender:
                member_to_kick.remove()
                kick_info = '成功移出 @{}'.format(member_to_kick.name)
            else:
                kick_info = '@{} 已不在群中'.format(member_to_kick.name)

            for ready_to_kick_group in  groups:
                if member_to_kick in ready_to_kick_group:
                    ready_to_kick_group.remove_members(member_to_kick)
                    logger.error(get_time()+ str("【"+member_to_kick.name + "】 被系统自动移出 " +  ready_to_kick_group.name))

            return kick_info

'''
邀请消息处理
'''
def get_new_member_name(msg):
    # itchat 1.2.32 版本未格式化群中的 Note 消息
    from itchat.utils import msg_formatter
    msg_formatter(msg.raw, 'Text')

    for rp in rp_new_member_name:
        match = rp.search(msg.text)
        if match:
            return match.group(1)

'''
定义邀请用户的方法。
按关键字搜索相应的群，如果存在相应的群，就向用户发起邀请。
'''
def invite(user, keyword):
    from random import randrange
    group = bot.groups().search(keyword_of_group[keyword])
    if len(group) > 0:
        for i in range(0, len(group)):
            if user in group[i]:
                content = "您已经加入了 {} [微笑]".format(group[i].nick_name)
                user.send(content)
                return
        if len(group) == 1:
            target_group = group[0]
        else:
            index = randrange(len(group))
            target_group = group[index]
        try:
            target_group.add_members(user, use_invitation=True)
        except:
            user.send("邀请错误！机器人邀请好友进群已达当日限制。请您明日再试")
    else:
        user.send("该群状态有误，您换个关键词试试？")

# 下方为消息处理

'''
处理加好友请求信息。
如果验证信息文本是字典的键值之一，则尝试拉群。
'''
@bot.register(msg_types=FRIENDS)
def new_friends(msg):
    user = msg.card.accept()
    if msg.text.lower() in keyword_of_group.keys():
        invite(user, msg.text.lower())
    else:
        user.send(invite_text)

@bot.register(Friend, msg_types=TEXT)
def exist_friends(msg):
    if msg.sender.name.find("黑名单") != -1:
        return "您已被拉黑！"
    else:
        if msg.text.lower() in keyword_of_group.keys():
            invite(msg.sender, msg.text.lower())
        else:
            return invite_text

# 管理群内的消息处理
@bot.register(groups, except_self=False)
def wxpy_group(msg):
    ret_msg = remote_kick(msg)
    if ret_msg:
        return ret_msg
    elif msg.is_at and not silence_mode:
        if turing_key :
            tuling = Tuling(api_key=turing_key)
            tuling.do_reply(msg)
        else:
            return "忙着呢，别烦我！";
            pass

@bot.register(groups, NOTE)
def welcome(msg):
    name = get_new_member_name(msg)
    if name and not silence_mode:
        return welcome_text.format(name)

@bot.register(alert_receiver, except_self=False)
def alert_command(msg):
    if from_admin(msg):
        if msg.text == "状态":
            return status()
        elif msg.text == "重启":
            _restart()

embed()
<|endoftext|>


<commit_before># Import various libraries
import discord
import asyncio
import subprocess
import requests
import json

# Initialize bot client
# TODO: Make bot a class like normal bots.
client = discord.Client()

# A dictionary of basic things for the bot to return. More complex (i.e.
# data-driven) interactions aren't stored here, those go below.
message_index = {
    '!ping': 'Pong!',
    '!hello': 'World!',
    'rickroll': 'https://www.youtube.com/watch?v=dQw4w9WgXcQ',
    'xcq': 'https://www.youtube.com/watch?v=dQw4w9WgXcQ',
    '!balloumoji': ':bigdissapointment::moustache::ballouminatti::1982::nope::notapproved::fedora1::happy::flowers::notbad::soundboard:',
    'it\'s time to stop': 'https://www.youtube.com/watch?v=2k0SmqbBIpQ',
    'stop': 'https://www.youtube.com/watch?v=2k0SmqbBIpQ'
}

@client.event
async def on_ready():
    """Run when the bot is ready."""
    print('Logged in as ' + client.user.name + ' (ID ' + client.user.id + ').')
    print('------')

@client.event
async def on_message(message):
    """Catch a user's messages and figure out what to return."""

    msg = message.content.lower()

    # Only send back message if user that sent the triggering message isn't a bot
    if not message.author.bot:
        # Special returns!
        if msg.startswith('!about'):
            await client.send_message(message.channel, 'Victibot is a chatbot for Team 1418\'s Discord server. Bot is currently running as ' + client.user.name + ' (ID ' + client.user.id + '). View on GitHub: https://github.com/ErikBoesen/victibot')
        elif msg.startswith('!xkcd'):
            # Store the number/other content after the '!xkcd '.
            comic = msg[6:]

            # If the user included a specific comic number in their message, get the JSON data for that comic. Otherwise, get the JSON data for the most recent comic.
            r = requests.get('http://xkcd.com/' + comic + '/info.0.json' if comic else 'http://xkcd.com/info.0.json')

            # Send the URL of the image from the JSON fetched above.
            await client.send_message(message.channel, r.json()['img'])
        elif msg == '!update':
            # Confirm that the bot is updating
            await client.send_message(message.channel, 'Updating...')
            # Start a git pull to update bot
            print(str(subprocess.Popen('git pull', shell=True, stdout=subprocess.PIPE).stdout.read()))
        else:
            # Respond if the message has a basic, static response.
            # TODO: Apparently 'await' has been replaced in py3 with 'yield from'.
            # Implement this change.
            try:
                await client.send_message(message.channel, message_index[msg])
            except:
                pass

@client.event
async def on_member_join(member):
    await client.send_message(member.server, member + ' has joined the server!')

@client.event
async def on_member_remove(member):
    await client.send_message(member.server, member + ' has left the server!')

# Get token from token.txt.
with open('token.txt', 'r') as token_file:
    token = token_file.read()

print('Starting with token ' + token + '...')

# Start bot!
client.run(token)

# That's all, folks.<commit_msg>Fix for trailing newlines in token.txt files<commit_after># Import various libraries
import discord
import asyncio
import subprocess
import requests
import json

# Initialize bot client
# TODO: Make bot a class like normal bots.
client = discord.Client()

# A dictionary of basic things for the bot to return. More complex (i.e.
# data-driven) interactions aren't stored here, those go below.
message_index = {
    '!ping': 'Pong!',
    '!hello': 'World!',
    'rickroll': 'https://www.youtube.com/watch?v=dQw4w9WgXcQ',
    'xcq': 'https://www.youtube.com/watch?v=dQw4w9WgXcQ',
    '!balloumoji': ':bigdissapointment::moustache::ballouminatti::1982::nope::notapproved::fedora1::happy::flowers::notbad::soundboard:',
    'it\'s time to stop': 'https://www.youtube.com/watch?v=2k0SmqbBIpQ',
    'stop': 'https://www.youtube.com/watch?v=2k0SmqbBIpQ'
}

@client.event
async def on_ready():
    """Run when the bot is ready."""
    print('Logged in as ' + client.user.name + ' (ID ' + client.user.id + ').')
    print('------')

@client.event
async def on_message(message):
    """Catch a user's messages and figure out what to return."""

    msg = message.content.lower()

    # Only send back message if user that sent the triggering message isn't a bot
    if not message.author.bot:
        # Special returns!
        if msg.startswith('!about'):
            await client.send_message(message.channel, 'Victibot is a chatbot for Team 1418\'s Discord server. Bot is currently running as ' + client.user.name + ' (ID ' + client.user.id + '). View on GitHub: https://github.com/ErikBoesen/victibot')
        elif msg.startswith('!xkcd'):
            # Store the number/other content after the '!xkcd '.
            comic = msg[6:]

            # If the user included a specific comic number in their message, get the JSON data for that comic. Otherwise, get the JSON data for the most recent comic.
            r = requests.get('http://xkcd.com/' + comic + '/info.0.json' if comic else 'http://xkcd.com/info.0.json')

            # Send the URL of the image from the JSON fetched above.
            await client.send_message(message.channel, r.json()['img'])
        elif msg == '!update':
            # Confirm that the bot is updating
            await client.send_message(message.channel, 'Updating...')
            # Start a git pull to update bot
            print(str(subprocess.Popen('git pull', shell=True, stdout=subprocess.PIPE).stdout.read()))
        else:
            # Respond if the message has a basic, static response.
            # TODO: Apparently 'await' has been replaced in py3 with 'yield from'.
            # Implement this change.
            try:
                await client.send_message(message.channel, message_index[msg])
            except:
                pass

@client.event
async def on_member_join(member):
    await client.send_message(member.server, member + ' has joined the server!')

@client.event
async def on_member_remove(member):
    await client.send_message(member.server, member + ' has left the server!')

# Get token from token.txt.
with open('token.txt', 'r') as token_file:
    # Parse into a string, and get rid of trailing newlines.
    token = token_file.read().replace('\n', '')

print('Starting with token ' + token + '...')

# Start bot!
client.run(token)

# That's all, folks.<|endoftext|>