<commit_before># -*- coding: utf-8 -*-

# csv2cmi
# version 0.9.11.1
# Copyright (c) 2015 Klaus Rettinghaus
# programmed by Klaus Rettinghaus
# licensed under MIT license

# needs Python3
import datetime
import csv
import logging
import os
import random
import string
import sys
import urllib.request
from xml.etree.ElementTree import Element, SubElement, Comment, tostring, ElementTree
from xml.dom import minidom

# enter here your name and the title of your edition
fileName = 'LetterTable.csv'
projectName = 'Letters project'
editorName = 'Klaus Rettinghaus'
edition = ''
editionType = 'print'  # 'hybrid' or 'online' if appropriate

# define log output
logging.basicConfig(format='%(levelname)s: %(message)s')

# define RDF types

# simple test for file
if sys.argv[-1] != sys.argv[0]:
    fileName = sys.argv[-1]
try:
    open(fileName, 'rt')
except:
    logging.error('File not found')
    exit()

# check internet connection via DNB
try:
    urllib.request.urlopen('http://193.175.100.220', timeout=1)
    connection = True
except:
    logging.error('No internet connection')
    connection = False

def isodate(datestring):
    try:
        datetime.datetime.strptime(datestring, '%Y-%m-%d')
    except:
        try:
            datetime.datetime.strptime(datestring, '%Y-%m')
        except:
            try:
                datetime.datetime.strptime(datestring, '%Y')
            except ValueError:
                return False
            else:
                return True
        else:
            return True
    else:
        return True

def createPerson(namestring):
    if letter[namestring]:
        rdf = {'rdf': 'http://www.w3.org/1999/02/22-rdf-syntax-ns#'}
        if (namestring + 'ID' in table.fieldnames) and (letter[namestring + 'ID']):
            if 'http://' not in str(letter[namestring + 'ID'].strip()):
                authID = 'http://d-nb.info/gnd/' + \
                    str(letter[namestring + 'ID'].strip())
            else:
                authID = str(letter[namestring + 'ID'].strip())
            if connection and (profileDesc.find('correspDesc/correspAction/persName[@ref="' + authID + '"]') == None):
                if 'viaf' in authID:
                    try:
                        gndrdf = ElementTree(
                            file=urllib.request.urlopen(authID + '/rdf.xml'))
                    except urllib.error.HTTPError:
                        logging.error(
                            'Authority file not found for %sID in line %s', namestring, table.line_num)
                        persName = SubElement(action, 'persName')
                        authID = ''
                    except urllib.error.URLError as argh:
                        logging.error('Failed to reach VIAF')
                        persName = SubElement(action, 'persName')
                    else:
                        gndrdf_root = gndrdf.getroot()
                        rdftype = gndrdf_root.find(
                            './rdf:Description[2]/rdf:type', rdf).get('{http://www.w3.org/1999/02/22-rdf-syntax-ns#}resource')
                        if 'Organization' in rdftype:
                            persName = SubElement(action, 'orgName')
                        elif 'Person' in rdftype:
                            persName = SubElement(action, 'persName')
                        else:
                            logging.warning(
                                '%sID in line %s links to unprocessable authority file', namestring, table.line_num)
                            persName = SubElement(action, 'persName')
                            authID = ''
                elif 'gnd' in authID:
                    try:
                        gndrdf = ElementTree(
                            file=urllib.request.urlopen(authID + '/about/rdf'))
                    except urllib.error.HTTPError:
                        logging.error(
                            'Authority file not found for %sID in line %s', namestring, table.line_num)
                        persName = SubElement(action, 'persName')
                        authID = ''
                    except urllib.error.URLError as argh:
                        logging.error('Failed to reach GND')
                        persName = SubElement(action, 'persName')
                    else:
                        gndrdf_root = gndrdf.getroot()
                        rdftype = gndrdf_root.find(
                            './/rdf:type', rdf).get('rdf:resource',rdf)
                        if 'Corporate' in rdftype:
                            persName = SubElement(action, 'orgName')
                        else:
                            persName = SubElement(action, 'persName')
                        if 'UndifferentiatedPerson' in rdftype:
                            logging.warning(
                                '%sID in line %s links to undifferentiated Person', namestring, table.line_num)
                            authID = ''
                else:
                    logging.error('No proper autority record in line %s', table.line_num)
                    persName = SubElement(action, 'persName')
                    authID = ''
            else:
                persName = SubElement(action, 'persName')
            if authID:
                persName.set('ref', authID)
        else:
            persName = SubElement(action, 'persName')
        if letter[namestring].startswith('[') and letter[namestring].endswith(']'):
            persName.set('evidence', 'conjecture')
            letter[namestring] = letter[namestring][1:-1]
            logging.info('Added @evidence for <persName> in line %s',
                         table.line_num)
        persName.text = str(letter[namestring])

def createPlace(placestring):
    # function for putting the place in <correspAction>
    if letter[placestring]:
        placeName = SubElement(action, 'placeName')
        letter[placestring] = letter[placestring].strip()
        if letter[placestring].startswith('[') and letter[placestring].endswith(']'):
            placeName.set('evidence', 'conjecture')
            letter[placestring] = letter[placestring][1:-1]
            logging.info('Added @evidence for <placeName> in line %s',
                         table.line_num)
        placeName.text = str(letter[placestring])
        if (placestring + 'ID' in table.fieldnames) and (letter[placestring + 'ID']):
            letter[placestring + 'ID'] = letter[placestring + 'ID'].strip()
            if 'http://www.geonames.org/' in letter[placestring + 'ID']:
                placeName.set('ref', str(letter[placestring + 'ID']))
            else:
                logging.warning("no standardized %sID in line %s",
                                placestring, table.line_num)
        else:
            logging.warning('ID for %s missing in line %s', letter[
                placestring], table.line_num)

def createEdition(biblText, biblID):
    # creates a new entry within <sourceDesc>
    bibl = SubElement(sourceDesc, 'bibl')
    bibl.text = biblText
    bibl.set('type', editionType)
    bibl.set('xml:id', biblID)

def getEditonID(editionTitle):
    editionID = ''
    for bibl in sourceDesc.findall('bibl'):
        if editionTitle == bibl.text:
            editionID = bibl.get('xml:id')
    if not editionID:
        editionID = createID('edition')
        createEdition(editionTitle, editionID)
    return editionID

def createID(id_prefix):
    fullID = id_prefix + '_' + ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(
        8)) + '_' + ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(8))
    return fullID

# building cmi
# generating root element
root = Element('TEI')
root.set('xmlns', 'http://www.tei-c.org/ns/1.0')
root.append(Comment('Generated from table of letters with csv2cmi'))

# teiHeader
teiHeader = SubElement(root, 'teiHeader')
# file description
fileDesc = SubElement(teiHeader, 'fileDesc')
titleStmt = SubElement(fileDesc, 'titleStmt')
title = SubElement(titleStmt, 'title')
title.text = projectName
editor = SubElement(titleStmt, 'editor')
editor.text = editorName
publicationStmt = SubElement(fileDesc, 'publicationStmt')
publisher = SubElement(publicationStmt, 'publisher')
publisher.text = editorName
idno = SubElement(publicationStmt, 'idno')
idno.set('type', 'url')
idno.text = os.path.splitext(os.path.basename(fileName))[0] + '.xml'
date = SubElement(publicationStmt, 'date')
date.set('when', str(datetime.datetime.now().isoformat()))
availability = SubElement(publicationStmt, 'availability')
licence = SubElement(availability, 'licence')
licence.set('target', 'https://creativecommons.org/licenses/by/4.0/')
licence.text = 'This file is licensed under the terms of the Creative-Commons-License CC-BY 4.0'
sourceDesc = SubElement(fileDesc, 'sourceDesc')

# filling in correspondance meta-data
profileDesc = SubElement(teiHeader, 'profileDesc')

with open(fileName, 'rt') as letterTable:
    table = csv.DictReader(letterTable)
    logging.info('Recognized columns: %s', table.fieldnames)
    if not('edition' in table.fieldnames):
        if edition:
            createEdition(edition, createID('edition'))
        else:
            logging.warning('No edition stated. Please set manually.')
            bibl = SubElement(sourceDesc, 'bibl')
    for letter in table:
        entry = SubElement(profileDesc, 'correspDesc')
        entry.set('xml:id', createID('letter'))
        if 'edition' in table.fieldnames:
            edition = letter['edition'].strip()
        if edition:
            entry.set('source', '#' + getEditonID(edition))
        if 'key' in table.fieldnames and letter['key']:
            if not(edition):
                logging.error("Key without edition in line %s", table.line_num)
            else:
                if 'http://' in str(letter['key']):
                    entry.set('ref', str(letter['key']).strip())
                else:
                    entry.set('key', str(letter['key']).strip())

        # sender info block
        if (letter['sender']) or (('senderPlace' in table.fieldnames) and (letter['senderPlace'])) or (letter['senderDate']):
            action = SubElement(entry, 'correspAction')
            action.set('type', 'sent')

            createPerson('sender')
            if 'senderPlace' in table.fieldnames:
                createPlace('senderPlace')

        if isodate(letter['senderDate']) or isodate(letter['senderDate'][1:-1]):
            senderDate = SubElement(action, 'date')
            if ('[' in str(letter['senderDate'])) and (']' in str(letter['senderDate'])):
                letter['senderDate'] = letter['senderDate'][1:-1]
                senderDate.set('cert', 'medium')
                logging.info('Added @cert for <date> in line %s',
                             table.line_num)
            senderDate.set('when', str(letter['senderDate']))
        else:
            logging.warning('senderDate in line %s not set (no ISO)',
                            table.line_num)

        # addressee info block
        if letter['addressee'] or 'addresseePlace' in table.fieldnames or (('addresseeDate') in table.fieldnames and (letter['addresseeDate'])):
            action = SubElement(entry, 'correspAction')
            action.set('type', 'received')

            createPerson('addressee')
            if 'addresseePlace' in table.fieldnames:
                createPlace('addresseePlace')

letterTable.close()

# generate empty body
text = SubElement(root, 'text')
body = SubElement(text, 'body')
p = SubElement(body, 'p')

# save cmi to file
tree = ElementTree(root)
tree.write(os.path.splitext(os.path.basename(fileName))[
           0] + '.xml', encoding="utf-8", xml_declaration=True, method="xml")
<commit_msg>minor gnd fix<commit_after># -*- coding: utf-8 -*-

# csv2cmi
# version 0.9.11.2
# Copyright (c) 2015 Klaus Rettinghaus
# programmed by Klaus Rettinghaus
# licensed under MIT license

# needs Python3
import datetime
import csv
import logging
import os
import random
import string
import sys
import urllib.request
from xml.etree.ElementTree import Element, SubElement, Comment, tostring, ElementTree
from xml.dom import minidom

# enter here your name and the title of your edition
fileName = 'LetterTable.csv'
projectName = 'Letters project'
editorName = 'Klaus Rettinghaus'
edition = ''
editionType = 'print'  # 'hybrid' or 'online' if appropriate

# define log output
logging.basicConfig(format='%(levelname)s: %(message)s')

# define RDF types

# simple test for file
if sys.argv[-1] != sys.argv[0]:
    fileName = sys.argv[-1]
try:
    open(fileName, 'rt')
except:
    logging.error('File not found')
    exit()

# check internet connection via DNB
try:
    urllib.request.urlopen('http://193.175.100.220', timeout=1)
    connection = True
except:
    logging.error('No internet connection')
    connection = False

def isodate(datestring):
    try:
        datetime.datetime.strptime(datestring, '%Y-%m-%d')
    except:
        try:
            datetime.datetime.strptime(datestring, '%Y-%m')
        except:
            try:
                datetime.datetime.strptime(datestring, '%Y')
            except ValueError:
                return False
            else:
                return True
        else:
            return True
    else:
        return True

def createPerson(namestring):
    if letter[namestring]:
        rdf = {'rdf': 'http://www.w3.org/1999/02/22-rdf-syntax-ns#'}
        if (namestring + 'ID' in table.fieldnames) and (letter[namestring + 'ID']):
            if 'http://' not in str(letter[namestring + 'ID'].strip()):
                authID = 'http://d-nb.info/gnd/' + \
                    str(letter[namestring + 'ID'].strip())
            else:
                authID = str(letter[namestring + 'ID'].strip())
            if connection and (profileDesc.find('correspDesc/correspAction/persName[@ref="' + authID + '"]') == None):
                if 'viaf' in authID:
                    try:
                        gndrdf = ElementTree(
                            file=urllib.request.urlopen(authID + '/rdf.xml'))
                    except urllib.error.HTTPError:
                        logging.error(
                            'Authority file not found for %sID in line %s', namestring, table.line_num)
                        persName = SubElement(action, 'persName')
                        authID = ''
                    except urllib.error.URLError as argh:
                        logging.error('Failed to reach VIAF')
                        persName = SubElement(action, 'persName')
                    else:
                        gndrdf_root = gndrdf.getroot()
                        rdftype = gndrdf_root.find(
                            './rdf:Description[2]/rdf:type', rdf).get('{http://www.w3.org/1999/02/22-rdf-syntax-ns#}resource')
                        if 'Organization' in rdftype:
                            persName = SubElement(action, 'orgName')
                        elif 'Person' in rdftype:
                            persName = SubElement(action, 'persName')
                        else:
                            logging.warning(
                                '%sID in line %s links to unprocessable authority file', namestring, table.line_num)
                            persName = SubElement(action, 'persName')
                            authID = ''
                elif 'gnd' in authID:
                    try:
                        gndrdf = ElementTree(
                            file=urllib.request.urlopen(authID + '/about/rdf'))
                    except urllib.error.HTTPError:
                        logging.error(
                            'Authority file not found for %sID in line %s', namestring, table.line_num)
                        persName = SubElement(action, 'persName')
                        authID = ''
                    except urllib.error.URLError as argh:
                        logging.error('Failed to reach GND')
                        persName = SubElement(action, 'persName')
                    else:
                        gndrdf_root = gndrdf.getroot()
                        rdftype = gndrdf_root.find(
                            './/rdf:type', rdf).get('{http://www.w3.org/1999/02/22-rdf-syntax-ns#}resource')
                        if 'Corporate' in rdftype:
                            persName = SubElement(action, 'orgName')
                        else:
                            persName = SubElement(action, 'persName')
                        if 'UndifferentiatedPerson' in rdftype:
                            logging.warning(
                                '%sID in line %s links to undifferentiated Person', namestring, table.line_num)
                            authID = ''
                else:
                    logging.error('No proper autority record in line %s', table.line_num)
                    persName = SubElement(action, 'persName')
                    authID = ''
            else:
                persName = SubElement(action, 'persName')
            if authID:
                persName.set('ref', authID)
        else:
            persName = SubElement(action, 'persName')
        if letter[namestring].startswith('[') and letter[namestring].endswith(']'):
            persName.set('evidence', 'conjecture')
            letter[namestring] = letter[namestring][1:-1]
            logging.info('Added @evidence for <persName> in line %s',
                         table.line_num)
        persName.text = str(letter[namestring])

def createPlace(placestring):
    # function for putting the place in <correspAction>
    if letter[placestring]:
        placeName = SubElement(action, 'placeName')
        letter[placestring] = letter[placestring].strip()
        if letter[placestring].startswith('[') and letter[placestring].endswith(']'):
            placeName.set('evidence', 'conjecture')
            letter[placestring] = letter[placestring][1:-1]
            logging.info('Added @evidence for <placeName> in line %s',
                         table.line_num)
        placeName.text = str(letter[placestring])
        if (placestring + 'ID' in table.fieldnames) and (letter[placestring + 'ID']):
            letter[placestring + 'ID'] = letter[placestring + 'ID'].strip()
            if 'http://www.geonames.org/' in letter[placestring + 'ID']:
                placeName.set('ref', str(letter[placestring + 'ID']))
            else:
                logging.warning("no standardized %sID in line %s",
                                placestring, table.line_num)
        else:
            logging.warning('ID for %s missing in line %s', letter[
                placestring], table.line_num)

def createEdition(biblText, biblID):
    # creates a new entry within <sourceDesc>
    bibl = SubElement(sourceDesc, 'bibl')
    bibl.text = biblText
    bibl.set('type', editionType)
    bibl.set('xml:id', biblID)

def getEditonID(editionTitle):
    editionID = ''
    for bibl in sourceDesc.findall('bibl'):
        if editionTitle == bibl.text:
            editionID = bibl.get('xml:id')
    if not editionID:
        editionID = createID('edition')
        createEdition(editionTitle, editionID)
    return editionID

def createID(id_prefix):
    fullID = id_prefix + '_' + ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(
        8)) + '_' + ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(8))
    return fullID

# building cmi
# generating root element
root = Element('TEI')
root.set('xmlns', 'http://www.tei-c.org/ns/1.0')
root.append(Comment('Generated from table of letters with csv2cmi'))

# teiHeader
teiHeader = SubElement(root, 'teiHeader')
# file description
fileDesc = SubElement(teiHeader, 'fileDesc')
titleStmt = SubElement(fileDesc, 'titleStmt')
title = SubElement(titleStmt, 'title')
title.text = projectName
editor = SubElement(titleStmt, 'editor')
editor.text = editorName
publicationStmt = SubElement(fileDesc, 'publicationStmt')
publisher = SubElement(publicationStmt, 'publisher')
publisher.text = editorName
idno = SubElement(publicationStmt, 'idno')
idno.set('type', 'url')
idno.text = os.path.splitext(os.path.basename(fileName))[0] + '.xml'
date = SubElement(publicationStmt, 'date')
date.set('when', str(datetime.datetime.now().isoformat()))
availability = SubElement(publicationStmt, 'availability')
licence = SubElement(availability, 'licence')
licence.set('target', 'https://creativecommons.org/licenses/by/4.0/')
licence.text = 'This file is licensed under the terms of the Creative-Commons-License CC-BY 4.0'
sourceDesc = SubElement(fileDesc, 'sourceDesc')

# filling in correspondance meta-data
profileDesc = SubElement(teiHeader, 'profileDesc')

with open(fileName, 'rt') as letterTable:
    table = csv.DictReader(letterTable)
    logging.info('Recognized columns: %s', table.fieldnames)
    if not('edition' in table.fieldnames):
        if edition:
            createEdition(edition, createID('edition'))
        else:
            logging.warning('No edition stated. Please set manually.')
            bibl = SubElement(sourceDesc, 'bibl')
    for letter in table:
        entry = SubElement(profileDesc, 'correspDesc')
        entry.set('xml:id', createID('letter'))
        if 'edition' in table.fieldnames:
            edition = letter['edition'].strip()
        if edition:
            entry.set('source', '#' + getEditonID(edition))
        if 'key' in table.fieldnames and letter['key']:
            if not(edition):
                logging.error("Key without edition in line %s", table.line_num)
            else:
                if 'http://' in str(letter['key']):
                    entry.set('ref', str(letter['key']).strip())
                else:
                    entry.set('key', str(letter['key']).strip())

        # sender info block
        if (letter['sender']) or (('senderPlace' in table.fieldnames) and (letter['senderPlace'])) or (letter['senderDate']):
            action = SubElement(entry, 'correspAction')
            action.set('type', 'sent')

            createPerson('sender')
            if 'senderPlace' in table.fieldnames:
                createPlace('senderPlace')

        if isodate(letter['senderDate']) or isodate(letter['senderDate'][1:-1]):
            senderDate = SubElement(action, 'date')
            if ('[' in str(letter['senderDate'])) and (']' in str(letter['senderDate'])):
                letter['senderDate'] = letter['senderDate'][1:-1]
                senderDate.set('cert', 'medium')
                logging.info('Added @cert for <date> in line %s',
                             table.line_num)
            senderDate.set('when', str(letter['senderDate']))
        else:
            logging.warning('senderDate in line %s not set (no ISO)',
                            table.line_num)

        # addressee info block
        if letter['addressee'] or 'addresseePlace' in table.fieldnames or (('addresseeDate') in table.fieldnames and (letter['addresseeDate'])):
            action = SubElement(entry, 'correspAction')
            action.set('type', 'received')

            createPerson('addressee')
            if 'addresseePlace' in table.fieldnames:
                createPlace('addresseePlace')

letterTable.close()

# generate empty body
text = SubElement(root, 'text')
body = SubElement(text, 'body')
p = SubElement(body, 'p')

# save cmi to file
tree = ElementTree(root)
tree.write(os.path.splitext(os.path.basename(fileName))[
           0] + '.xml', encoding="utf-8", xml_declaration=True, method="xml")
<|endoftext|>


<commit_before>#!/usr/bin/python

# Note: python >= 2.7 required
# Note: you will need s3cmd installed from http://s3tools.org/, I was using 1.5.0-rc1 (http://sourceforge.net/projects/s3tools/files/s3cmd/1.5.0-rc1/s3cmd-1.5.0-rc1.tar.gz)

import pp
import time         # for sleep
import os
import shutil
import sys
import string
import urllib2
import random
import errno
import argparse
import zipfile

parser = argparse.ArgumentParser(description='Crawl a set of ISBNs on amazon.com')
parser.add_argument('-i','--isbns', type=file, required=True, help='The file that contains the ISBNs to crawl')
parser.add_argument('-n','--nodes', type=file, help='The file that contains the cluster slave nodes to use')
parser.add_argument('-k','--secretkey', type=str, help='The secret key used to communicate with the cluster slave nodes')
parser.add_argument('-s','--s3credentials', type=file, help='The file containing the S3 credentials required to upload:\n\tbucket=<bucketname>\n\taccessKey=<access key>\n\tsecretKey=<secret key>')
args = parser.parse_args()

LocalRun = (args.nodes == None)
if LocalRun:
    print "No slave nodes provided, assuming you just want to run on this machine and this is a LOCAL RUN"

s3info = {}
if (args.s3credentials != None):
    s3lines = args.s3credentials.read().splitlines()
    for line in s3lines:
        kv = line.split('=')
        s3info[str.lower(kv[0])] = kv[1]
else:
    print "No S3 credentials provided, we will store the files locally"

userAgents = [
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/37.0.2062.94 Safari/537.36",
    "Mozilla/5.0 (compatible; MSIE 10.0; Windows NT 6.2; Trident/6.0)",
    "Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:25.0) Gecko/20100101 Firefox/29.0"
]
basePath = '/tmp/books/' if LocalRun else '/vol/temp_data/'
basePath = basePath + ''.join(random.choice(string.ascii_uppercase) for i in range(8)) + '/'

# download the contents of a url to the book's directory

#def _dummy_job(bookid, headers, fpath): # useful for debugging to short circuit the fetching
def _fetch_url_ (bookid, headers, fpath):
    url = 'http://www.amazon.com/dp/'+str(bookid)
    captchaText = '<form method="get" action="/errors/validateCaptcha" name="">'
    success = False
    failfast = False
    maxTries = 5
    ntries = 0
    sleep = 2
    request = urllib2.Request(url,None,headers)

    while not success and not failfast and ntries < maxTries:
        text = ''
        try:
            response = urllib2.urlopen(request)
            if response.getcode() == 200:
                text = response.read()
                if captchaText in text:
                    print "Got captcha, sleeping and retrying"
                else:
                    with open(fpath, 'wb') as f:
                        f.write(text)
                    success = True
            else:
                print "Received {0} code when retrieving {1}".format(response.getcode(), url)
        except urllib2.HTTPError, e:
            print "When retrieving {0} received HTTP Error: {1}".format(url, e.code)
            if e.code == 404:
                print "received 404, so skipping further retries"
                failfast = True
        except urllib2.URLError, e:
            print "URL Error:", e.reason , url
        except IOError, e:
            print "Could not write to file {0} when fetching {1}... skipping this book".format(fpath,url)
            print "Detailed Exception: ", e
            failfast = True
        finally:
            try:
                response.close()
            except NameError:
                pass

        if success:
            time.sleep(sleep) # sleep before trying next task
        else:
            if not failfast:
                # exponential backoff
                print "Attempt {0} for url {1} failed, backing off {2} seconds".format(ntries, url, sleep**ntries)
                time.sleep(sleep**ntries)
        ntries+=1
    return (bookid, success)

#def _fetch_url_ (bookid, headers, fpath):
def _dummy_job(bookid, headers, fpath):
    time.sleep(1)
    text = 'Fake book data'
    return (bookid, text)

def startJob(bookid,jobserver, outputPath):
    path = outputPath + str(bookid)
    index = random.randint(0,len(userAgents)-1)
    headers = {
        'User-Agent' : userAgents[index],
        'Accept' : 'text/html,application/xhtml+xml,application/xml;q=0.9',
        'Accept-Language' : 'en-US,en;q=0.8'
    }
    def inner():
        return _fetch_url_(bookid, headers, path)

    return inner if LocalRun else jobserver.submit(_fetch_url_, (bookid, headers, path),(),('time','urllib2','pp'))

def slicer(seq, size):
    return (seq[pos:pos + size] for pos in xrange(0, len(seq), size))

def rmall(filepath):
    if os.path.isdir(filepath):
        shutil.rmtree(filepath)
    else:
        os.remove(filepath)

def uploadToS3(filepath):
    if not bool(s3info): # test if the dictionary is empty
        print "S3 credentials not defined... not uploading", filepath
        return False

    destination = os.path.basename(filepath)
    cmd = "/usr/local/bin/s3cmd --access_key={0} --secret_key={1} --mime-type=application/zip --no-progress --quiet put {2} s3://{3}/{4}".format(
            s3info['accesskey'], s3info['secretkey'], filepath, s3info['bucket'], destination)
    exitcode = os.system(cmd)
    print "uploading",filepath,"to S3"
    return (exitcode == 0)

########## main ##########

starttime = time.gmtime()
print "Start Time:", time.strftime("%Y-%m-%d %H:%M:%S UTC", starttime)

book_ids = args.isbns.read().splitlines()
book_ids = map(lambda x: string.strip(x), book_ids)

# start servers
totalServers = 1
if LocalRun:
    job_server = pp.Server()
else:
    ppservers = tuple(args.nodes.read().splitlines())
    totalServers = len(ppservers)+1 # +1 for the local node
    job_server = pp.Server(ppservers=ppservers) if (args.secretkey == None) else pp.Server(ppservers=ppservers, secret=args.secretkey)

print "Starting pp with", job_server.get_ncpus(), "cpus"
currentNodeCount = len(job_server.get_active_nodes())
print "We have", currentNodeCount, "nodes"
while (currentNodeCount < totalServers):
    print "Expecting {0} cluster nodes, have {1}, waiting for others to join...".format(totalServers, currentNodeCount)
    time.sleep(2)
    currentNodeCount = len(job_server.get_active_nodes())
print "All expected nodes have joined"

# set up the temp directory
try:
    os.makedirs(basePath)
except OSError as exc: # Python >2.5
    if exc.errno == errno.EEXIST and os.path.isdir(basePath):
        # if the directory is already there, delete all contents
        files=os.listdir(basePath)
        for f in files:
            rmall(basePath + f)
    else: raise

# grab the books 1000 at a time and fetch the urls; wait for the results and then zip them up
batchSize = 10 if LocalRun else 1000
print "Total number of books to process is",len(book_ids)

failedBooksPath = basePath + str(time.mktime(time.gmtime())).replace('.','_') + '_incompleteBooks'
bookManifestPath = basePath + str(time.mktime(time.gmtime())).replace('.','_') + '_book_manifest'
count = 0

for bslice in slicer(book_ids, batchSize):
    failedBooks = []
    succeededBooks = []
    jobs=[]
    sliceName = str(min(bslice)) + '_' + str(max(bslice)) + '_' + str(count)
    zpath = basePath + sliceName +  '.zip'
    outputPath = basePath + sliceName + '/'
    processedBooks = []

    # Kick off the jobs
    jobs = map(lambda x: (x, startJob(x, job_server, outputPath)), bslice)

    # set up the slice's directory
    try:
        os.makedirs(outputPath)
    except OSError as exc: # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(outputPath):
            # if the directory is already there, delete all contents
            files=os.listdir(outputPath)
            for f in files:
                rmall(outputPath + f)
            pass
        else:
            raise exc

    # block waiting for all the jobs to complete
    for (bkid,job) in jobs:
        success = False
        try:
            res = job()
            bkid = res[0]
            success = res[1]
        except (KeyboardInterrupt, SystemExit):
            print "Caught keyboard interrupt/system exit"
            raise
        except Exception, e:
            print "Encountered exception when executing job: ", e

        lst = processedBooks if success else failedBooks
        lst.append(bkid)

    time.sleep(2)

    if (len(processedBooks) > 0):
        # zip everything up
        zout = zipfile.ZipFile(zpath,'w',zipfile.ZIP_DEFLATED)
        for root,dirs,files in os.walk(outputPath):
            for f in files:
                zout.write(os.path.join(root, f))
        zout.close()

        #delete all the files we just zipped
        rmall(outputPath)

        uploadSuccess = uploadToS3(zpath)

        # if we stored it in s3, delete the zipfile and the slice's directory
        if (uploadSuccess or LocalRun):
            if not LocalRun: # if it's a local run, leave the zip file where it is
                rmall(zpath)
            for b in processedBooks:
                succeededBooks.append((b, zpath))
        else:
            failedBooks = failedBooks + processedBooks

    # write to the appropriate files
    if (len(failedBooks) > 0):
        with open(failedBooksPath , "a") as local_file:
            for bkid in failedBooks:
              local_file.write(str(bkid) + "\n")

    if (len(succeededBooks) > 0):
        with open(bookManifestPath, "a") as local_file:
            for (bkid, zp) in succeededBooks:
              local_file.write("{0},{1}\n".format(bkid, os.path.basename(zp)))

    print "Finished batch of size", len(bslice)
    count += len(bslice)

if (os.path.isfile(failedBooksPath)):
    upsucc = uploadToS3(failedBooksPath)
    if (upsucc):
        rmall(failedBooksPath)

if (os.path.isfile(bookManifestPath)):
    upsucc = uploadToS3(bookManifestPath)
    if (upsucc):
        rmall(bookManifestPath)

endtime = time.gmtime()
print "End Time:", time.strftime("%Y-%m-%d %H:%M:%S UTC", endtime)
print "Total Time:", time.mktime(endtime)-time.mktime(starttime)
<commit_msg>Changing exponential backoff to be linear since we are still getting good results with the more aggressive retry<commit_after>#!/usr/bin/python

# Note: python >= 2.7 required
# Note: you will need s3cmd installed from http://s3tools.org/, I was using 1.5.0-rc1 (http://sourceforge.net/projects/s3tools/files/s3cmd/1.5.0-rc1/s3cmd-1.5.0-rc1.tar.gz)

import pp
import time         # for sleep
import os
import shutil
import sys
import string
import urllib2
import random
import errno
import argparse
import zipfile

parser = argparse.ArgumentParser(description='Crawl a set of ISBNs on amazon.com')
parser.add_argument('-i','--isbns', type=file, required=True, help='The file that contains the ISBNs to crawl')
parser.add_argument('-n','--nodes', type=file, help='The file that contains the cluster slave nodes to use')
parser.add_argument('-k','--secretkey', type=str, help='The secret key used to communicate with the cluster slave nodes')
parser.add_argument('-s','--s3credentials', type=file, help='The file containing the S3 credentials required to upload:\n\tbucket=<bucketname>\n\taccessKey=<access key>\n\tsecretKey=<secret key>')
args = parser.parse_args()

LocalRun = (args.nodes == None)
if LocalRun:
    print "No slave nodes provided, assuming you just want to run on this machine and this is a LOCAL RUN"

s3info = {}
if (args.s3credentials != None):
    s3lines = args.s3credentials.read().splitlines()
    for line in s3lines:
        kv = line.split('=')
        s3info[str.lower(kv[0])] = kv[1]
else:
    print "No S3 credentials provided, we will store the files locally"

userAgents = [
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/37.0.2062.94 Safari/537.36",
    "Mozilla/5.0 (compatible; MSIE 10.0; Windows NT 6.2; Trident/6.0)",
    "Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:25.0) Gecko/20100101 Firefox/29.0"
]
basePath = '/tmp/books/' if LocalRun else '/vol/temp_data/'
basePath = basePath + ''.join(random.choice(string.ascii_uppercase) for i in range(8)) + '/'

# download the contents of a url to the book's directory

#def _dummy_job(bookid, headers, fpath): # useful for debugging to short circuit the fetching
def _fetch_url_ (bookid, headers, fpath):
    url = 'http://www.amazon.com/dp/'+str(bookid)
    captchaText = '<form method="get" action="/errors/validateCaptcha" name="">'
    success = False
    failfast = False
    maxTries = 5
    ntries = 1 # 0 doesn't let me multiply well
    sleep = 2
    request = urllib2.Request(url,None,headers)

    while not success and not failfast and ntries <= maxTries:
        text = ''
        try:
            response = urllib2.urlopen(request)
            if response.getcode() == 200:
                text = response.read()
                if captchaText in text:
                    print "Got captcha, sleeping and retrying"
                else:
                    with open(fpath, 'wb') as f:
                        f.write(text)
                    success = True
            else:
                print "Received {0} code when retrieving {1}".format(response.getcode(), url)
        except urllib2.HTTPError, e:
            print "When retrieving {0} received HTTP Error: {1}".format(url, e.code)
            if e.code == 404:
                print "received 404, so skipping further retries"
                failfast = True
        except urllib2.URLError, e:
            print "URL Error:", e.reason , url
        except IOError, e:
            print "Could not write to file {0} when fetching {1}... skipping this book".format(fpath,url)
            print "Detailed Exception: ", e
            failfast = True
        finally:
            try:
                response.close()
            except NameError:
                pass

        if success:
            time.sleep(sleep) # sleep before trying next task
        else:
            if not failfast:
                # exponential backoff
                print "Attempt {0} for url {1} failed, backing off {2} seconds".format(ntries, url, sleep*ntries)
                time.sleep(sleep*ntries)
        ntries+=1
    return (bookid, success)

#def _fetch_url_ (bookid, headers, fpath):
def _dummy_job(bookid, headers, fpath):
    time.sleep(1)
    text = 'Fake book data'
    return (bookid, text)

def startJob(bookid,jobserver, outputPath):
    path = outputPath + str(bookid)
    index = random.randint(0,len(userAgents)-1)
    headers = {
        'User-Agent' : userAgents[index],
        'Accept' : 'text/html,application/xhtml+xml,application/xml;q=0.9',
        'Accept-Language' : 'en-US,en;q=0.8'
    }
    def inner():
        return _fetch_url_(bookid, headers, path)

    return inner if LocalRun else jobserver.submit(_fetch_url_, (bookid, headers, path),(),('time','urllib2','pp'))

def slicer(seq, size):
    return (seq[pos:pos + size] for pos in xrange(0, len(seq), size))

def rmall(filepath):
    if os.path.isdir(filepath):
        shutil.rmtree(filepath)
    else:
        os.remove(filepath)

def uploadToS3(filepath):
    if not bool(s3info): # test if the dictionary is empty
        print "S3 credentials not defined... not uploading", filepath
        return False

    destination = os.path.basename(filepath)
    cmd = "/usr/local/bin/s3cmd --access_key={0} --secret_key={1} --mime-type=application/zip --no-progress --quiet put {2} s3://{3}/{4}".format(
            s3info['accesskey'], s3info['secretkey'], filepath, s3info['bucket'], destination)
    exitcode = os.system(cmd)
    print "uploading",filepath,"to S3"
    return (exitcode == 0)

########## main ##########

starttime = time.gmtime()
print "Start Time:", time.strftime("%Y-%m-%d %H:%M:%S UTC", starttime)

book_ids = args.isbns.read().splitlines()
book_ids = map(lambda x: string.strip(x), book_ids)

# start servers
totalServers = 1
if LocalRun:
    job_server = pp.Server()
else:
    ppservers = tuple(args.nodes.read().splitlines())
    totalServers = len(ppservers)+1 # +1 for the local node
    job_server = pp.Server(ppservers=ppservers) if (args.secretkey == None) else pp.Server(ppservers=ppservers, secret=args.secretkey)

print "Starting pp with", job_server.get_ncpus(), "cpus"
currentNodeCount = len(job_server.get_active_nodes())
print "We have", currentNodeCount, "nodes"
while (currentNodeCount < totalServers):
    print "Expecting {0} cluster nodes, have {1}, waiting for others to join...".format(totalServers, currentNodeCount)
    time.sleep(2)
    currentNodeCount = len(job_server.get_active_nodes())
print "All expected nodes have joined"

# set up the temp directory
try:
    os.makedirs(basePath)
except OSError as exc: # Python >2.5
    if exc.errno == errno.EEXIST and os.path.isdir(basePath):
        # if the directory is already there, delete all contents
        files=os.listdir(basePath)
        for f in files:
            rmall(basePath + f)
    else: raise

# grab the books 1000 at a time and fetch the urls; wait for the results and then zip them up
batchSize = 10 if LocalRun else 1000
print "Total number of books to process is",len(book_ids)

failedBooksPath = basePath + str(time.mktime(time.gmtime())).replace('.','_') + '_incompleteBooks'
bookManifestPath = basePath + str(time.mktime(time.gmtime())).replace('.','_') + '_book_manifest'
count = 0

for bslice in slicer(book_ids, batchSize):
    failedBooks = []
    succeededBooks = []
    jobs=[]
    sliceName = str(min(bslice)) + '_' + str(max(bslice)) + '_' + str(count)
    zpath = basePath + sliceName +  '.zip'
    outputPath = basePath + sliceName + '/'
    processedBooks = []

    # Kick off the jobs
    jobs = map(lambda x: (x, startJob(x, job_server, outputPath)), bslice)

    # set up the slice's directory
    try:
        os.makedirs(outputPath)
    except OSError as exc: # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(outputPath):
            # if the directory is already there, delete all contents
            files=os.listdir(outputPath)
            for f in files:
                rmall(outputPath + f)
            pass
        else:
            raise exc

    # block waiting for all the jobs to complete
    for (bkid,job) in jobs:
        success = False
        try:
            res = job()
            bkid = res[0]
            success = res[1]
        except (KeyboardInterrupt, SystemExit):
            print "Caught keyboard interrupt/system exit"
            raise
        except Exception, e:
            print "Encountered exception when executing job: ", e

        lst = processedBooks if success else failedBooks
        lst.append(bkid)

    time.sleep(2)

    if (len(processedBooks) > 0):
        # zip everything up
        zout = zipfile.ZipFile(zpath,'w',zipfile.ZIP_DEFLATED)
        for root,dirs,files in os.walk(outputPath):
            for f in files:
                zout.write(os.path.join(root, f))
        zout.close()

        #delete all the files we just zipped
        rmall(outputPath)

        uploadSuccess = uploadToS3(zpath)

        # if we stored it in s3, delete the zipfile and the slice's directory
        if (uploadSuccess or LocalRun):
            if not LocalRun: # if it's a local run, leave the zip file where it is
                rmall(zpath)
            for b in processedBooks:
                succeededBooks.append((b, zpath))
        else:
            failedBooks = failedBooks + processedBooks

    # write to the appropriate files
    if (len(failedBooks) > 0):
        with open(failedBooksPath , "a") as local_file:
            for bkid in failedBooks:
              local_file.write(str(bkid) + "\n")

    if (len(succeededBooks) > 0):
        with open(bookManifestPath, "a") as local_file:
            for (bkid, zp) in succeededBooks:
              local_file.write("{0},{1}\n".format(bkid, os.path.basename(zp)))

    print "Finished batch of size", len(bslice)
    count += len(bslice)

if (os.path.isfile(failedBooksPath)):
    upsucc = uploadToS3(failedBooksPath)
    if (upsucc):
        rmall(failedBooksPath)

if (os.path.isfile(bookManifestPath)):
    upsucc = uploadToS3(bookManifestPath)
    if (upsucc):
        rmall(bookManifestPath)

endtime = time.gmtime()
print "End Time:", time.strftime("%Y-%m-%d %H:%M:%S UTC", endtime)
print "Total Time:", time.mktime(endtime)-time.mktime(starttime)
<|endoftext|>


<commit_before>#!/usr/bin/python2
# Copyright 2016 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# A script which builds, flashes, and runs EC CTS
#
# Software prerequisites:
# - openocd version 0.10 or above
# - lsusb
# - udevadm
#
# To try it out, hook two boards (DEFAULT_TH and DEFAULT_DUT) with USB cables
# to the host and execute the script:
#   $ ./cts.py --debug
# It'll run mock tests. The result will be stored in CTS_TEST_RESULT_DIR.

import argparse
import collections
import os
import time
import common.board as board
from copy import deepcopy
import xml.etree.ElementTree as et
from twisted.python.syslog import DEFAULT_FACILITY

CTS_CORRUPTED_CODE = -2  # The test didn't execute correctly
CTS_CONFLICTING_CODE = -1 # Error codes should never conflict
CTS_SUCCESS_CODE = 0
CTS_COLOR_RED = '#fb7d7d'
CTS_COLOR_GREEN = '#7dfb9f'
DEFAULT_TH = 'stm32l476g-eval'
DEFAULT_DUT = 'nucleo-f072rb'
MAX_SUITE_TIME_SEC = 3
CTS_DEBUG_START = '[DEBUG]'
CTS_DEBUG_END = '[DEBUG_END]'
CTS_TEST_RESULT_DIR = '/tmp/cts'

class Cts(object):
  """Class that represents a CTS testing setup and provides
  interface to boards (building, flashing, etc.)

  Attributes:
    dut: DeviceUnderTest object representing dut
    th: TestHarness object representing th
    module: Name of module to build/run tests for
    test_names: List of strings of test names contained in given module
    test_results: Dictionary of results of each test from module, with
        keys being test name strings and values being test result integers
    return_codes: Dict of strings of return codes, with a code's integer
      value being the index for the corresponding string representation
    debug: Boolean that indicates whether or not on-board debug message
      printing should be enabled when building.
    debug_output: Dictionary mapping test name to an array contain debug
      messages sent while it was running
  """

  def __init__(self, ec_dir, dut, module, debug=False):
    """Initializes cts class object with given arguments.

    Args:
      dut: Name of Device Under Test (DUT) board
      ec_dir: String path to ec directory
      dut: Name of board to use for DUT
      module: Name of module to build/run tests for
      debug: Boolean that indicates whether or not on-board debug message
        printing should be enabled.
    """
    self.results_dir = CTS_TEST_RESULT_DIR
    self.ec_dir = ec_dir
    self.module = module
    self.debug = debug
    serial_path = os.path.join(self.ec_dir, 'build', 'cts_th_serial')
    self.th = board.TestHarness(DEFAULT_TH, serial_path)
    self.dut = board.DeviceUnderTest(dut, self.th)
    cts_dir = os.path.join(self.ec_dir, 'cts')
    testlist_path = os.path.join(cts_dir, self.module, 'cts.testlist')
    self.test_names = Cts.get_macro_args(testlist_path, 'CTS_TEST')

    self.debug_output = {}
    for test in self.test_names:
      self.debug_output[test] = []

    return_codes_path = os.path.join(cts_dir, 'common', 'cts.rc')
    self.return_codes = dict(enumerate(Cts.get_macro_args(
        return_codes_path, 'CTS_RC_')))

    self.return_codes[CTS_CONFLICTING_CODE] = 'RESULTS CONFLICT'
    self.return_codes[CTS_CORRUPTED_CODE] = 'CORRUPTED'
    self.test_results = collections.OrderedDict()

  def build(self):
    """Build images for DUT and TH"""
    if self.dut.build(self.module, self.ec_dir, self.debug):
      raise RuntimeError('Building module %s for DUT failed' % (self.module))
    if self.th.build(self.module, self.ec_dir, self.debug):
      raise RuntimeError('Building module %s for TH failed' % (self.module))

  def flash_boards(self):
    """Flashes th and dut boards with their most recently build ec.bin"""
    cts_module = 'cts_' + self.module
    image_path = os.path.join('build', self.th.board, cts_module, 'ec.bin')
    self.identify_boards()
    print 'Flashing TH with', image_path
    if self.th.flash(image_path):
      raise RuntimeError('Flashing TH failed')
    image_path = os.path.join('build', self.dut.board, cts_module, 'ec.bin')
    print 'Flashing DUT with', image_path
    if self.dut.flash(image_path):
      raise RuntimeError('Flashing DUT failed')

  def setup(self):
    """Setup boards"""
    self.th.save_serial()

  def identify_boards(self):
    """Updates serials of both th and dut, in that order (order matters)"""
    self.th.get_serial()
    self.dut.get_serial()

  def reset_boards(self):
    """Resets the boards and allows them to run tests
    Due to current (7/27/16) version of sync function,
    both boards must be rest and halted, with the th
    resuming first, in order for the test suite to run
    in sync
    """
    self.identify_boards()
    self.th.send_open_ocd_commands(['init', 'reset halt'])
    self.dut.send_open_ocd_commands(['init', 'reset halt'])
    self.th.send_open_ocd_commands(['init', 'resume'])
    self.dut.send_open_ocd_commands(['init', 'resume'])

  @staticmethod
  def get_macro_args(filepath, macro):
    """Get list of args of a certain macro in a file when macro is used
    by itself on a line

    Args:
      filepath: String containing absolute path to the file
      macro: String containing text of macro to get args of
    """
    args = []
    with open(filepath, 'r') as f:
      for l in f.readlines():
        if not l.strip().startswith(macro):
          continue
        l = l.strip()[len(macro):]
        args.append(l.strip('()').replace(',', ''))
    return args

  def extract_debug_output(self, output):
    """Append the debug messages from output to self.debug_output

    Args:
      output: String containing output from which to extract debug
        messages
    """
    lines = [ln.strip() for ln in output.split('\n')]
    test_num = 0
    i = 0
    message_buf = []
    while i < len(lines):
      if test_num >= len(self.test_names):
        break
      if lines[i].strip() == CTS_DEBUG_START:
        i += 1
        msg = ''
        while i < len(lines):
          if lines[i] == CTS_DEBUG_END:
            break
          else:
            msg += lines[i] + '\n'
            i += 1
        message_buf.append(msg)
      else:
        current_test = self.test_names[test_num]
        if lines[i].strip().startswith(current_test):
          self.debug_output[current_test] += message_buf
          message_buf = []
          test_num += 1
      i += 1

  def parse_output(self, r1, r2):
    """Parse the outputs of the DUT and TH together

    Args;
      r1: String output of one of the DUT or the TH (order does not matter)
      r2: String output of one of the DUT or the TH (order does not matter)
    """
    self.test_results.clear()  # empty out any old results

    first_corrupted_test = len(self.test_names)

    self.extract_debug_output(r1)
    self.extract_debug_output(r2)

    for output_str in [r1, r2]:
      test_num = 0
      for ln in [ln.strip() for ln in output_str.split('\n')]:
        tokens = ln.split()
        if len(tokens) != 2:
          continue
        test = tokens[0].strip()
        if test not in self.test_names:
          continue
        try:
          return_code = int(tokens[1])
        except ValueError: # Second token is not an int
          continue
        if test != self.test_names[test_num]:
          first_corrupted_test = test_num
          break # Results after this test are corrupted
        elif self.test_results.get(
            test,
            CTS_SUCCESS_CODE) == CTS_SUCCESS_CODE:
          self.test_results[test] = return_code
        elif return_code == CTS_SUCCESS_CODE:
          pass
        elif return_code != self.test_results[test]:
          self.test_results[test] = CTS_CONFLICTING_CODE
        test_num += 1

      if test_num != len(self.test_names): # If a suite didn't finish
        first_corrupted_test = min(first_corrupted_test, test_num)

    if first_corrupted_test < len(self.test_names):
      for test in self.test_names[first_corrupted_test:]:
        self.test_results[test] = CTS_CORRUPTED_CODE

  def _results_as_string(self):
    """Takes saved results and returns a duplicate of their dictionary
    with the return codes replaces with their string representation

    Returns:
      dictionary with test name strings as keys and test result strings
          as values
    """
    result = deepcopy(self.test_results)
    # Convert codes to strings
    for test, code in result.items():
        result[test] = self.return_codes.get(code, 'UNKNOWN %d' % code)
    return result

  def prettify_results(self):
    """Takes saved results and returns a string representation of them

    Return: Dictionary similar to self.test_results, but with strings
        instead of error codes
    """
    res = self._results_as_string()
    t_long = max(len(s) for s in res.keys())
    e_max_len = max(len(s) for s in res.values())

    pretty_results = 'CTS Test Results for ' + self.module + ' module:\n'

    for test, code in res.items():
      align_str = '\n{0:<' + str(t_long) + \
        '} {1:>' + str(e_max_len) + '}'
      pretty_results += align_str.format(test, code)

    return pretty_results

  def results_as_html(self):
    res = self._results_as_string()
    root = et.Element('html')
    head = et.SubElement(root, 'head')
    style = et.SubElement(head, 'style')
    style.text = ('table, td, th {border: 1px solid black;}'
                  'body {font-family: \"Lucida Console\", Monaco, monospace')
    body = et.SubElement(root, 'body')
    table = et.SubElement(body, 'table')
    table.set('style','width:100%')
    title_row = et.SubElement(table, 'tr')
    test_name_title = et.SubElement(title_row, 'th')
    test_name_title.text = 'Test Name'
    test_name_title.set('style', 'white-space : nowrap')
    test_results_title = et.SubElement(title_row, 'th')
    test_results_title.text = 'Test Result'
    test_results_title.set('style', 'white-space : nowrap')
    test_debug_title = et.SubElement(title_row, 'th')
    test_debug_title.text = 'Debug Output'
    test_debug_title.set('style', 'width:99%')

    for name, result in res.items():
      row = et.SubElement(table, 'tr')
      name_e = et.SubElement(row, 'td')
      name_e.text = name
      name_e.set('style', 'white-space : nowrap')
      result_e = et.SubElement(row, 'td')
      result_e.text = result
      result_e.set('style', 'white-space : nowrap')
      debug_e = et.SubElement(row, 'td')
      debug_e.set('style', 'width:99%')
      debug_e.set('style', 'white-space : pre-wrap')
      if len(self.debug_output[name]) == 0:
        debug_e.text = 'None'
      else:
        combined_message = ''
        for msg in self.debug_output[name]:
          combined_message += msg
        combined_message = combined_message
        debug_e.text = combined_message
      if result == self.return_codes[CTS_SUCCESS_CODE]:
        result_e.set('bgcolor', CTS_COLOR_GREEN)
      else:
        result_e.set('bgcolor', CTS_COLOR_RED)

    return et.tostring(root, method='html')

  def run(self):
    """Resets boards, records test results in results dir"""
    self.identify_boards()
    self.dut.setup_tty()
    self.th.setup_tty()

    # clear buffers
    self.dut.read_tty()
    self.th.read_tty()

    self.reset_boards()

    time.sleep(MAX_SUITE_TIME_SEC)

    dut_results = self.dut.read_tty()
    th_results = self.th.read_tty()

    if not dut_results or not th_results:
      raise ValueError('Output missing from boards. If you have a process '
                       'reading ttyACMx, please kill that process and try '
                       'again.')

    self.parse_output(dut_results, th_results)
    pretty_results = self.prettify_results()
    html_results = self.results_as_html()

    dest = os.path.join(self.results_dir, self.dut.board, self.module + '.html')
    if not os.path.exists(os.path.dirname(dest)):
      os.makedirs(os.path.dirname(dest))

    with open(dest, 'w') as fl:
      fl.write(html_results)

    print pretty_results

def main():
  """Main entry point for CTS script from command line"""
  ec_dir = os.path.realpath(os.path.join(
      os.path.dirname(os.path.abspath(__file__)), '..'))
  os.chdir(ec_dir)

  dut = DEFAULT_DUT
  module = 'meta'

  parser = argparse.ArgumentParser(description='Used to build/flash boards')
  parser.add_argument('-d',
                      '--dut',
                      help='Specify DUT you want to build/flash')
  parser.add_argument('-m',
                      '--module',
                      help='Specify module you want to build/flash')
  parser.add_argument('--debug',
                      action='store_true',
                      help=('If building, build with debug printing enabled. '
                            'This may change test results'))
  parser.add_argument('-s',
                      '--setup',
                      action='store_true',
                      help='Connect only the TH to save its serial')
  parser.add_argument('-b',
                      '--build',
                      action='store_true',
                      help='Build test suite (no flashing)')
  parser.add_argument('-f',
                      '--flash',
                      action='store_true',
                      help='Flash boards with most recent images')
  parser.add_argument('-r',
                      '--run',
                      action='store_true',
                      help='Run tests without flashing')

  args = parser.parse_args()

  if args.module:
    module = args.module

  if args.dut:
    dut = args.dut

  cts = Cts(ec_dir, dut=dut, module=module, debug=args.debug)

  if args.setup:
    cts.setup()
  elif args.build:
    cts.build()
  elif args.flash:
    cts.flash_boards()
  elif args.run:
    cts.run()
  else:
    cts.build()
    cts.flash_boards()
    cts.run()

if __name__ == "__main__":
  main()
<commit_msg>eCTS: Pause a few seconds before flushing tty<commit_after>#!/usr/bin/python2
# Copyright 2016 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# A script which builds, flashes, and runs EC CTS
#
# Software prerequisites:
# - openocd version 0.10 or above
# - lsusb
# - udevadm
#
# To try it out, hook two boards (DEFAULT_TH and DEFAULT_DUT) with USB cables
# to the host and execute the script:
#   $ ./cts.py --debug
# It'll run mock tests. The result will be stored in CTS_TEST_RESULT_DIR.

import argparse
import collections
import os
import time
import common.board as board
from copy import deepcopy
import xml.etree.ElementTree as et
from twisted.python.syslog import DEFAULT_FACILITY

CTS_CORRUPTED_CODE = -2  # The test didn't execute correctly
CTS_CONFLICTING_CODE = -1 # Error codes should never conflict
CTS_SUCCESS_CODE = 0
CTS_COLOR_RED = '#fb7d7d'
CTS_COLOR_GREEN = '#7dfb9f'
DEFAULT_TH = 'stm32l476g-eval'
DEFAULT_DUT = 'nucleo-f072rb'
MAX_SUITE_TIME_SEC = 3
CTS_DEBUG_START = '[DEBUG]'
CTS_DEBUG_END = '[DEBUG_END]'
CTS_TEST_RESULT_DIR = '/tmp/cts'

class Cts(object):
  """Class that represents a CTS testing setup and provides
  interface to boards (building, flashing, etc.)

  Attributes:
    dut: DeviceUnderTest object representing dut
    th: TestHarness object representing th
    module: Name of module to build/run tests for
    test_names: List of strings of test names contained in given module
    test_results: Dictionary of results of each test from module, with
        keys being test name strings and values being test result integers
    return_codes: Dict of strings of return codes, with a code's integer
      value being the index for the corresponding string representation
    debug: Boolean that indicates whether or not on-board debug message
      printing should be enabled when building.
    debug_output: Dictionary mapping test name to an array contain debug
      messages sent while it was running
  """

  def __init__(self, ec_dir, dut, module, debug=False):
    """Initializes cts class object with given arguments.

    Args:
      dut: Name of Device Under Test (DUT) board
      ec_dir: String path to ec directory
      dut: Name of board to use for DUT
      module: Name of module to build/run tests for
      debug: Boolean that indicates whether or not on-board debug message
        printing should be enabled.
    """
    self.results_dir = CTS_TEST_RESULT_DIR
    self.ec_dir = ec_dir
    self.module = module
    self.debug = debug
    serial_path = os.path.join(self.ec_dir, 'build', 'cts_th_serial')
    self.th = board.TestHarness(DEFAULT_TH, serial_path)
    self.dut = board.DeviceUnderTest(dut, self.th)
    cts_dir = os.path.join(self.ec_dir, 'cts')
    testlist_path = os.path.join(cts_dir, self.module, 'cts.testlist')
    self.test_names = Cts.get_macro_args(testlist_path, 'CTS_TEST')

    self.debug_output = {}
    for test in self.test_names:
      self.debug_output[test] = []

    return_codes_path = os.path.join(cts_dir, 'common', 'cts.rc')
    self.return_codes = dict(enumerate(Cts.get_macro_args(
        return_codes_path, 'CTS_RC_')))

    self.return_codes[CTS_CONFLICTING_CODE] = 'RESULTS CONFLICT'
    self.return_codes[CTS_CORRUPTED_CODE] = 'CORRUPTED'
    self.test_results = collections.OrderedDict()

  def build(self):
    """Build images for DUT and TH"""
    if self.dut.build(self.module, self.ec_dir, self.debug):
      raise RuntimeError('Building module %s for DUT failed' % (self.module))
    if self.th.build(self.module, self.ec_dir, self.debug):
      raise RuntimeError('Building module %s for TH failed' % (self.module))

  def flash_boards(self):
    """Flashes th and dut boards with their most recently build ec.bin"""
    cts_module = 'cts_' + self.module
    image_path = os.path.join('build', self.th.board, cts_module, 'ec.bin')
    self.identify_boards()
    print 'Flashing TH with', image_path
    if self.th.flash(image_path):
      raise RuntimeError('Flashing TH failed')
    image_path = os.path.join('build', self.dut.board, cts_module, 'ec.bin')
    print 'Flashing DUT with', image_path
    if self.dut.flash(image_path):
      raise RuntimeError('Flashing DUT failed')

  def setup(self):
    """Setup boards"""
    self.th.save_serial()

  def identify_boards(self):
    """Updates serials of both th and dut, in that order (order matters)"""
    self.th.get_serial()
    self.dut.get_serial()

  @staticmethod
  def get_macro_args(filepath, macro):
    """Get list of args of a certain macro in a file when macro is used
    by itself on a line

    Args:
      filepath: String containing absolute path to the file
      macro: String containing text of macro to get args of
    """
    args = []
    with open(filepath, 'r') as f:
      for l in f.readlines():
        if not l.strip().startswith(macro):
          continue
        l = l.strip()[len(macro):]
        args.append(l.strip('()').replace(',', ''))
    return args

  def extract_debug_output(self, output):
    """Append the debug messages from output to self.debug_output

    Args:
      output: String containing output from which to extract debug
        messages
    """
    lines = [ln.strip() for ln in output.split('\n')]
    test_num = 0
    i = 0
    message_buf = []
    while i < len(lines):
      if test_num >= len(self.test_names):
        break
      if lines[i].strip() == CTS_DEBUG_START:
        i += 1
        msg = ''
        while i < len(lines):
          if lines[i] == CTS_DEBUG_END:
            break
          else:
            msg += lines[i] + '\n'
            i += 1
        message_buf.append(msg)
      else:
        current_test = self.test_names[test_num]
        if lines[i].strip().startswith(current_test):
          self.debug_output[current_test] += message_buf
          message_buf = []
          test_num += 1
      i += 1

  def parse_output(self, r1, r2):
    """Parse the outputs of the DUT and TH together

    Args;
      r1: String output of one of the DUT or the TH (order does not matter)
      r2: String output of one of the DUT or the TH (order does not matter)
    """
    self.test_results.clear()  # empty out any old results

    first_corrupted_test = len(self.test_names)

    self.extract_debug_output(r1)
    self.extract_debug_output(r2)

    for output_str in [r1, r2]:
      test_num = 0
      for ln in [ln.strip() for ln in output_str.split('\n')]:
        tokens = ln.split()
        if len(tokens) != 2:
          continue
        test = tokens[0].strip()
        if test not in self.test_names:
          continue
        try:
          return_code = int(tokens[1])
        except ValueError: # Second token is not an int
          continue
        if test != self.test_names[test_num]:
          first_corrupted_test = test_num
          break # Results after this test are corrupted
        elif self.test_results.get(
            test,
            CTS_SUCCESS_CODE) == CTS_SUCCESS_CODE:
          self.test_results[test] = return_code
        elif return_code == CTS_SUCCESS_CODE:
          pass
        elif return_code != self.test_results[test]:
          self.test_results[test] = CTS_CONFLICTING_CODE
        test_num += 1

      if test_num != len(self.test_names): # If a suite didn't finish
        first_corrupted_test = min(first_corrupted_test, test_num)

    if first_corrupted_test < len(self.test_names):
      for test in self.test_names[first_corrupted_test:]:
        self.test_results[test] = CTS_CORRUPTED_CODE

  def _results_as_string(self):
    """Takes saved results and returns a duplicate of their dictionary
    with the return codes replaces with their string representation

    Returns:
      dictionary with test name strings as keys and test result strings
          as values
    """
    result = deepcopy(self.test_results)
    # Convert codes to strings
    for test, code in result.items():
        result[test] = self.return_codes.get(code, 'UNKNOWN %d' % code)
    return result

  def prettify_results(self):
    """Takes saved results and returns a string representation of them

    Return: Dictionary similar to self.test_results, but with strings
        instead of error codes
    """
    res = self._results_as_string()
    t_long = max(len(s) for s in res.keys())
    e_max_len = max(len(s) for s in res.values())

    pretty_results = 'CTS Test Results for ' + self.module + ' module:\n'

    for test, code in res.items():
      align_str = '\n{0:<' + str(t_long) + \
        '} {1:>' + str(e_max_len) + '}'
      pretty_results += align_str.format(test, code)

    return pretty_results

  def results_as_html(self):
    res = self._results_as_string()
    root = et.Element('html')
    head = et.SubElement(root, 'head')
    style = et.SubElement(head, 'style')
    style.text = ('table, td, th {border: 1px solid black;}'
                  'body {font-family: \"Lucida Console\", Monaco, monospace')
    body = et.SubElement(root, 'body')
    table = et.SubElement(body, 'table')
    table.set('style','width:100%')
    title_row = et.SubElement(table, 'tr')
    test_name_title = et.SubElement(title_row, 'th')
    test_name_title.text = 'Test Name'
    test_name_title.set('style', 'white-space : nowrap')
    test_results_title = et.SubElement(title_row, 'th')
    test_results_title.text = 'Test Result'
    test_results_title.set('style', 'white-space : nowrap')
    test_debug_title = et.SubElement(title_row, 'th')
    test_debug_title.text = 'Debug Output'
    test_debug_title.set('style', 'width:99%')

    for name, result in res.items():
      row = et.SubElement(table, 'tr')
      name_e = et.SubElement(row, 'td')
      name_e.text = name
      name_e.set('style', 'white-space : nowrap')
      result_e = et.SubElement(row, 'td')
      result_e.text = result
      result_e.set('style', 'white-space : nowrap')
      debug_e = et.SubElement(row, 'td')
      debug_e.set('style', 'width:99%')
      debug_e.set('style', 'white-space : pre-wrap')
      if len(self.debug_output[name]) == 0:
        debug_e.text = 'None'
      else:
        combined_message = ''
        for msg in self.debug_output[name]:
          combined_message += msg
        combined_message = combined_message
        debug_e.text = combined_message
      if result == self.return_codes[CTS_SUCCESS_CODE]:
        result_e.set('bgcolor', CTS_COLOR_GREEN)
      else:
        result_e.set('bgcolor', CTS_COLOR_RED)

    return et.tostring(root, method='html')

  def run(self):
    """Resets boards, records test results in results dir"""
    self.identify_boards()
    self.dut.setup_tty()
    self.th.setup_tty()

    # Boards might be still writing to tty. Wait a few seconds before flashing.
    time.sleep(3)

    # clear buffers
    self.dut.read_tty()
    self.th.read_tty()

    # Resets the boards and allows them to run tests
    # Due to current (7/27/16) version of sync function,
    # both boards must be rest and halted, with the th
    # resuming first, in order for the test suite to run in sync
    print 'Halting TH...'
    self.th.send_open_ocd_commands(['init', 'reset halt'])
    print 'Resetting DUT...'
    self.dut.send_open_ocd_commands(['init', 'reset halt'])
    print 'Resuming TH...'
    self.th.send_open_ocd_commands(['init', 'resume'])
    print 'Resuming DUT...'
    self.dut.send_open_ocd_commands(['init', 'resume'])

    time.sleep(MAX_SUITE_TIME_SEC)

    dut_results = self.dut.read_tty()
    th_results = self.th.read_tty()

    if not dut_results or not th_results:
      raise ValueError('Output missing from boards. If you have a process '
                       'reading ttyACMx, please kill that process and try '
                       'again.')

    self.parse_output(dut_results, th_results)
    pretty_results = self.prettify_results()
    html_results = self.results_as_html()

    dest = os.path.join(self.results_dir, self.dut.board, self.module + '.html')
    if not os.path.exists(os.path.dirname(dest)):
      os.makedirs(os.path.dirname(dest))

    with open(dest, 'w') as fl:
      fl.write(html_results)

    print pretty_results

def main():
  """Main entry point for CTS script from command line"""
  ec_dir = os.path.realpath(os.path.join(
      os.path.dirname(os.path.abspath(__file__)), '..'))
  os.chdir(ec_dir)

  dut = DEFAULT_DUT
  module = 'meta'

  parser = argparse.ArgumentParser(description='Used to build/flash boards')
  parser.add_argument('-d',
                      '--dut',
                      help='Specify DUT you want to build/flash')
  parser.add_argument('-m',
                      '--module',
                      help='Specify module you want to build/flash')
  parser.add_argument('--debug',
                      action='store_true',
                      help=('If building, build with debug printing enabled. '
                            'This may change test results'))
  parser.add_argument('-s',
                      '--setup',
                      action='store_true',
                      help='Connect only the TH to save its serial')
  parser.add_argument('-b',
                      '--build',
                      action='store_true',
                      help='Build test suite (no flashing)')
  parser.add_argument('-f',
                      '--flash',
                      action='store_true',
                      help='Flash boards with most recent images')
  parser.add_argument('-r',
                      '--run',
                      action='store_true',
                      help='Run tests without flashing')

  args = parser.parse_args()

  if args.module:
    module = args.module

  if args.dut:
    dut = args.dut

  cts = Cts(ec_dir, dut=dut, module=module, debug=args.debug)

  if args.setup:
    cts.setup()
  elif args.build:
    cts.build()
  elif args.flash:
    cts.flash_boards()
  elif args.run:
    cts.run()
  else:
    cts.build()
    cts.flash_boards()
    cts.run()

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


<commit_before>from email import message_from_string
import email.utils as email_utils
from email.Parser import Parser

# import imaplib as imaplib_original
import auth
from email.header import decode_header
from imapclient import IMAPClient
import time
import datetime
import logging as log

from webify import plaintext2html, fix_links, gravatar_url

base_gmail_url = 'https://mail.google.com/mail/b/' + auth.ACCOUNT + '/imap/'
HOST = 'imap.gmail.com'
ssl = True

server = None

        # {
        # tos: [ <Contacts>, ... ]
        # from:
        #     [ <Contact> name, address, ... ]
        # headers:
        #     { 'someheader', value}

        # body-text { 'content-type', bodyvalue}
        #     >>> somehow have resource keys in body-text

        # resources {'file-key', value}

        # }

     # BODY.PEEK[HEADER]
         # X-GM-THRID
         # X-GM-MSGID
         # X-GM-LABELS
        # BODY.PEEK[HEADER.FIELDS (Message-Id)]
        # BODY.PEEK[HEADER.FIELDS (From)]
        # ENVELOPE
        # RFC822.SIZE
        # UID
        # FLAGS
        # INTERNALDATE
        # X-GM-THRID
        # X-GM-MSGID
        # X-GM-LABELS

class Message():
    def __init__(self):
        self.to_contacts = []
        self.from_contacts = None
        self.subject = None
        self.date = None
        self.body_text = {}

        self.thread_id = None
        self.size = None
        self.uid = None

    def gravatar(self):
        return gravatar_url(self.from_contacts[0]['address'])

    def gmail_url(self):
        if not self.uid: return
        return "https://mail.google.com/mail/u/0/#inbox/" + hex(self.uid)

# use decorators to make sure this happens?
def connect():
    global server
    log.info('Connecting to %s ...' % auth.HOST,)

    try:
        server.noop()
        log.info('Already connected to host.')
        return True
    except Exception, e:
        log.info('No active connection. Opening connection...')

    try:
        server = IMAPClient(HOST, use_uid=True, ssl=ssl)
        server.oauth_login(base_gmail_url,
                    auth.OAUTH_TOKEN,
                    auth.OAUTH_TOKEN_SECRET,
                    auth.CONSUMER_KEY,
                    auth.CONSUMER_SECRET)
    except IMAPClient.Error, err:
        log.error("Could not connect. %s", e)
        return False

    log.info('Connection successful.')
    return True

def setup():
    connect()
    select_folder("Inbox")

def list_folders():
    global server

    try:
        resp = server.xlist_folders()
    except Exception, e:
        raise e
    return [dict(flags = f[0], delimiter = f[1], name = f[2]) for f in resp]

def message_count(folder):
    global server
    select_info = server.select_folder(folder, readonly=True)
    return int(select_info['EXISTS'])

def select_folder(folder):
    global server

    # TOFIX catch exception here
    select_info = server.select_folder(folder, readonly=True)
    # Format of select_info
    # {'EXISTS': 3597, 'PERMANENTFLAGS': (),
    # 'UIDNEXT': 3719,
    # 'FLAGS': ('\\Answered', '\\Flagged', '\\Draft', '\\Deleted', '\\Seen', '$Pending', 'Junk', 'NonJunk', 'NotJunk', '$Junk', 'Forwarded', '$Forwarded', 'JunkRecorded', '$NotJunk'),
    # 'UIDVALIDITY': 196, 'READ-ONLY': [''], 'RECENT': 0}
    log.info('Selected folder %s with %d messages.' % (folder, select_info['EXISTS']) )
    return select_info

# TODO this shit is broken
def create_draft(message_string):
    global server

    print 'Adding test draft...'
    server.append("[Gmail]/Drafts",
                str(email.message_from_string(message_string)))
    print 'Done!'

def latest_message_uid():
    global server
    messages = server.search(['NOT DELETED'])
    return messages[-1]

def fetch_latest_message():
    global server
    latest_email_uid = latest_message_uid();
    response = server.fetch(latest_email_uid, ['RFC822', 'X-GM-THRID', 'X-GM-MSGID'])
    return response[latest_email_uid]['RFC822']

def fetch_msg(msg_uid):
    global server
    response = server.fetch(msg_uid, ['RFC822', 'X-GM-THRID', 'X-GM-MSGID'])
    raw_response = response[msg_uid]['RFC822']

    new_msg = Message()

    parser = Parser()
    msg = parser.parsestr(raw_response)

    def make_uni(txt, default_encoding="ascii"):
        return u"".join([unicode(text, charset or default_encoding)
                   for text, charset in decode_header(txt)])

    def parse_contact(headers):
        # Works with both strings and lists
        try: headers += []
        except: headers = [headers]

        # combine and flatten header values
        addrs = reduce(lambda x,y: x+y, [msg.get_all(a, []) for a in headers])

        if len(addrs) > 0:
            return [ dict(name = make_uni(t[0]),
                          address=make_uni(t[1]))
                    for t in email_utils.getaddresses(addrs)]
        else:
            return [ dict(name = "undisclosed recipients", address = "") ]

    new_msg.to_contacts = parse_contact(['to', 'cc'])
    new_msg.from_contacts = parse_contact(['from'])

    subject = make_uni(msg['subject'])
    # Headers will wrap when longer than 78 lines per RFC822_2
    subject = subject.replace('\n\t', '')
    subject = subject.replace('\r\n', '')

    # Remove "RE" or whatever
    if subject[:4] == u'RE: ' or subject[:4] == u'Re: ' :
        subject = subject[4:]

    new_msg.subject = subject

    # TODO : upgrade to python3?
    # new_msg.date = email_utils.parsedate_to_datetime(msg["Date"])

    # date_tuple = email_utils.parsedate_tz(message["Date"])
    # sent_time = time.strftime("%b %m, %Y — %I:%M %p", date_tuple[:9])

    time_epoch = time.mktime( email_utils.parsedate_tz(msg["Date"])[:9] )
    new_msg.date = datetime.datetime.fromtimestamp(time_epoch)

    log.info('To: %s' % new_msg.to_contacts)
    log.info('From: %s' % new_msg.from_contacts)
    log.info('Subject: %s' % new_msg.subject)
    log.info('Date: %s' % new_msg.date.strftime('%b %m, %Y %I:%M %p') )

    # log.info('Date: %s' % new_msg.date )

    msg_text = None
    content_type = None

    maintype = msg.get_content_maintype()
    if maintype == 'multipart':
        for part in msg.get_payload():
            if part.get_content_maintype() == 'text':
                msg_text = part.get_payload(decode=True)
                content_type = part.get_content_type()
    elif maintype == 'text':
        msg_text = msg.get_payload(decode=True)
        content_type = msg.get_content_type()
    else:
        log.error("Message doesn't have text Content-Type: %s" % msg)

        # msg_text = quopri.decodestring(msg_text)

    # My old way of doing this

    # for part in message.walk():
    #     content_type = part.get_content_type()
    #     if content_type == "text/html":
    #         msg_text = part.get_payload(decode=True)
    #         break
    #     if part.get_content_type() == 'text/plain':
    #         continue
    #         # break

    if msg_text == None:
        log.error("Couldn't find message text! %s" % msg)
        return ""

    msg_text = msg_text.decode('iso-8859-1').encode('utf8')

    # TODO: This doesn't always work right.
    # This is so broken
    # msg_text = trim_quoted_text(msg_text, content_type)

    if content_type == 'text/plain':
        msg_text = plaintext2html(msg_text)

    msg_text = fix_links(msg_text)

    new_msg.body_text = msg_text

    return new_msg

def fetch_all_udids(self):
    global server
    UIDs = server.search(['NOT DELETED'])
    return UIDs

def fetch_thread(self, thread_id):
    global server
    threads = server.search('X-GM-THRID %s' % str(thread_id) )
    print threads

def fetch_headers(self, folder_name):
    global server

    select_info = self.select_folder(folder_name)
    UIDs = self.fetch_all_udids()

    query = 'BODY.PEEK[HEADER.FIELDS (TO CC FROM DATE SUBJECT)]'
    query_key = 'BODY[HEADER.FIELDS (TO CC FROM DATE SUBJECT)]'

    print 'Fetching all message headers:', query
    messages = server.fetch(UIDs, [query, 'X-GM-THRID'])
    print "   found %i messages." % len( messages.values() )

    parser = Parser()

    subjects = []

    for message_dict in messages.values():
        raw_header = message_dict[query_key]
        thread_id = message_dict['X-GM-THRID']

        message = parser.parsestr(raw_header)

        tos = message.get_all('to', [])
        ccs = message.get_all('cc', [])
        from_address = message.get_all('from', [])
        # from_name_decoded = decode_header(from_name)
        header_text = decode_header( message['Subject'] )
        default_encoding="ascii"
        header_sections = [unicode(text, charset or default_encoding)
                           for text, charset in header_text]

        # Headers will wrap when longer than 78 lines per RFC822_2
        subject = u"".join(header_sections).replace('\n\t', '').replace('\r\n', '')

        date_tuple = email_utils.parsedate_tz(message["Date"])
        sent_time = time.strftime("%I:%M %p - %b %d, %Y", date_tuple[:9])

        print 'From:', from_address[0]
        print 'To:', tos
        if len(ccs) > 0:
            print 'CC:', ccs
        print 'Date:', time.strftime("%c", date_tuple[:9])
        print 'THRID', thread_id
        print 'Subject', subject
        print

        subjects.append(subject)
    return subjects

def main():

    log.basicConfig(level=log.DEBUG)

    if not ( connect() ):
        print "Couldn't connect. :("
        return

    setup()
    uid = latest_message_uid()
    msg = fetch_msg(uid)

    print msg.subject

    # select_info = m.select_folder(u'Awesome')

    # UIDs = m.fetch_all_udids()
    # latest_email_uid = UIDs[-1]
    # print '   Latest UID:', latest_email_uid
    # print '   Total UIDs: ', len(UIDs)

    # thread_id = message_dict['X-GM-THRID']

    # m.create_draft("Test hello world")

    # folders =  list_folders()
    # other_folders = []
    # print '\nSpecial mailboxes:'
    # for f in folders:
    #     if u'\\AllMail' in f['flags']:
    #         print "    ALL MAIL --> ", f['name']
    #     elif u'\\Drafts' in f['flags']:
    #         print "    DRAFTS --> ", f['name']
    #     elif u'\\Important' in f['flags']:
    #         print "    IMPORTANT --> ", f['name']
    #     elif u'\\Sent' in f['flags']:
    #         print "    SENT --> ", f['name']
    #     elif u'\\Starred' in f['flags']:
    #         print "    STARRED --> ", f['name']
    #     elif u'\\Trash' in f['flags']:
    #         print "    TRASH --> ", f['name']
    #     else:
    #         other_folders.append(f)
    # print '\Other mailboxes:'
    # for f in other_folders:
    #     print "   ", f['name']

    # print 'Unread counts:'
    # for f in folders:
    #     if u'\\Noselect' in f['flags']: continue
    #     print f['flags']
    #     print "    " + f['name'] + '...',
    #     print str(message_count(f['name']))

if __name__ == "__main__":
    main()
<commit_msg>Separate some header parsing functions.<commit_after>from email import message_from_string
import email.utils as email_utils
from email.Parser import Parser

# import imaplib as imaplib_original
import auth
from email.header import decode_header
from imapclient import IMAPClient
import time
import datetime
import logging as log

from webify import plaintext2html, fix_links, gravatar_url

base_gmail_url = 'https://mail.google.com/mail/b/' + auth.ACCOUNT + '/imap/'
HOST = 'imap.gmail.com'
ssl = True

server = None

# {
# tos: [ <Contacts>, ... ]
# from:
#     [ <Contact> name, address, ... ]
# headers:
#     { 'someheader', value}

# body-text { 'content-type', bodyvalue}
#     >>> somehow have resource keys in body-text

# resources {'file-key', value}

# }

# BODY.PEEK[HEADER]
# X-GM-THRID
# X-GM-MSGID
# X-GM-LABELS
# BODY.PEEK[HEADER.FIELDS (Message-Id)]
# BODY.PEEK[HEADER.FIELDS (From)]
# ENVELOPE
# RFC822.SIZE
# UID
# FLAGS
# INTERNALDATE
# X-GM-THRID
# X-GM-MSGID
# X-GM-LABELS

class Message():
    def __init__(self):
        self.to_contacts = []
        self.from_contacts = None
        self.subject = None
        self.date = None
        self.body_text = {}

        self.thread_id = None
        self.size = None
        self.uid = None

    def gravatar(self):
        return gravatar_url(self.from_contacts[0]['address'])

    def gmail_url(self):
        if not self.uid: return
        return "https://mail.google.com/mail/u/0/#inbox/" + hex(self.uid)

# use decorators to make sure this happens?
def connect():
    global server
    log.info('Connecting to %s ...' % auth.HOST,)

    try:
        server.noop()
        log.info('Already connected to host.')
        return True
    except Exception, e:
        log.info('No active connection. Opening connection...')

    try:
        server = IMAPClient(HOST, use_uid=True, ssl=ssl)
        server.oauth_login(base_gmail_url,
                    auth.OAUTH_TOKEN,
                    auth.OAUTH_TOKEN_SECRET,
                    auth.CONSUMER_KEY,
                    auth.CONSUMER_SECRET)
    except IMAPClient.Error, err:
        log.error("Could not connect. %s", e)
        return False

    log.info('Connection successful.')
    return True

def setup():
    connect()
    select_folder("Inbox")

def list_folders():
    global server

    try:
        resp = server.xlist_folders()
    except Exception, e:
        raise e
    return [dict(flags = f[0], delimiter = f[1], name = f[2]) for f in resp]

def message_count(folder):
    global server
    select_info = server.select_folder(folder, readonly=True)
    return int(select_info['EXISTS'])

def select_folder(folder):
    global server

    # TOFIX catch exception here
    select_info = server.select_folder(folder, readonly=True)
    # Format of select_info
    # {'EXISTS': 3597, 'PERMANENTFLAGS': (),
    # 'UIDNEXT': 3719,
    # 'FLAGS': ('\\Answered', '\\Flagged', '\\Draft', '\\Deleted', '\\Seen', '$Pending', 'Junk', 'NonJunk', 'NotJunk', '$Junk', 'Forwarded', '$Forwarded', 'JunkRecorded', '$NotJunk'),
    # 'UIDVALIDITY': 196, 'READ-ONLY': [''], 'RECENT': 0}
    log.info('Selected folder %s with %d messages.' % (folder, select_info['EXISTS']) )
    return select_info

# TODO this shit is broken
def create_draft(message_string):
    global server

    print 'Adding test draft...'
    server.append("[Gmail]/Drafts",
                str(email.message_from_string(message_string)))
    print 'Done!'

def latest_message_uid():
    global server
    messages = server.search(['NOT DELETED'])
    return messages[-1]

def fetch_latest_message():
    global server
    latest_email_uid = latest_message_uid();
    response = server.fetch(latest_email_uid, ['RFC822', 'X-GM-THRID', 'X-GM-MSGID'])
    return response[latest_email_uid]['RFC822']

def parse_main_headers(msg, new_msg = Message()):
    # new_msg = Message()

    def make_uni(txt, default_encoding="ascii"):
        return u"".join([unicode(text, charset or default_encoding)
                   for text, charset in decode_header(txt)])

    def parse_contact(headers):
        # Works with both strings and lists
        try: headers += []
        except: headers = [headers]

        # combine and flatten header values
        addrs = reduce(lambda x,y: x+y, [msg.get_all(a, []) for a in headers])

        if len(addrs) > 0:
            return [ dict(name = make_uni(t[0]),
                          address=make_uni(t[1]))
                    for t in email_utils.getaddresses(addrs)]
        else:
            return [ dict(name = "undisclosed recipients", address = "") ]

    new_msg.to_contacts = parse_contact(['to', 'cc'])
    new_msg.from_contacts = parse_contact(['from'])

    subject = make_uni(msg['subject'])
    # Headers will wrap when longer than 78 lines per RFC822_2
    subject = subject.replace('\n\t', '')
    subject = subject.replace('\r\n', '')

    # Remove "RE" or whatever
    if subject[:4] == u'RE: ' or subject[:4] == u'Re: ' :
        subject = subject[4:]

    new_msg.subject = subject

    # TODO : upgrade to python3?
    # new_msg.date = email_utils.parsedate_to_datetime(msg["Date"])

    # date_tuple = email_utils.parsedate_tz(message["Date"])
    # sent_time = time.strftime("%b %m, %Y — %I:%M %p", date_tuple[:9])

    time_epoch = time.mktime( email_utils.parsedate_tz(msg["Date"])[:9] )
    new_msg.date = datetime.datetime.fromtimestamp(time_epoch)

    return new_msg

def parse_body(msg, new_msg = Message()):

    msg_text = None
    content_type = None

    maintype = msg.get_content_maintype()
    if maintype == 'multipart':
        for part in msg.get_payload():
            if part.get_content_maintype() == 'text':
                msg_text = part.get_payload(decode=True)
                content_type = part.get_content_type()
    elif maintype == 'text':
        msg_text = msg.get_payload(decode=True)
        content_type = msg.get_content_type()
    else:
        log.error("Message doesn't have text Content-Type: %s" % msg)

    if msg_text == None:
        log.error("Couldn't find message text! %s" % msg)
        return ""

    msg_text = msg_text.decode('iso-8859-1').encode('utf8')

    # TODO: This is so broken
    # msg_text = trim_quoted_text(msg_text, content_type)

    if content_type == 'text/plain':
        msg_text = plaintext2html(msg_text)

    msg_text = fix_links(msg_text)

    new_msg.body_text = msg_text

    return new_msg

def fetch_msg(msg_uid):
    global server

    log.info("Fetching message. UID: %i" % msg_uid)
    response = server.fetch(msg_uid, ['RFC822', 'X-GM-THRID', 'X-GM-MSGID'])
    raw_response = response[msg_uid]['RFC822']
    log.info("Received response. Size: %i" % len(raw_response))

    msg = Parser().parsestr(raw_response)

    # headers
    new_msg = parse_main_headers(msg)  # returns Message()

    # body
    new_msg = parse_body(msg, new_msg)

    log.info('To: %s' % new_msg.to_contacts)
    log.info('From: %s' % new_msg.from_contacts)
    log.info('Subject: %s' % new_msg.subject)
    log.info('Date: %s' % new_msg.date.strftime('%b %m, %Y %I:%M %p') )

    return new_msg

def fetch_all_udids():
    global server
    UIDs = server.search(['NOT DELETED'])
    return UIDs

def fetch_thread(self, thread_id):
    global server
    threads = server.search('X-GM-THRID %s' % str(thread_id) )
    print threads

def fetch_headers(folder_name):
    global server

    query = 'BODY.PEEK[HEADER.FIELDS (TO CC FROM DATE SUBJECT)]'
    query_key = 'BODY[HEADER.FIELDS (TO CC FROM DATE SUBJECT)]'

    log.info("Fetching message headers. Query: %s" % query)
    select_info = select_folder(folder_name)
    UIDs = fetch_all_udids()

    log.info("Fetching message headers. Query: %s" % query)
    messages = server.fetch(UIDs, [query, 'X-GM-THRID'])
    log.info("found %i messages." % len(messages.values()))

    parser = Parser()
    subjects = []

    for message_dict in messages.values():
        raw_header = message_dict[query_key]
        thread_id = message_dict['X-GM-THRID']

        msg = parser.parsestr(raw_header)

        # headers
        new_msg = parse_main_headers(msg)  # returns Message()
        subjects.append(new_msg.subject)

    return subjects

def main():

    log.basicConfig(level=log.DEBUG)

    if not ( connect() ):
        print "Couldn't connect. :("
        return

    setup()

    uid = latest_message_uid()
    msg = fetch_msg(uid)

    print msg.subject

    # select_info = m.select_folder(u'Awesome')

    # UIDs = m.fetch_all_udids()
    # latest_email_uid = UIDs[-1]
    # print '   Latest UID:', latest_email_uid
    # print '   Total UIDs: ', len(UIDs)

    # thread_id = message_dict['X-GM-THRID']

    # m.create_draft("Test hello world")

    # folders =  list_folders()
    # other_folders = []
    # print '\nSpecial mailboxes:'
    # for f in folders:
    #     if u'\\AllMail' in f['flags']:
    #         print "    ALL MAIL --> ", f['name']
    #     elif u'\\Drafts' in f['flags']:
    #         print "    DRAFTS --> ", f['name']
    #     elif u'\\Important' in f['flags']:
    #         print "    IMPORTANT --> ", f['name']
    #     elif u'\\Sent' in f['flags']:
    #         print "    SENT --> ", f['name']
    #     elif u'\\Starred' in f['flags']:
    #         print "    STARRED --> ", f['name']
    #     elif u'\\Trash' in f['flags']:
    #         print "    TRASH --> ", f['name']
    #     else:
    #         other_folders.append(f)
    # print '\Other mailboxes:'
    # for f in other_folders:
    #     print "   ", f['name']

    # print 'Unread counts:'
    # for f in folders:
    #     if u'\\Noselect' in f['flags']: continue
    #     print f['flags']
    #     print "    " + f['name'] + '...',
    #     print str(message_count(f['name']))

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