Mercurial > public > scc_access / file revision

#!/usr/bin/env python
"""
The MIT License (MIT)

Copyright (c) 2015, Ioannis Binietoglou

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

__version__ = "0.5.0"


# Try to read the settings from the settings.py file
try:
    from settings import *
except:
    raise ImportError(
            """A settings file (setting.py) is required to run the script. 
             You can use settings.sample.py for a template.""")


import requests
import urlparse
import argparse
import os
import re
import time
import StringIO
from zipfile import ZipFile
import datetime


# Construct the absolute URLs
LOGIN_URL = urlparse.urljoin(BASE_URL, 'accounts/login/')
UPLOAD_URL = urlparse.urljoin(BASE_URL, 'data_processing/measurements/quick/')
DOWNLOAD_PREPROCESSED = urlparse.urljoin(BASE_URL, 'data_processing/measurements/{0}/download-preprocessed/')
DOWNLOAD_OPTICAL = urlparse.urljoin(BASE_URL, 'data_processing/measurements/{0}/download-optical/')
DOWNLOAD_GRAPH = urlparse.urljoin(BASE_URL, 'data_processing/measurements/{0}/download-plots/')
DELETE_MEASUREMENT =  urlparse.urljoin(BASE_URL, 'admin/database/measurements/{0}/delete/')
API_BASE_URL = urlparse.urljoin(BASE_URL, 'api/v1/')

# The regex to find the measurement id from the measurement page
# This should be read from the uploaded file, but would require an extra module
regex = "<h3>Measurement (?P<measurement_id>.{12}) <small>"


class SCC:
    """ A simple class that will attempt to upload a file on the SCC server.
    The uploading is done by simulation a normal browser session. In the current
    version no check is performed, and no feedback is given if the upload 
    was successful. If everything is setup correctly, it will work. 
    """
    def __init__(self, auth = BASIC_LOGIN, output_dir = OUTPUT_DIR):
        self.auth = auth
        self.output_dir = OUTPUT_DIR
        self.session = requests.Session()
        
    def login(self, credential = DJANGO_LOGIN):
        """ Login the the website. """
        self.login_credentials = {'username': credential[0], 
                                  'password': credential[1]}
        
        # Get upload form
        login_page = self.session.get(LOGIN_URL, 
                                        auth = self.auth, verify = False)
        
        # Submit the login data
        login_submit = self.session.post(LOGIN_URL, 
                                          data = self.login_credentials, 
                                          headers = {'X-CSRFToken': login_page.cookies['csrftoken'], 
                                                     'referer': LOGIN_URL}, 
                                          verify = False, 
                                          auth = self.auth)
        return login_submit

    def logout(self):
        pass
    
    def upload_file(self, filename, system_id):
        """ Upload a filename for processing with a specific system. If the 
        upload is successful, it returns the measurement id. """
        # Get submit page
        upload_page = self.session.get(UPLOAD_URL, 
                                       auth = self.auth, 
                                       verify = False)
       
       # Submit the data
        upload_data = {'system': system_id}
        files  = {'data': open(filename, 'rb')}
        
        print "Uploading of file %s started." % filename
        
        upload_submit = self.session.post(UPLOAD_URL, 
                                          data = upload_data, 
                                          files = files, 
                                          headers = {'X-CSRFToken': upload_page.cookies['csrftoken'], 
                                                     'referer': UPLOAD_URL}, 
                                          verify = False, 
                                          auth = self.auth)
        
        if upload_submit.status_code != 200:
            print "Connection error. Status code: %s" % upload_submit.status_code
            return False
        
        measurement_id = True    
        
        # Check if there was a redirect to a new page.
        if upload_submit.url == UPLOAD_URL:
            measurement_id = False
            print "Uploaded file rejected! Try to upload manually to see the error."
        else:
            measurement_id = re.findall(regex, upload_submit.text)[0]
            print "Successfully uploaded measurement with id %s." % measurement_id

        return measurement_id
    
    def download_files(self, measurement_id, subdir, download_url):
        """ Downloads some files from the download_url to the specified
        subdir. This method is used to download preprocessed file, optical 
        files etc.
        """        
        # Get the file
        request = self.session.get(download_url, auth = self.auth, 
                                                 verify = False,
                                                 stream=True)
        
        # Create the dir if it does not exist
        local_dir = os.path.join(self.output_dir, measurement_id, subdir)
        if not os.path.exists(local_dir):
            os.makedirs(local_dir)
        
        
        # Save the file by chunk, needed if the file is big.
        memory_file = StringIO.StringIO()
        
        for chunk in request.iter_content(chunk_size=1024): 
            if chunk: # filter out keep-alive new chunks
                memory_file.write(chunk)
                memory_file.flush()
        
        zip_file = ZipFile(memory_file)
        
        for ziped_name in zip_file.namelist():
            basename = os.path.basename(ziped_name)
           
            local_file = os.path.join(local_dir, basename)
        
            with open(local_file, 'wb') as f:
                f.write(zip_file.read(ziped_name))

    def download_preprocessed(self, measurement_id):
        """ Download preprocessed files for the measurement id. """
        # Construct the download url
        download_url = DOWNLOAD_PREPROCESSED.format(measurement_id)
        self.download_files(measurement_id, 'scc_preprocessed', download_url)
    
    def download_optical(self, measurement_id):
        """ Download optical files for the measurement id. """
        # Construct the download url
        download_url = DOWNLOAD_OPTICAL.format(measurement_id)
        self.download_files(measurement_id, 'scc_optical', download_url)
    
    def download_graphs(self, measurement_id):
        """ Download profile graphs for the measurement id. """
        # Construct the download url
        download_url = DOWNLOAD_GRAPH.format(measurement_id)
        self.download_files(measurement_id, 'scc_plots', download_url)
    
    def process(self, filename, system_id):
        """ Upload a file for processing and wait for the processing to finish.
        If the processing is successful, it will download all produced files.
        """
        print "--- Processing started on %s. ---" % datetime.datetime.now()
        # Upload file
        measurement_id = self.upload_file(filename, system_id) 
        
        measurement = None
        if measurement_id:
            measurement = self.get_measurement(measurement_id)
            while measurement.is_running:
                print "Measurement is being processed (status: %s, %s, %s). Please wait." % (measurement.upload, 
                                                                                             measurement.pre_processing, 
                                                                                             measurement.opt_retrievals)
                time.sleep(10)
                measurement = self.get_measurement(measurement_id)
            print "Measurement processing finished (status: %s, %s, %s)." % (measurement.upload, 
                                                                             measurement.pre_processing, 
                                                                             measurement.opt_retrievals)
            if measurement.pre_processing == 127:
                print "Downloading preprocessed files."
                self.download_preprocessed(measurement_id)
            if measurement.opt_retrievals == 127:
                print "Downloading optical files."
                self.download_optical(measurement_id)
                print "Downloading graphs."
                self.download_graphs(measurement_id)
        print "--- Processing finished. ---"
        return measurement
        
    def get_status(self, measurement_id):
        """ Get the processing status for a measurement id through the API. """
        measurement_url = urlparse.urljoin(API_BASE_URL, 'measurements/?id__exact=%s' % measurement_id)
        
        response = self.session.get(measurement_url, 
                                    auth = self.auth, 
                                    verify = False)
        
        response_dict = response.json()
        
        if response_dict['objects']:
            measurement_list = response_dict['objects']
            measurement = Measurement(measurement_list[0])
            return (measurement.upload, measurement.pre_processing, measurement.opt_retrievals)
        else:
            print "No measurement with id %s found on the SCC." % measurement_id
            return None
    
    def get_measurement(self, measurement_id):
        measurement_url = urlparse.urljoin(API_BASE_URL, 'measurements/%s/' % measurement_id)
        
        response = self.session.get(measurement_url, 
                                    auth = self.auth, 
                                    verify = False)
        
        response_dict = response.json()
        
        if response_dict:
            measurement = Measurement(response_dict)
            return measurement
        else:
            print "No measurement with id %s found on the SCC." % measurement_id
            return None
    
    def delete_measurement(self, measurement_id):
        """ Deletes a measurement with the provided measurement id. The user
        should have the appropriate permissions. 
        
        The procedures is performed directly through the web interface and
        NOT through the API.
        """
        # Get the measurement object
        measurement = self.get_measurement(measurement_id)
        
        # Check that it exists
        if measurement is None:
            print "Nothing to delete."
            return None
        
        # Go the the page confirming the deletion
        delete_url = DELETE_MEASUREMENT.format(measurement.id)
        
        confirm_page = self.session.get(delete_url, 
                                    auth = self.auth, 
                                    verify = False)
         
        # Check that the page opened properly
        if confirm_page.status_code != 200:
            print "Could not open delete page. Status: {0}".format(confirm_page.status_code)
            return None
        
        # Delete the measurement
        delete_page = self.session.post(delete_url, 
                                       auth=self.auth, 
                                       verify=False,
                                       data={'post':'yes'}, 
                                       headers={'X-CSRFToken': confirm_page.cookies['csrftoken'],
                                                 'referer': delete_url}
                                                 )
        if delete_page.status_code != 200:
            print "Something went wrong. Delete page status: {0}".format(
                                                        delete_page.status_code)
            return None
        
        print "Deleted measurement {0}".format(measurement_id)
        return True    

    def available_measurements(self):
        """ Get a list of available measurement on the SCC. """
        measurement_url = urlparse.urljoin(API_BASE_URL, 'measurements')
        response = self.session.get(measurement_url, 
                                    auth = self.auth, 
                                    verify = False)
        response_dict = response.json()
        
        measurements = None
        if response_dict:
            measurement_list = response_dict['objects']
            measurements = [Measurement(measurement_dict) for measurement_dict in measurement_list]
            print "Found %s measurements on the SCC." % len(measurements)
        else:
            print "No response received from the SCC when asked for available measurements."

        return measurements
    
    def measurement_id_for_date(self, t1, call_sign = 'bu', base_number = 0):
        """ Give the first available measurement id on the SCC for the specific
        date. 
        """
        date_str = t1.strftime('%Y%m%d')
        search_url = urlparse.urljoin(API_BASE_URL, 'measurements/?id__startswith=%s' % date_str)
        
        response = self.session.get(search_url, 
                                    auth = self.auth, 
                                    verify = False)
        
        response_dict = response.json()
        
        measurement_id = None
        
        if response_dict:
            measurement_list = response_dict['objects']
            existing_ids = [measurement_dict['id'] for measurement_dict in measurement_list]
            
            measurement_number = base_number
            measurement_id = "%s%s%02i" % (date_str, call_sign, measurement_number)
           
            while measurement_id in existing_ids:
                measurement_number = measurement_number + 1
                measurement_id = "%s%s%02i" % (date_str, call_sign, measurement_number)
                if measurement_number == 100:
                    raise ValueError('No available measurement id found.')

        return measurement_id     
        
        
class ApiObject:
    """ A generic class object. """
    
    def __init__(self, dict_response):
        
        if dict_response:
            # Add the dictionary key value pairs as object properties
            for key, value in dict_response.items():
                setattr(self, key, value)
            self.exists = True
        else:
            self.exists = False


class Measurement(ApiObject):
    """ This class represents the measurement object as returned in the SCC API.
    """
    @property
    def is_running(self):
        """ Returns True if the processing has not finished.
        """
        if self.upload == 0:
            return False
        if self.pre_processing == -127:
            return False
        if self.pre_processing == 127:
            if self.opt_retrievals in [127, -127]:
                return False
        return True
    
    def delete(self):
        """ Delete the entry from the SCC database. """
        
        
    def __str__(self):
        return "%s: %s, %s, %s" % (self.id, 
                                   self.upload, 
                                   self.pre_processing, 
                                   self.opt_retrievals)
                                   
                                   
def upload_file(filename, system_id, auth = BASIC_LOGIN, credential = DJANGO_LOGIN):
    """ Shortcut function to upload a file to the SCC. """
    scc = SCC(auth)
    scc.login(credential)
    measurement_id = scc.upload_file(filename, system_id)
    scc.logout()
    return measurement_id

def process_file(filename, system_id, auth = BASIC_LOGIN, credential = DJANGO_LOGIN):
    """ Shortcut function to process a file to the SCC. """
    scc = SCC(auth)
    scc.login(credential)
    measurement = scc.process(filename, system_id)
    scc.logout()
    return measurement   

def delete_measurement(measurement_id, auth = BASIC_LOGIN, credential = DJANGO_LOGIN):
    """ Shortcut function to delete a measurement from the SCC. """
    scc = SCC(auth)
    scc.login(credential)
    scc.delete_measurement(measurement_id)
    scc.logout()

# When running through terminal
if __name__ == '__main__':
    
    # Define the command line arguments.
    parser = argparse.ArgumentParser()
    parser.add_argument("filename",  nargs='?', help = "Measurement file name or path.", default='')
    parser.add_argument("system", nargs='?', help = "Processing system id.", default=0)
    parser.add_argument("-p", "--process", help="Wait for the results of the processing.",
                    action="store_true")
    parser.add_argument("-d", "--delete", help="Measurement ID to delete.")
    args = parser.parse_args()
    
    # If the arguments are OK, try to login on the site and upload.
    if args.delete:
        # If the delete is provided, do nothing else
        delete_measurement(args.delete)
    else:
        if (args.filename == '') or  (args.system == 0):
            parser.error('Provide a valid filename and system parameters.\nRun with -h for help.\n')
            
        if args.process:
            process_file(args.filename, args.system)
        else:
            upload_file(args.filename, args.system)