#! /usr/bin/env python
'''
This file is a copy of "wpd" from https://github.com/borismus/webplatform-tools,
modified and upgraded.
'''

import os
from optparse import OptionParser
import urllib, urllib2
import json
import sys
import re
import time
import logging

# We require json format
DEFAULT_PARAMS = {'format': 'json'}

# Usage for this program
usage = """

    %prog <options> command <arguments>

Commands:
    ls <prefix> (<limit>) (<pagefrom>)
        Lists file beginning with given prefix.
        Limit is the option max number of page names returned.
        Pagefrom is the page to continue listing, useful if the
        limit is too small and you want to continue listing where
        the previous ls left off.

    edit <page> <filename> <editsummary>
        Replaces the content of the specified page, using the file contents,
        and using the edit summary as the explanation for the edit.

    batchupload <mappingfile> <editsummary>
        Upload multiple pages from multiple files, for the summary edit reason
        given.  The mappingfile must have repeated lines in the format:
            FILENAME || PAGENAME
        Any line not in this format is ignored.  Leading or trailing BAR
        characters or spaces are ignored.

    delete <page> <reason>
        Delete a single page for the given reason.

    batchdelete <pagenamesfile> <deletereason>
        Delete multiple pages, for the given reason.
        The pagenamesfile must have repeated lines in the format:
            PAGENAME
        One page name per line.

    cat <page>
        Extract content of given page

Environment variables:
    WPD_USERNAME - username to login to wiki, overridden by --username option
    WPD_PASSWORD - password to login to wiki, overridden by --password option
    WPD_API_ROOT - root url of wiki api (do not include /api.php at end),
                   overridden by --api-root option
"""

class WPDConnection:
  '''
  Embodies the connection to the wiki, allowing operations to be performed on it.
  '''
  def __init__(self, username, password, apiroot):
    self.username = username
    self.password = password
    self.apiroot = apiroot
    self.changetoken = None

    cookie_handler = urllib2.HTTPCookieProcessor()
    redirect_handler = urllib2.HTTPRedirectHandler()
    self.opener = urllib2.build_opener(redirect_handler, cookie_handler)

  def logout(self):
    '''Logout from wiki, assumption is that this always succeeds.'''
    params = {
      'action': 'logout',
    }
    params.update(DEFAULT_PARAMS)
    logging.debug ("LOGOUT REQUEST:\n{}".format (params))
    resp = json.loads(self.opener.open(self.apiroot, urllib.urlencode(params)).read())
    logging.debug ("LOGOUT RESPONSE:\n{}".format (resp))

  def login(self):
    '''
    Login to wiki.

    The side effect is that a cookie gets set used on subsequent requests to
    validate logged-in user.
    '''
    params = {
      'action': 'login',
      'lgname': self.username,
      'lgpassword': self.password
    }
    params.update(DEFAULT_PARAMS)

    # Issue an initial login request
    logging.debug ("LOGIN REQUEST:\n{}".format (params))
    resp = json.loads(self.opener.open(self.apiroot, urllib.urlencode(params)).read())
    logging.debug ("LOGIN RESPONSE:\n{}".format (resp))

    # See if there was an error
    if 'error' in resp:
        if 'info' in resp['error']:
            raise Exception, 'Error logging in:\n{}'.format (resp['error']['info'])
        else:
            raise Exception, 'Unknown error logging in:\n{}'.format (resp['error'])

    # See if there was a login response
    if 'login' in resp:
        if 'result' in resp['login']:
            if resp['login']['result'] == 'NeedToken' and 'token' in resp['login']:
                # Now use the response token to do another login request
                params['lgtoken'] = resp['login']['token']
                logging.debug ("SECOND LOGIN REQUEST:\n{}".format (params))
                resp = json.loads(self.opener.open(self.apiroot, urllib.urlencode(params)).read())
                logging.debug ("SECOND LOGIN RESPONSE:\n{}".format (resp))
        else:
            raise Exception, 'Unknown login results:\n{}'.format (resp['login'])
    else:
        # Missing login section of response
        raise Exception, 'Unknown response:\n{}'.format (resp)

    # See if the last response was a success
    if 'login' in resp and 'result' in resp['login'] and resp['login']['result'] == 'Success' and 'lgtoken' in resp['login']:
        self.is_logged_in = True
    else:
        raise Exception, 'Failed to login:\n{}'.format (resp)

  def ls(self, prefix='', limit=20, pagefrom=''):
    '''
    Lists all pages under the specified prefix.
    '''
    if limit > 500:
      limit = 500  # sanity check

    params = {
        'action': 'query',
        'list': 'allpages',
        'apprefix': prefix,
        'aplimit': limit
    }
    if pagefrom != '':
      params.update ({'apfrom': pagefrom})
    params.update(DEFAULT_PARAMS)
    logging.debug ("LS REQUEST:\n{}".format (params))
    resp = json.loads(self.opener.open(self.apiroot, urllib.urlencode(params)).read())
    logging.debug ("LS RESPONSE:\n{}".format (resp))
    if 'query' in resp and 'allpages' in resp['query']:
      page_json = resp['query']['allpages']
      page_list = [p['title'] for p in page_json]
      return page_list
    else:
      raise Exception, 'Unknown list response:\n{}'.format (resp)

  def get_change_token (self, type='edit'):
    '''
    In order to edit or delete pages, obtain a token specified by type.

    Called from edit and delete, so no need to call this manually.  The
    token can be used multiple times for different changes.
    '''
    params = {
        'action': 'tokens',
        'type': type
    }
    params.update(DEFAULT_PARAMS)
    url = self.apiroot + urllib.urlencode(params)
    logging.debug ("GET {} TOKEN REQUEST:\n{}".format (type, params))
    resp = json.loads(self.opener.open(url).read())
    logging.debug ("GET {} TOKEN RESPONSE:\n{}".format (type, resp))

    # See if there was an error
    if 'error' in resp:
        if 'info' in resp['error']:
            raise Exception, 'Error getting {} token:\n{}'.format (type, resp['error']['info'])
        else:
            raise Exception, 'Unknown error getting {} token:\n{}'.format (type, resp['error'])

    # Get the token and return it
    if 'tokens' in resp and (type + 'token') in resp['tokens']:
      return resp['tokens'][type + 'token']
    else:
      raise Exception, 'Unknown response getting {} token:\n{}'.format (type, resp)

  def edit(self, page, content, edit_summary):
    '''
    Replaces the content for the particular page.
    Returns True if successful otherwise throws exception.
    '''
    if not self.is_logged_in:
      raise Exception, 'User %s is not logged in.' %self.username
    if not self.changetoken:
      self.changetoken = self.get_change_token ('edit')
    params = {
        'action': 'edit',
        'title': page,
        'summary': edit_summary,
        'text': content,
        'bot': '1',
        'watchlist': 'nochange',
        'token': self.changetoken
    }
    params.update(DEFAULT_PARAMS)
    logging.debug ("EDIT REQUEST:\n{}".format (params))
    resp = json.loads(self.opener.open(self.apiroot, urllib.urlencode(params)).read())
    logging.debug ("EDIT RESPONSE:\n{}".format (resp))

    # See if there was an error
    if 'error' in resp:
        if 'info' in resp['error']:
            raise Exception, 'Error editing page:\n{}'.format (resp['error']['info'])
        else:
            raise Exception, 'Unknown error editing page:\n{}'.format (resp['error'])

    if 'edit' in resp and 'result' in resp['edit'] and resp['edit']['result'] == 'Success':
      return True
    else:
      raise Exception ("Edit not successful:\n{}".format (resp))

  def delete (self, page, delete_reason):
    '''
    Deletes a page.
    Returns True if successful otherwise throws exception.
    '''
    if not self.is_logged_in:
      raise Exception, 'User %s is not logged in.' %self.username
    if not self.changetoken:
      self.changetoken = self.get_change_token('delete')
    params = {
        'action': 'delete',
        'title': page,
        'reason': delete_reason,
        'token': self.changetoken
    }
    params.update(DEFAULT_PARAMS)
    logging.debug ("DELETE REQUEST:\n{}".format (params))
    resp = json.loads(self.opener.open(self.apiroot, urllib.urlencode(params)).read())
    logging.debug ("DELETE RESPONSE:\n{}".format (resp))

    # See if there was an error
    if 'error' in resp:
        if 'info' in resp['error']:
            raise Exception, 'Error deleting page:\n{}'.format (resp['error']['info'])
        else:
            raise Exception, 'Unknown error deleting page:\n{}'.format (resp['error'])
    # See if there is a success
    if 'delete' in resp and 'logid' in resp['delete']:
      return True
    else:
      raise Exception ("Delete not successful:\n{}".format (resp))

  def cat (self, page):
    '''
    Extract content from a page.
    '''
    params = {
        'action': 'query',
        'prop': 'revisions',
        'rvprop': 'content',
        'titles': page
    }
    params.update(DEFAULT_PARAMS)
    logging.debug ("CAT REQUEST:\n{}".format (params))
    resp = json.loads(self.opener.open(self.apiroot, urllib.urlencode(params)).read())
    logging.debug ("CAT RESPONSE:\n{}".format (resp))
    if ('query' in resp and
        'pages' in resp['query'] and
        len (resp['query']['pages']) == 1):
      # Get the only element in the dictionary, keyed by a revision
      # number we don't know, like {u'query': {u'pages': {u'5147': {THIS}}}}
      rev = next (resp['query']['pages'].itervalues ())
      if ('revisions' in rev and len (rev['revisions']) == 1 and '*' in rev['revisions'][0]):
        return rev['revisions'][0]['*']
    raise Exception ("Unknown response:\n{}".format (resp))

if __name__ == '__main__':
  parser = OptionParser (usage=usage)
  parser.add_option (
    '-u', '--username',
    help='Login username, if missing use WPD_USERNAME env var')
  parser.add_option (
    '-p', '--password',
    help='Login password, if missing use WPD_PASSWORD env var')
  parser.add_option (
    '-a', '--apiroot',
    help='API root, if missing use WPD_API_ROOT env var')
  parser.add_option (
    '-d', '--debug',
    action="store_true", default=False,
    help='Print debugging information including all requests and responses')

  # Extract options and separate them from arguments
  (options, args) = parser.parse_args ()

  # No args means print help and exit
  if len(args) == 0:
    parser.print_help ()
    sys.exit(1)

  # Use option first, and if missing use env var
  command = args[0]
  username = options.username or os.getenv ("WPD_USERNAME")
  password = options.password or os.getenv ("WPD_PASSWORD")
  apiroot = options.apiroot or os.getenv ("WPD_API_ROOT")
  if apiroot is not None:
    apiroot += "/api.php?"

  # Set debug level
  if options.debug:
    logging.basicConfig (stream=sys.stderr, level=logging.DEBUG)
    logging.debug ('Logging level set to logging.DEBUG')
  else:
    logging.basicConfig (stream=sys.stderr, level=logging.ERROR)
    logging.debug ('Logging level set to logging.ERROR')

  # Ensure either env var set or option set
  if apiroot is None:
    print 'ERROR: missing apiroot\n'
    parser.print_help ()
    sys.exit (1)
  # Only ls and cat do NOT require username and password
  if command not in ('ls', 'cat') and (username is None or password is None):
    print 'ERROR: missing username or password\n'
    parser.print_help ()
    sys.exit (1)

  conn = WPDConnection (username, password, apiroot)

  if command == 'ls':
    prefix=''
    pagefrom=''
    limit=20
    if len(args) < 1 or len(args) > 4:
      print ('ERROR: ls arguments invalid\n')
      parser.print_help ()
      sys.exit(1)
    if len(args) > 1:
      prefix = args[1]
    if len(args) > 2:
      limit = int(args[2])
    if len(args) > 3:
      pagefrom = args[3]
    print ('\n'.join (conn.ls (prefix=prefix, limit=limit, pagefrom=pagefrom)))
  elif command == 'edit':
    if len(args) != 4:
      print ('ERROR: edit arguments invalid\n')
      parser.print_help ()
      sys.exit(1)
    page = args[1]
    content = open (args[2], 'r').read ()
    edit_summary = args[3]
    conn.login()
    conn.edit(page, content, edit_summary)
    print ("Success")
    conn.logout()
  elif command == 'batchupload':
    if len(args) != 3 and len(args) != 4:
      print ('ERROR: batchupload arguments invalid')
      parser.print_help ()
      sys.exit(1)
    begintime = int (time.time ())
    mappingfile = args[1]
    edit_summary = args[2]
    count = 0
    conn.login()
    for line in open (mappingfile, 'r'):
      pattern = re.compile ("^[| ]*(.*) \|\| (.*?)[| ]*$")
      match = re.match (pattern, line)
      if match and match.lastindex == 2:
        filename = match.group (1)
        pagename = match.group (2)
        print ("Upload {} as {}".format (filename, pagename))
        content = open (filename, 'r').read ()
        conn.edit(pagename, content, edit_summary)
        count += 1
    conn.logout()
    duration = int (time.time ()) - begintime
    print ("SUMMARY: uploaded {} pages in {} seconds".format (count, duration))
  elif command == 'delete':
    if len(args) != 3:
      print ('ERROR: delete arguments invalid')
      parser.print_help ()
      sys.exit(1)
    page = args[1]
    delete_reason = args[2]
    conn.login()
    conn.delete(page, delete_reason)
    print ("Success")
    conn.logout()
  elif command == 'batchdelete':
    if len(args) != 3:
      print ('ERROR: batchdelete arguments invalid')
      parser.print_help ()
      sys.exit(1)
    begintime = int (time.time ())
    pagenamesfile = args[1]
    delete_reason = args[2]
    count = 0
    conn.login()
    for page in open (pagenamesfile, 'r'):
      pagename = page.strip ()
      print ("Delete {}".format (pagename))
      conn.delete (pagename, delete_reason)
      count += 1
    conn.logout()
    duration = int (time.time ()) - begintime
    print ("SUMMARY: deleted {} pages in {} seconds".format (count, duration))
  elif command == 'cat':
    if len (args) != 2:
      print ('ERROR: cat arguments invalid')
      parser.print_help ()
      sys.exit (1)
    print (conn.cat (args[1]))
