<commit_before>#!/usr/bin/env python

import sys, os, glob, string
import numpy as np
import astropy as ast
import matplotlib.pyplot as plt
from pyraf import iraf
from tqdm import tqdm
import odi_config as odi
import glob
import shutil
import pandas as pd
import time

try:
    object_str, filters, instrument, images, illcor_flag, skyflat_src, wcs_flag, reproject_flag, scale_flag, stack_flag, gaia_flag, cluster_flag, ra_center, dec_center, min_radius = odi.cfgparse('config.yaml')
except IOError:
    print 'config.yaml does not exist, quitting...'
    exit()

# for basic processing, filter shouldn't matter enough to consider separately
# or rather, this script already handles that just fine
# so just stick all the image names together into one long list
images_ = []
for filt in images:
    for key in images[filt]:
        images_.append(images[filt][key])

rad, decd = odi.get_targ_ra_dec(images_[0], 'OTA33.SCI')
source = 'sdss'
inst = odi.instrument(images_[0])
#source = 'twomass'

#Create offline catalogs
for img in images_:
    print 'Retrieving QR SDSS and 2MASS catalogs for:', img
    for key in odi.OTA_dictionary:
        ota = odi.OTA_dictionary[key]
        outputsd = odi.sdsspath+'offline_'+ota+'.'+str(img[16:-5])+'.sdss'
        if not os.path.isfile(outputsd):
            x,y = odi.get_sdss_coords_offline(img,ota,inst,output=outputsd)
            output2m = odi.twomasspath+'offline_'+ota+'.'+str(img[16:-5])+'.mass'
            x,y = odi.get_2mass_coords_offline(img,ota,inst,output=output2m)
        if gaia_flag == True:
            outputg = odi.gaiapath+'offline_'+ota+'.'+str(img[16:-5])+'.gaia'
            if not os.path.isfile(outputg):
                if cluster_flag == True:
                    odi.get_gaia_coords(img,ota,inst,
                                        output=outputg,
                                        cluster=cluster_flag,
                                        racenter=float(ra_center),
                                        deccenter=float(dec_center),
                                        min_radius=float(min_radius),
                                        G_lim = 20.75)
                else:
                    odi.get_gaia_coords(img,ota,inst,
                                        output=outputg,
                                        cluster=cluster_flag)

listfiles = glob.glob('*.lis')
if len(listfiles) == 0:
    odi.imcombine_lists(images_, filters)
else:
    print 'imcombine lists done'

if not os.path.isfile('bpms.done'):
    for img in images_:
        print 'updating bpms for', img
        for key in tqdm(odi.OTA_dictionary):
            ota = odi.OTA_dictionary[key]
            odi.make_bpms(img, ota)

listfiles = glob.glob(odi.skyflatpath+'*.med.fits')
if len(listfiles) == 0:
    for filter in filters:
        odi.dark_sky_flat(filter)
else:
    print 'dark sky flats done'

if not os.path.isfile('derived_props.txt'):
    f1 = open('derived_props.txt','w+')
    print >> f1, '# img  ota  filter fwhm  zp_med  zp_std  bg_mean  bg_med  bg_std'
    for img in images_:
        for key in tqdm(odi.OTA_dictionary):
            ota = odi.OTA_dictionary[key]
            hdulist = odi.fits.open(img)
            hdr = hdulist[0].header
            filt = hdr['filter']
            image_to_correct = img+'['+ota+']'
            correction_image = ota+'.'+filt+'.med.fits'
            corrected_image = 'illcor_'+ota+'.'+str(img[16:])
            if not os.path.isfile(odi.illcorpath+corrected_image):
                odi.illumination_corrections(image_to_correct, correction_image, corrected_image, do_correction=illcor_flag)
            gaps = odi.get_gaps(img, ota)
            reprojed_image = 'reproj_'+ota+'.'+str(img[16:])
            if not os.path.isfile(odi.reprojpath+reprojed_image):
                if wcs_flag:
                    pixcrd3 = odi.list_wcs_coords(img, ota, gaps, inst,output=img[:-5]+'.'+ota+'.radec.coo', gmaglim=23., stars_only=True, offline = True, source = source)
                    try:
                        odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                    except:
                        try:
                            print 'msccmatch failed, wait a second and try again'
                            time.sleep(1.0)
                            odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                        except:
                            print 'there might be too few stars for msccmatch, just skip it.'
                if reproject_flag:
                    wcsref = odi.illcorpath+'illcor_OTA33.SCI.'+str((images_[0])[16:])
                    odi.reproject_ota(img, ota, rad, decd, wcsref)
            gaps = odi.get_gaps_rep(img, ota)
            odi.refetch_sdss_coords(img, ota, gaps, inst,gmaglim=21.5,offline = True,source=source)
            #run an additional refetch to get the xy for 2mass so they can be used for scaling
            odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='twomass')
            odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='gaia')
            fwhm = odi.getfwhm_ota(img, ota)
            if 'odi_NB695' in filters:
                zp_med, zp_std = 99.99,99.99
            elif source == 'sdss':
                zp_med, zp_std = 99.99,99.99
            elif source == 'twomass':
                zp_med, zp_std = 99.99,99.99
            elif source == 'gaia':
                zp_med, zp_std = 99.99,99.99
            if not os.path.isfile(odi.bgsubpath+'bgsub_'+ota+'.'+str(img[16:])):
                bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=True)
            else:
                bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=False)
            print >> f1, img[16], ota, filt, fwhm, zp_med, zp_std, bg_mean, bg_median, bg_std
            dim_stats = odi.check_mask_dim(img,ota)
            if not dim_stats:
                print 'mask dimensions do not match image'
                print 'redo', img, ota
                raise ValueError
    f1.close()
else:
    imgnum,fwhm,zp_med, zp_std, bg_mean, bg_median, bg_std = np.loadtxt('derived_props.txt',usecols=(0,3,4,5,6,7,8),unpack=True)
    ota_d, filt_d = np.loadtxt('derived_props.txt',usecols=(1,2),unpack=True,dtype=str)
    finished = zip(imgnum,ota_d,filt_d)
    f1 = open('derived_props.txt','a+')
    for img in images_:
        for key in tqdm(odi.OTA_dictionary):
            ota = odi.OTA_dictionary[key]
            hdulist = odi.fits.open(img)
            hdr = hdulist[0].header
            filt = hdr['filter']
            finishcheck = (int(str(img[16])),ota,filt)
            if finishcheck in finished:
                already = 0
            else:
                image_to_correct = img+'['+ota+']'
                correction_image = ota+'.'+filt+'.med.fits'
                corrected_image = 'illcor_'+ota+'.'+str(img[16:])
                if not os.path.isfile(odi.illcorpath+corrected_image):
                    odi.illumination_corrections(image_to_correct, correction_image, corrected_image, do_correction=illcor_flag)
                gaps = odi.get_gaps(img, ota)
                reprojed_image = 'reproj_'+ota+'.'+str(img[16:])
                if not os.path.isfile(odi.reprojpath+reprojed_image):
                    if wcs_flag:
                        pixcrd3 = odi.list_wcs_coords(img, ota, gaps, inst,output=img[:-5]+'.'+ota+'.radec.coo', gmaglim=23., stars_only=True, offline = True, source = source)
                        try:
                            odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                        except IndexError:
                            print 'not enough stars to fix wcs, skipping for this ota:', img, ota
                        except:
                            print 'msccmatch failed, wait a second and try again'
                            time.sleep(1.0)
                            odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                    if reproject_flag:
                        wcsref = odi.illcorpath+'illcor_OTA33.SCI.'+str((images_[0])[16:])
                        odi.reproject_ota(img, ota, rad, decd, wcsref)
                gaps = odi.get_gaps_rep(img, ota)
                odi.refetch_sdss_coords(img, ota, gaps, inst,gmaglim=21.5,offline = True,source=source)
                #run an additional refetch to get the xy for 2mass so they can be used for scaling
                odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='twomass')
                odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='gaia')
                fwhm = odi.getfwhm_ota(img, ota)
                if source == 'sdss':
                    zp_med, zp_std = 99.99,99.99
                if source == 'twomass':
                    zp_med, zp_std = 99.99,99.99
                elif source == 'gaia':
                    zp_med, zp_std = 99.99,99.99
                if not os.path.isfile(odi.bgsubpath+'bgsub_'+ota+'.'+str(img[16:])):
                    bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=True)
                else:
                    bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=False)
                print >> f1, img[16], ota, filt, fwhm, zp_med, zp_std, bg_mean, bg_median, bg_std
                dim_stats = odi.check_mask_dim(img,ota)
                if not dim_stats:
                    print 'mask dimensions do not match image'
                    print 'redo', img, ota
                    raise ValueError
    f1.close()
<commit_msg>turn off 2mass offline cats<commit_after>#!/usr/bin/env python

import sys, os, glob, string
import numpy as np
import astropy as ast
import matplotlib.pyplot as plt
from pyraf import iraf
from tqdm import tqdm
import odi_config as odi
import glob
import shutil
import pandas as pd
import time

try:
    object_str, filters, instrument, images, illcor_flag, skyflat_src, wcs_flag, reproject_flag, scale_flag, stack_flag, gaia_flag, cluster_flag, ra_center, dec_center, min_radius = odi.cfgparse('config.yaml')
except IOError:
    print 'config.yaml does not exist, quitting...'
    exit()

# for basic processing, filter shouldn't matter enough to consider separately
# or rather, this script already handles that just fine
# so just stick all the image names together into one long list
images_ = []
for filt in images:
    for key in images[filt]:
        images_.append(images[filt][key])

rad, decd = odi.get_targ_ra_dec(images_[0], 'OTA33.SCI')
source = 'sdss'
inst = odi.instrument(images_[0])
#source = 'twomass'

#Create offline catalogs
for img in images_:
    print 'Retrieving QR SDSS and 2MASS catalogs for:', img
    for key in odi.OTA_dictionary:
        ota = odi.OTA_dictionary[key]
        outputsd = odi.sdsspath+'offline_'+ota+'.'+str(img[16:-5])+'.sdss'
        if not os.path.isfile(outputsd):
            x,y = odi.get_sdss_coords_offline(img,ota,inst,output=outputsd)
            output2m = odi.twomasspath+'offline_'+ota+'.'+str(img[16:-5])+'.mass'
            # x,y = odi.get_2mass_coords_offline(img,ota,inst,output=output2m)
        if gaia_flag == True:
            outputg = odi.gaiapath+'offline_'+ota+'.'+str(img[16:-5])+'.gaia'
            if not os.path.isfile(outputg):
                if cluster_flag == True:
                    odi.get_gaia_coords(img,ota,inst,
                                        output=outputg,
                                        cluster=cluster_flag,
                                        racenter=float(ra_center),
                                        deccenter=float(dec_center),
                                        min_radius=float(min_radius),
                                        G_lim = 20.75)
                else:
                    odi.get_gaia_coords(img,ota,inst,
                                        output=outputg,
                                        cluster=cluster_flag)

listfiles = glob.glob('*.lis')
if len(listfiles) == 0:
    odi.imcombine_lists(images_, filters)
else:
    print 'imcombine lists done'

if not os.path.isfile('bpms.done'):
    for img in images_:
        print 'updating bpms for', img
        for key in tqdm(odi.OTA_dictionary):
            ota = odi.OTA_dictionary[key]
            odi.make_bpms(img, ota)

listfiles = glob.glob(odi.skyflatpath+'*.med.fits')
if len(listfiles) == 0:
    for filter in filters:
        odi.dark_sky_flat(filter)
else:
    print 'dark sky flats done'

if not os.path.isfile('derived_props.txt'):
    f1 = open('derived_props.txt','w+')
    print >> f1, '# img  ota  filter fwhm  zp_med  zp_std  bg_mean  bg_med  bg_std'
    for img in images_:
        for key in tqdm(odi.OTA_dictionary):
            ota = odi.OTA_dictionary[key]
            hdulist = odi.fits.open(img)
            hdr = hdulist[0].header
            filt = hdr['filter']
            image_to_correct = img+'['+ota+']'
            correction_image = ota+'.'+filt+'.med.fits'
            corrected_image = 'illcor_'+ota+'.'+str(img[16:])
            if not os.path.isfile(odi.illcorpath+corrected_image):
                odi.illumination_corrections(image_to_correct, correction_image, corrected_image, do_correction=illcor_flag)
            gaps = odi.get_gaps(img, ota)
            reprojed_image = 'reproj_'+ota+'.'+str(img[16:])
            if not os.path.isfile(odi.reprojpath+reprojed_image):
                if wcs_flag:
                    pixcrd3 = odi.list_wcs_coords(img, ota, gaps, inst,output=img[:-5]+'.'+ota+'.radec.coo', gmaglim=23., stars_only=True, offline = True, source = source)
                    try:
                        odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                    except:
                        try:
                            print 'msccmatch failed, wait a second and try again'
                            time.sleep(1.0)
                            odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                        except:
                            print 'there might be too few stars for msccmatch, just skip it.'
                if reproject_flag:
                    wcsref = odi.illcorpath+'illcor_OTA33.SCI.'+str((images_[0])[16:])
                    odi.reproject_ota(img, ota, rad, decd, wcsref)
            gaps = odi.get_gaps_rep(img, ota)
            odi.refetch_sdss_coords(img, ota, gaps, inst,gmaglim=21.5,offline = True,source=source)
            #run an additional refetch to get the xy for 2mass so they can be used for scaling
            odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='twomass')
            odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='gaia')
            fwhm = odi.getfwhm_ota(img, ota)
            if 'odi_NB695' in filters:
                zp_med, zp_std = 99.99,99.99
            elif source == 'sdss':
                zp_med, zp_std = 99.99,99.99
            elif source == 'twomass':
                zp_med, zp_std = 99.99,99.99
            elif source == 'gaia':
                zp_med, zp_std = 99.99,99.99
            if not os.path.isfile(odi.bgsubpath+'bgsub_'+ota+'.'+str(img[16:])):
                bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=True)
            else:
                bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=False)
            print >> f1, img[16], ota, filt, fwhm, zp_med, zp_std, bg_mean, bg_median, bg_std
            dim_stats = odi.check_mask_dim(img,ota)
            if not dim_stats:
                print 'mask dimensions do not match image'
                print 'redo', img, ota
                raise ValueError
    f1.close()
else:
    imgnum,fwhm,zp_med, zp_std, bg_mean, bg_median, bg_std = np.loadtxt('derived_props.txt',usecols=(0,3,4,5,6,7,8),unpack=True)
    ota_d, filt_d = np.loadtxt('derived_props.txt',usecols=(1,2),unpack=True,dtype=str)
    finished = zip(imgnum,ota_d,filt_d)
    f1 = open('derived_props.txt','a+')
    for img in images_:
        for key in tqdm(odi.OTA_dictionary):
            ota = odi.OTA_dictionary[key]
            hdulist = odi.fits.open(img)
            hdr = hdulist[0].header
            filt = hdr['filter']
            finishcheck = (int(str(img[16])),ota,filt)
            if finishcheck in finished:
                already = 0
            else:
                image_to_correct = img+'['+ota+']'
                correction_image = ota+'.'+filt+'.med.fits'
                corrected_image = 'illcor_'+ota+'.'+str(img[16:])
                if not os.path.isfile(odi.illcorpath+corrected_image):
                    odi.illumination_corrections(image_to_correct, correction_image, corrected_image, do_correction=illcor_flag)
                gaps = odi.get_gaps(img, ota)
                reprojed_image = 'reproj_'+ota+'.'+str(img[16:])
                if not os.path.isfile(odi.reprojpath+reprojed_image):
                    if wcs_flag:
                        pixcrd3 = odi.list_wcs_coords(img, ota, gaps, inst,output=img[:-5]+'.'+ota+'.radec.coo', gmaglim=23., stars_only=True, offline = True, source = source)
                        try:
                            odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                        except IndexError:
                            print 'not enough stars to fix wcs, skipping for this ota:', img, ota
                        except:
                            print 'msccmatch failed, wait a second and try again'
                            time.sleep(1.0)
                            odi.fix_wcs(img, ota, coords=img[:-5]+'.'+ota+'.radec.coo', iters=1)
                    if reproject_flag:
                        wcsref = odi.illcorpath+'illcor_OTA33.SCI.'+str((images_[0])[16:])
                        odi.reproject_ota(img, ota, rad, decd, wcsref)
                gaps = odi.get_gaps_rep(img, ota)
                odi.refetch_sdss_coords(img, ota, gaps, inst,gmaglim=21.5,offline = True,source=source)
                #run an additional refetch to get the xy for 2mass so they can be used for scaling
                # odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='twomass')
                odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='gaia')
                fwhm = odi.getfwhm_ota(img, ota)
                if source == 'sdss':
                    zp_med, zp_std = 99.99,99.99
                if source == 'twomass':
                    zp_med, zp_std = 99.99,99.99
                elif source == 'gaia':
                    zp_med, zp_std = 99.99,99.99
                if not os.path.isfile(odi.bgsubpath+'bgsub_'+ota+'.'+str(img[16:])):
                    bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=True)
                else:
                    bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=False)
                print >> f1, img[16], ota, filt, fwhm, zp_med, zp_std, bg_mean, bg_median, bg_std
                dim_stats = odi.check_mask_dim(img,ota)
                if not dim_stats:
                    print 'mask dimensions do not match image'
                    print 'redo', img, ota
                    raise ValueError
    f1.close()
<|endoftext|>


<commit_before>"""
Documentation Decorators
~~~~~~~~~~~~~~~~~~~~~~~~

Additional decorators for improving documentation of APIs.

"""
from odin.utils import getmeta

from . import _compat
from .constants import *

__all__ = (
    # Docs
    'OperationDoc', 'operation', 'parameter', 'response', 'produces'
)

class OperationDoc(object):
    """
    Utility class for building/storing documentation on callback endpoints.
    """
    @classmethod
    def get(cls, func):
        docs = getattr(func, '_api_docs', None)
        if docs is None:
            docs = cls(func)
            setattr(func, '_api_docs', docs)
        return docs

    __slots__ = 'callback _parameters summary deprecated tags responses produces'.split()

    def __init__(self, callback):
        self.callback = callback
        self.summary = None
        self.deprecated = False
        self.tags = set()
        self._parameters = {}
        self.responses = {
            'default': {
                'description': 'Error',
                'schema': {'$ref': '#/definitions/Error'}
            }
        }
        self.produces = set()

    def add_parameter(self, name, in_, **options):
        # Ensure there are no duplicates
        param = self._parameters.setdefault("{}:{}".format(in_, name), {})
        param['name'] = name
        param['in'] = in_
        param.update((k, v) for k, v in options.items() if v is not None)

    def add_response(self, status, description, resource=None):
        response_spec = {'description': description}
        if resource:
            response_spec['schema'] = {'$ref': '#/definitions/{}'.format(getmeta(resource).resource_name)}

        self.responses[status] = response_spec

    @property
    def parameters(self):
        return list(self._parameters.values())

    @property
    def description(self):
        return self.callback.__doc__.strip()

    def to_dict(self):
        d = {
            "operationId": self.callback.__name__,
            "description": (self.callback.__doc__ or '').strip(),
        }
        if self.deprecated:
            d['deprecated'] = True
        if self.produces:
            d['produces'] = list(self.produces)
        if self.tags:
            d['tags'] = list(self.tags)
        if self.responses:
            d['responses'] = self.responses
        if self.parameters:
            d['parameters'] = self.parameters

        return d

def operation(summary=None, tags=None, deprecated=False):
    """
    Decorator for applying operation documentation to a callback.

    The values are based off `Swagger <https://swagger.io/specification>`_.

    """
    def inner(func):
        docs = OperationDoc.get(func)
        docs.summary = summary
        docs.tags.update(tags)
        docs.deprecated = deprecated
        return func
    return inner

def parameter(name, in_, required=None, type_=None, default=None):
    """
    Decorator for applying parameter documentation to a callback.

    The values are based off `Swagger <https://swagger.io/specification>`_.

    """
    if in_ not in In:
        raise ValueError("In parameter not a valid value.")

    def inner(func):
        OperationDoc.get(func).add_parameter(name, in_.value, required=required, type=type_.value, default=default)
        return func
    return inner

def body_param(resource=None, description=None):
    """
    Decorator for defining request body
    """
    def inner(func):
        args = {
            'description': description,
        }
        if resource:
            args['schema'] = {'$ref': '#/definitions/{}'.format(getmeta(resource).resource_name)}
        OperationDoc.get(func).add_parameter('', In.Body.value, **args)
        return func
    return inner

def response(status, description, resource=None):
    """
    Define an expected responses.

    The values are based off `Swagger <https://swagger.io/specification>`_.

    """
    def inner(func):
        OperationDoc.get(func).add_response(status, description, resource)
        return func
    return inner

def produces(*content_types):
    """
    Define content types produced by an endpoint.
    """
    if not all(isinstance(content_type, _compat.string_types) for content_type in content_types):
        raise ValueError("In parameter not a valid value.")

    def inner(func):
        OperationDoc.get(func).produces.update(content_types)
        return func
    return inner
<commit_msg>Updated doc decorators to reflect more of the swagger spec<commit_after>"""
Documentation Decorators
~~~~~~~~~~~~~~~~~~~~~~~~

Additional decorators for improving documentation of APIs.

"""
from odin.utils import getmeta
from odinweb.utils import dict_update_values, dict_filter_update

from . import _compat
from .constants import *

__all__ = (
    # Docs
    'OperationDoc', 'operation',
    'query_param', 'body_param', 'path_param', 'header_param',
    'response', 'produces'
)

class OperationDoc(object):
    """
    Utility class for building/storing documentation on callback endpoints.
    """
    @classmethod
    def get(cls, func):
        docs = getattr(func, '_api_docs', None)
        if docs is None:
            docs = cls(func)
            setattr(func, '_api_docs', docs)
        return docs

    __slots__ = ('callback', 'summary', 'tags', 'consumes', 'produces', 'parameters', 'responses', 'deprecated')

    def __init__(self, callback):
        self.callback = callback
        self.summary = None
        self.tags = set()
        self.consumes = {}
        self.produces = set()
        self.parameters = {}
        self.responses = {
            'default': {
                'description': 'Error',
                'schema': {'$ref': '#/definitions/Error'}
            }
        }
        self.deprecated = None

    def add_parameter(self, name, in_, **options):
        # Ensure there are no duplicates
        param = self.parameters.setdefault("{}:{}".format(in_, name), {'name': name, 'in': in_})
        dict_filter_update(param, options)

    def add_response(self, status, description, resource=None):
        response_spec = {'description': description}
        if resource:
            response_spec['schema'] = {'$ref': '#/definitions/{}'.format(getmeta(resource).resource_name)}

        self.responses[status] = response_spec

    def to_dict(self):
        return dict_update_values(
            operationId=self.callback.__name__,
            description=(self.callback.__doc__ or '').strip(),
            deprecated=True if self.deprecated else None,
            parameters=self.parameters.values() if self.parameters else None,
            responses=self.responses if self.responses else None,
            produces=list(self.produces) if self.produces else None,
            tags=list(self.tags) if self.tags else None,
        )

def operation(summary=None, tags=None, deprecated=False):
    """
    Decorator for applying operation documentation to a callback.

    The values are based off `Swagger <https://swagger.io/specification>`_.

    """
    def inner(func):
        docs = OperationDoc.get(func)
        docs.summary = summary
        docs.tags.update(tags)
        docs.deprecated = deprecated
        return func
    return inner

def query_param(name, type, description=None, required=False,
                default=None, minimum=None, maximum=None, enum=None, **options):
    """
    Query parameter documentation.

    :param name:
    :param type:
    :param description:
    :param default:
    :param minimum:
    :param maximum:
    :param enum:

    """
    def inner(func):
        OperationDoc.get(func).add_parameter(
            name, In.Query.value, type=type.value, description=description, required=required or None,
            default=default, minimum=minimum, maximum=maximum, enum=enum, **options)
        return func
    return inner

def path_param(name, type, description=None,
               default=None, minimum=None, maximum=None, enum=None, **options):
    """
    Path parameter documentation.

    :param name:
    :param type:
    :param description:
    :param default:
    :param minimum:
    :param maximum:
    :param enum:
    :param options:

    """
    def inner(func):
        OperationDoc.get(func).add_parameter(
            name, In.Path.value, type=type.value, description=description,
            default=default, minimum=minimum, maximum=maximum, enum=enum, **options
        )
        return func
    return inner

def body_param(resource=None, description=None,
               default=None, minimum=None, maximum=None, enum=None, **options):
    """
    Body parameter documentation.

    :param resource:
    :param description:
    :param default:
    :param minimum:
    :param maximum:
    :param enum:
    :param options:

    """
    def inner(func):
        schema = {'$ref': '#/definitions/{}'.format(getmeta(resource).resource_name)} if resource else None
        OperationDoc.get(func).add_parameter(
            '', In.Body.value, description=description, schema=schema,
            default=default, minimum=minimum, maximum=maximum, enum=enum, **options
        )
        return func
    return inner

def header_param(name, type, description=None, required=False, **options):
    """
    Header parameter documentation.

    :param name:
    :param type:
    :param description:
    :param required:
    :param options:

    """
    def inner(func):
        OperationDoc.get(func).add_parameter(
            name, In.Query.value, type=type.value, description=description, required=required or None,
            **options)
        return func
    return inner

def response(status, description, resource=None):
    """
    Define an expected responses.

    The values are based off `Swagger <https://swagger.io/specification>`_.

    """
    def inner(func):
        OperationDoc.get(func).add_response(status, description, resource)
        return func
    return inner

def produces(*content_types):
    """
    Define content types produced by an endpoint.
    """
    if not all(isinstance(content_type, _compat.string_types) for content_type in content_types):
        raise ValueError("In parameter not a valid value.")

    def inner(func):
        OperationDoc.get(func).produces.update(content_types)
        return func
    return inner
<|endoftext|>


<commit_before>__version__ = '0.1.0.dev1'
<commit_msg>set to version 0.2.0 for pre release<commit_after>__version__ = '0.2.0'
<|endoftext|>


<commit_before>from collections import OrderedDict, defaultdict

from z3 import eq
import claripy

try:
    import pycparser
except ImportError:
    pycparser = None

class SimType(object):
    '''
    SimType exists to track type information for SimProcedures.
    '''

    _fields = ()
    base = True

    def __init__(self, label=None):
        '''
        @param label: the type label
        '''
        self.label = label

    def __eq__(self, other):
        if type(self) != type(other):
            return False

        for attr in self._fields:
            if getattr(self, attr) != getattr(other, attr):
                return False

        return True

    def __ne__(self, other):
        # wow many efficient
        return not self == other

    def __hash__(self):
        # very hashing algorithm many secure wow
        out = hash(type(self))
        for attr in self._fields:
            out ^= hash(getattr(self, attr))
        return out

    def view(self, state, addr):
        return SimMemView(ty=self, addr=addr, state=state)

    @property
    def name(self):
        return repr(self)

    def _refine_dir(self):
        return []

    def _refine(self, view, k):
        if k == 'array':
            return lambda length: view._deeper(ty=SimTypeFixedSizeArray(self, length))
        raise KeyError("{} is not a valid refinement".format(k))

class SimTypeBottom(SimType):
    '''
    SimTypeBottom basically repesents a type error.
    '''

    def __repr__(self):
        return 'BOT'

class SimTypeTop(SimType):
    '''
    SimTypeTop represents any type (mostly used with a pointer for void*).
    '''

    _fields = ('size',)

    def __init__(self, size=None):
        SimType.__init__(self)
        self.size = size

    def __repr__(self):
        return 'TOP'

class SimTypeReg(SimType):
    '''
    SimTypeReg is the base type for all types that are register-sized.
    '''

    _fields = ('size',)

    def __init__(self, size, label=None):
        '''
        @param label: the type label
        @param size: the size of the type (e.g. 32bit, 8bit, etc.)
        '''
        SimType.__init__(self, label=label)
        self.size = size

    def __repr__(self):
        return "reg{}_t".format(self.size)

    def extract(self, state, addr):
        return state.memory.load(addr, self.size / 8, endness=state.arch.memory_endness)

    def store(self, state, addr, value):
        if isinstance(value, claripy.ast.BV):
            if value.size() != self.size:
                raise ValueError("size of expression is wrong size for type")
        elif isinstance(value, (int, long)):
            value = state.se.BVV(value, self.size)
        else:
            raise TypeError("unrecognized expression type for SimType {}".format(type(self).__name__))

        state.memory.store(addr, value, endness=state.arch.memory_endness)

class SimTypeInt(SimTypeReg):
    '''
    SimTypeInt is a type that specifies a signed or unsigned integer of some size.
    '''

    _fields = SimTypeReg._fields + ('signed',)

    def __init__(self, size, signed, label=None):
        '''
        @param label: the type label
        @param size: the size of the integer (e.g. 32bit, 8bit, etc.)
        @param signed: True if signed, False if unsigned
        '''
        SimTypeReg.__init__(self, size, label=label)
        self.signed = signed

    def __repr__(self):
        return ('' if self.signed else 'u') + 'int{}_t'.format(self.size)

class SimTypeChar(SimTypeReg):
    '''
    SimTypeChar is a type that specifies a character;
    this could be represented by an 8-bit int, but this is meant to be interpreted as a character.
    '''

    _fields = SimTypeReg._fields

    def __init__(self, label=None):
        '''
        @param label: the type label
        '''
        SimTypeReg.__init__(self, 8, label=label) # a char better be 8 bits (I'm looking at you, DCPU-16)
        self.signed = False

    def __repr__(self):
        return 'char'

    def store(self, state, addr, value):
        try:
            super(SimTypeChar, self).store(state, addr, value)
        except TypeError:
            if isinstance(value, str) and len(value) == 1:
                value = state.se.BVV(ord(value), 8)
                super(SimTypeChar, self).store(state, addr, value)
            else:
                raise

class SimTypeFd(SimTypeReg):
    '''
    SimTypeFd is a type that specifies a file descriptor.
    '''

    _fields = SimTypeReg._fields

    def __init__(self, label=None):
        '''
        @param label: the type label
        '''
        # file descriptors are always 32 bits, right?
        SimTypeReg.__init__(self, 32, label=label)

    def __repr__(self):
        return 'fd_t'

class SimTypePointer(SimTypeReg):
    '''
    SimTypePointer is a type that specifies a pointer to some other type.
    '''

    _fields = SimTypeReg._fields + ('pts_to',)

    def __init__(self, arch, pts_to, label=None):
        '''
        @param label: the type label
        @param pts_to: the type to which this pointer points to
        '''
        SimTypeReg.__init__(self, arch.bits, label=label)
        self._arch = arch
        self.pts_to = pts_to
        self.signed = False

    def __repr__(self):
        return '{}*'.format(self.pts_to)

    def make(self, pts_to):
        new = type(self)(self._arch, pts_to)
        return new

class SimTypeFixedSizeArray(SimType):
    '''
    SimTypeFixedSizeArray is a literal (i.e. not a pointer) fixed-size array.
    '''

    def __init__(self, elem_type, length):
        SimType.__init__(self)
        self.elem_type = elem_type
        self.length = length
        self.size = elem_type.size * length

    def __repr__(self):
        return '{}[{}]'.format(self.elem_type, self.length)

    def extract(self, state, addr):
        return [self.elem_type.extract(state, addr + i*self.elem_type.size) for i in xrange(self.length)]

    def store(self, state, addr, values):
        for i, val in enumerate(values):
            self.elem_type.store(state, addr + i*self.elem_type.size, val)

class SimTypeArray(SimType):
    '''
    SimTypeArray is a type that specifies a pointer to an array; while it is a pointer, it has a semantic difference.
    '''

    _fields = ('elem_type', 'length')

    def __init__(self, elem_type, length=None, label=None):
        '''
        @param label: the type label
        @param elem_type: the type of each element in the array
        @param length: an expression of the length of the array, if known
        '''
        SimType.__init__(self, label=label)
        self.elem_type = elem_type
        self.length = length

    def __repr__(self):
        return '{}[{}]'.format(self.elem_type, '' if self.length is None else self.length)

class SimTypeString(SimTypeArray):
    '''
    SimTypeString is a type that represents a C-style string,
    i.e. a NUL-terminated array of bytes.
    '''

    _fields = SimTypeArray._fields + ('length',)

    def __init__(self, length=None, label=None):
        '''
        @param label: the type label
        @param length: an expression of the length of the string, if known
        '''
        SimTypeArray.__init__(self, SimTypeChar(), label=label)
        self.length = length

    def __repr__(self):
        return 'string_t'

    def extract(self, state, addr):
        mem = state.memory.load(addr, self.length)
        if state.se.symbolic(mem):
            return mem
        else:
            return repr(state.se.any_str(mem))

class SimTypeFunction(SimType):
    '''
    SimTypeFunction is a type that specifies an actual function (i.e. not a pointer) with certain types of arguments and a certain return value.
    '''

    _fields = ('args', 'returnty')
    base = False

    def __init__(self, args, returnty, label=None):
        '''
        @param label: the type label
        @param args: a tuple of types representing the arguments to the function
        @param returnty: the return type of the function
        '''
        SimType.__init__(self, label=label)
        self.args = args
        self.returnty = returnty

    def __repr__(self):
        return '({}) -> {}'.format(', '.join(str(a) for a in self.args), self.returnty)

class SimTypeLength(SimTypeInt):
    '''
    SimTypeLength is a type that specifies the length of some buffer in memory.
    '''

    _fields = SimTypeInt._fields + ('addr', 'length') # ?

    def __init__(self, arch, addr=None, length=None, label=None):
        '''
        @param label: the type label
        @param addr: the memory address (expression)
        @param length: the length (expression)
        '''
        SimTypeInt.__init__(self, arch.bits, False, label=label)
        self.addr = addr
        self.length = length
        self.signed = False

    def __repr__(self):
        return 'size_t'

class SimStructValue(object):
    def __init__(self, struct, values=None):
        self._struct = struct
        self._values = defaultdict(lambda: None, values or ())

    def __repr__(self):
        fields = ('.{} = {}'.format(name, self._values[name]) for name in self._struct.fields)
        return '{{\n  {}\n}}'.format(',\n  '.join(fields))

_C_TYPE_TO_SIMTYPE = {
    ('int',): lambda _: SimTypeInt(32, True),
    ('unsigned', 'int'): lambda _: SimTypeInt(32, False),
    ('long',): lambda arch: SimTypeInt(arch.bits, True),
    ('unsigned', 'long'): lambda arch: SimTypeInt(arch.bits, False),
    ('char',): lambda _: SimTypeChar(),
    ('int8_t',): lambda _: SimTypeInt(8, True),
    ('uint8_t',): lambda _: SimTypeInt(8, False),
    ('int16_t',): lambda _: SimTypeInt(16, True),
    ('uint16_t',): lambda _: SimTypeInt(16, False),
    ('int32_t',): lambda _: SimTypeInt(32, True),
    ('uint32_t',): lambda _: SimTypeInt(32, False),
    ('int64_t',): lambda _: SimTypeInt(64, True),
    ('uint64_t',): lambda _: SimTypeInt(64, False),
    ('ptrdiff_t',): lambda arch: SimTypeInt(arch.bits, False),
    ('size_t',): lambda arch: SimTypeInt(arch.bits, False),
    ('ssize_t',): lambda arch: SimTypeInt(arch.bits, True),
    ('uintptr_t',) : lambda arch: SimTypeInt(arch.bits, False),
}

def _decl_to_type(decl):
    if isinstance(decl, pycparser.c_ast.TypeDecl):
        return _C_TYPE_TO_SIMTYPE[tuple(decl.type.names)]
    elif isinstance(decl, pycparser.c_ast.PtrDecl):
        pts_to = _decl_to_type(decl.type)
        return lambda arch: SimTypePointer(arch, pts_to(arch))
    elif isinstance(decl, pycparser.c_ast.ArrayDecl):
        elem_type = _decl_to_type(decl.type)
        size = int(decl.dim.value)
        return lambda arch: SimTypeFixedSizeArray(elem_type(arch), size)

# these are all bogus, on purpose
_C_STRUCT_PREAMBLE = """
typedef int int8_t;
typedef int uint8_t;
typedef int int16_t;
typedef int uint16_t;
typedef int int32_t;
typedef int uint32_t;
typedef int int64_t;
typedef int uint64_t;
"""

class SimStruct(SimType):
    _fields = ('name', 'fields')

    def __init__(self, name, fields, pack=True):
        if not pack:
            raise ValueError("you think I've implemented padding, how cute")

        self._name = name
        self.fields = fields
        self._arch_offsets_cache = {}

    def _arch_offsets(self, arch):
        if arch in self._arch_offsets_cache:
            return self._arch_offsets_cache[arch]

        offsets = {}
        offset_so_far = 0
        for name, almost_ty in self.fields.iteritems():
            ty = almost_ty(arch)
            offsets[name] = (ty, offset_so_far)
            offset_so_far += ty.size / 8

        self._arch_offsets_cache[arch] = offsets
        return offsets

    def extract(self, state, addr):
        values = {name: ty.view(state, addr + offset)
                  for (name, (ty, offset))
                  in self._arch_offsets(state.arch).iteritems()}
        return SimStructValue(self, values=values)

    @property
    def name(self):
        return self._name

    @classmethod
    def from_c(cls, defn):
        if pycparser is None:
            raise ImportError("pycparser is needed to use SimStruct.from_c!")

        # if preprocess:
        #     defn = subprocess.Popen(['cpp'], stdin=subprocess.PIPE, stdout=subprocess.PIPE).communicate(input=defn)[0]

        node = pycparser.c_parser.CParser().parse(_C_STRUCT_PREAMBLE + defn)

        if not isinstance(node, pycparser.c_ast.FileAST) or \
           not isinstance(node.ext[-1], pycparser.c_ast.Decl) or \
           not isinstance(node.ext[-1].type, pycparser.c_ast.Struct):
            raise ValueError("invalid struct definition")

        struct = node.ext[-1].type
        fields = OrderedDict((decl.name, _decl_to_type(decl.type)) for decl in struct.decls)

        return cls(struct.name, fields)

    def _refine_dir(self):
        return self.fields.keys()

    def _refine(self, view, k):
        ty, offset = self._arch_offsets(view.state.arch)[k]
        return view._deeper(ty=ty, addr=view._addr + offset)

try:
    _example_struct = SimStruct.from_c("""
struct example {
  int foo;
  int bar;
  char *hello;
};
""")
except ImportError:
    _example_struct = None

ALL_TYPES = {
    'char': lambda _: SimTypeInt(8, True),
    'int8_t': lambda _: SimTypeInt(8, True),
    'uchar': lambda _: SimTypeInt(8, False),
    'uint8_t': lambda _: SimTypeInt(8, False),
    'byte': lambda _: SimTypeInt(8, False),

    'short': lambda _: SimTypeInt(16, True),
    'int16_t': lambda _: SimTypeInt(16, True),
    'ushort': lambda _: SimTypeInt(16, False),
    'uint16_t': lambda _: SimTypeInt(16, False),
    'word': lambda _: SimTypeInt(16, False),

    'int': lambda _: SimTypeInt(32, True),
    'int32_t': lambda _: SimTypeInt(32, True),
    'uint': lambda _: SimTypeInt(32, False),
    'uint32_t': lambda _: SimTypeInt(32, False),
    'dword': lambda _: SimTypeInt(32, False),

    'long': lambda _: SimTypeInt(64, True),
    'int64_t': lambda _: SimTypeInt(64, True),
    'ulong': lambda _: SimTypeInt(64, False),
    'uint64_t': lambda _: SimTypeInt(64, False),
    'qword': lambda _: SimTypeInt(64, False),

    'string': lambda arch: SimTypePointer(arch, SimTypeChar()),
    'example': lambda _: _example_struct,

}

def define_struct(defn):
    struct = SimStruct.from_c(defn)
    ALL_TYPES[struct.name] = lambda _: struct
    return struct

from .plugins.view import SimMemView
<commit_msg>SimType: allow storing strings directly.<commit_after>from collections import OrderedDict, defaultdict

from z3 import eq
import claripy

try:
    import pycparser
except ImportError:
    pycparser = None

class SimType(object):
    '''
    SimType exists to track type information for SimProcedures.
    '''

    _fields = ()
    base = True

    def __init__(self, label=None):
        '''
        @param label: the type label
        '''
        self.label = label

    def __eq__(self, other):
        if type(self) != type(other):
            return False

        for attr in self._fields:
            if getattr(self, attr) != getattr(other, attr):
                return False

        return True

    def __ne__(self, other):
        # wow many efficient
        return not self == other

    def __hash__(self):
        # very hashing algorithm many secure wow
        out = hash(type(self))
        for attr in self._fields:
            out ^= hash(getattr(self, attr))
        return out

    def view(self, state, addr):
        return SimMemView(ty=self, addr=addr, state=state)

    @property
    def name(self):
        return repr(self)

    def _refine_dir(self):
        return []

    def _refine(self, view, k):
        if k == 'array':
            return lambda length: view._deeper(ty=SimTypeFixedSizeArray(self, length))
        raise KeyError("{} is not a valid refinement".format(k))

class SimTypeBottom(SimType):
    '''
    SimTypeBottom basically repesents a type error.
    '''

    def __repr__(self):
        return 'BOT'

class SimTypeTop(SimType):
    '''
    SimTypeTop represents any type (mostly used with a pointer for void*).
    '''

    _fields = ('size',)

    def __init__(self, size=None):
        SimType.__init__(self)
        self.size = size

    def __repr__(self):
        return 'TOP'

class SimTypeReg(SimType):
    '''
    SimTypeReg is the base type for all types that are register-sized.
    '''

    _fields = ('size',)

    def __init__(self, size, label=None):
        '''
        @param label: the type label
        @param size: the size of the type (e.g. 32bit, 8bit, etc.)
        '''
        SimType.__init__(self, label=label)
        self.size = size

    def __repr__(self):
        return "reg{}_t".format(self.size)

    def extract(self, state, addr):
        return state.memory.load(addr, self.size / 8, endness=state.arch.memory_endness)

    def store(self, state, addr, value):
        store_endness = state.arch.memory_endness

        if isinstance(value, claripy.ast.BV):
            if value.size() != self.size:
                raise ValueError("size of expression is wrong size for type")
        elif isinstance(value, (int, long)):
            value = state.se.BVV(value, self.size)
        elif isinstance(value, str):
            store_endness = 'Iend_BE'
        else:
            raise TypeError("unrecognized expression type for SimType {}".format(type(self).__name__))

        state.memory.store(addr, value, endness=store_endness)

class SimTypeInt(SimTypeReg):
    '''
    SimTypeInt is a type that specifies a signed or unsigned integer of some size.
    '''

    _fields = SimTypeReg._fields + ('signed',)

    def __init__(self, size, signed, label=None):
        '''
        @param label: the type label
        @param size: the size of the integer (e.g. 32bit, 8bit, etc.)
        @param signed: True if signed, False if unsigned
        '''
        SimTypeReg.__init__(self, size, label=label)
        self.signed = signed

    def __repr__(self):
        return ('' if self.signed else 'u') + 'int{}_t'.format(self.size)

class SimTypeChar(SimTypeReg):
    '''
    SimTypeChar is a type that specifies a character;
    this could be represented by an 8-bit int, but this is meant to be interpreted as a character.
    '''

    _fields = SimTypeReg._fields

    def __init__(self, label=None):
        '''
        @param label: the type label
        '''
        SimTypeReg.__init__(self, 8, label=label) # a char better be 8 bits (I'm looking at you, DCPU-16)
        self.signed = False

    def __repr__(self):
        return 'char'

    def store(self, state, addr, value):
        try:
            super(SimTypeChar, self).store(state, addr, value)
        except TypeError:
            if isinstance(value, str) and len(value) == 1:
                value = state.se.BVV(ord(value), 8)
                super(SimTypeChar, self).store(state, addr, value)
            else:
                raise

class SimTypeFd(SimTypeReg):
    '''
    SimTypeFd is a type that specifies a file descriptor.
    '''

    _fields = SimTypeReg._fields

    def __init__(self, label=None):
        '''
        @param label: the type label
        '''
        # file descriptors are always 32 bits, right?
        SimTypeReg.__init__(self, 32, label=label)

    def __repr__(self):
        return 'fd_t'

class SimTypePointer(SimTypeReg):
    '''
    SimTypePointer is a type that specifies a pointer to some other type.
    '''

    _fields = SimTypeReg._fields + ('pts_to',)

    def __init__(self, arch, pts_to, label=None):
        '''
        @param label: the type label
        @param pts_to: the type to which this pointer points to
        '''
        SimTypeReg.__init__(self, arch.bits, label=label)
        self._arch = arch
        self.pts_to = pts_to
        self.signed = False

    def __repr__(self):
        return '{}*'.format(self.pts_to)

    def make(self, pts_to):
        new = type(self)(self._arch, pts_to)
        return new

class SimTypeFixedSizeArray(SimType):
    '''
    SimTypeFixedSizeArray is a literal (i.e. not a pointer) fixed-size array.
    '''

    def __init__(self, elem_type, length):
        SimType.__init__(self)
        self.elem_type = elem_type
        self.length = length
        self.size = elem_type.size * length

    def __repr__(self):
        return '{}[{}]'.format(self.elem_type, self.length)

    def extract(self, state, addr):
        return [self.elem_type.extract(state, addr + i*self.elem_type.size) for i in xrange(self.length)]

    def store(self, state, addr, values):
        for i, val in enumerate(values):
            self.elem_type.store(state, addr + i*self.elem_type.size, val)

class SimTypeArray(SimType):
    '''
    SimTypeArray is a type that specifies a pointer to an array; while it is a pointer, it has a semantic difference.
    '''

    _fields = ('elem_type', 'length')

    def __init__(self, elem_type, length=None, label=None):
        '''
        @param label: the type label
        @param elem_type: the type of each element in the array
        @param length: an expression of the length of the array, if known
        '''
        SimType.__init__(self, label=label)
        self.elem_type = elem_type
        self.length = length

    def __repr__(self):
        return '{}[{}]'.format(self.elem_type, '' if self.length is None else self.length)

class SimTypeString(SimTypeArray):
    '''
    SimTypeString is a type that represents a C-style string,
    i.e. a NUL-terminated array of bytes.
    '''

    _fields = SimTypeArray._fields + ('length',)

    def __init__(self, length=None, label=None):
        '''
        @param label: the type label
        @param length: an expression of the length of the string, if known
        '''
        SimTypeArray.__init__(self, SimTypeChar(), label=label)
        self.length = length

    def __repr__(self):
        return 'string_t'

    def extract(self, state, addr):
        mem = state.memory.load(addr, self.length)
        if state.se.symbolic(mem):
            return mem
        else:
            return repr(state.se.any_str(mem))

class SimTypeFunction(SimType):
    '''
    SimTypeFunction is a type that specifies an actual function (i.e. not a pointer) with certain types of arguments and a certain return value.
    '''

    _fields = ('args', 'returnty')
    base = False

    def __init__(self, args, returnty, label=None):
        '''
        @param label: the type label
        @param args: a tuple of types representing the arguments to the function
        @param returnty: the return type of the function
        '''
        SimType.__init__(self, label=label)
        self.args = args
        self.returnty = returnty

    def __repr__(self):
        return '({}) -> {}'.format(', '.join(str(a) for a in self.args), self.returnty)

class SimTypeLength(SimTypeInt):
    '''
    SimTypeLength is a type that specifies the length of some buffer in memory.
    '''

    _fields = SimTypeInt._fields + ('addr', 'length') # ?

    def __init__(self, arch, addr=None, length=None, label=None):
        '''
        @param label: the type label
        @param addr: the memory address (expression)
        @param length: the length (expression)
        '''
        SimTypeInt.__init__(self, arch.bits, False, label=label)
        self.addr = addr
        self.length = length
        self.signed = False

    def __repr__(self):
        return 'size_t'

class SimStructValue(object):
    def __init__(self, struct, values=None):
        self._struct = struct
        self._values = defaultdict(lambda: None, values or ())

    def __repr__(self):
        fields = ('.{} = {}'.format(name, self._values[name]) for name in self._struct.fields)
        return '{{\n  {}\n}}'.format(',\n  '.join(fields))

_C_TYPE_TO_SIMTYPE = {
    ('int',): lambda _: SimTypeInt(32, True),
    ('unsigned', 'int'): lambda _: SimTypeInt(32, False),
    ('long',): lambda arch: SimTypeInt(arch.bits, True),
    ('unsigned', 'long'): lambda arch: SimTypeInt(arch.bits, False),
    ('char',): lambda _: SimTypeChar(),
    ('int8_t',): lambda _: SimTypeInt(8, True),
    ('uint8_t',): lambda _: SimTypeInt(8, False),
    ('int16_t',): lambda _: SimTypeInt(16, True),
    ('uint16_t',): lambda _: SimTypeInt(16, False),
    ('int32_t',): lambda _: SimTypeInt(32, True),
    ('uint32_t',): lambda _: SimTypeInt(32, False),
    ('int64_t',): lambda _: SimTypeInt(64, True),
    ('uint64_t',): lambda _: SimTypeInt(64, False),
    ('ptrdiff_t',): lambda arch: SimTypeInt(arch.bits, False),
    ('size_t',): lambda arch: SimTypeInt(arch.bits, False),
    ('ssize_t',): lambda arch: SimTypeInt(arch.bits, True),
    ('uintptr_t',) : lambda arch: SimTypeInt(arch.bits, False),
}

def _decl_to_type(decl):
    if isinstance(decl, pycparser.c_ast.TypeDecl):
        return _C_TYPE_TO_SIMTYPE[tuple(decl.type.names)]
    elif isinstance(decl, pycparser.c_ast.PtrDecl):
        pts_to = _decl_to_type(decl.type)
        return lambda arch: SimTypePointer(arch, pts_to(arch))
    elif isinstance(decl, pycparser.c_ast.ArrayDecl):
        elem_type = _decl_to_type(decl.type)
        size = int(decl.dim.value)
        return lambda arch: SimTypeFixedSizeArray(elem_type(arch), size)

# these are all bogus, on purpose
_C_STRUCT_PREAMBLE = """
typedef int int8_t;
typedef int uint8_t;
typedef int int16_t;
typedef int uint16_t;
typedef int int32_t;
typedef int uint32_t;
typedef int int64_t;
typedef int uint64_t;
"""

class SimStruct(SimType):
    _fields = ('name', 'fields')

    def __init__(self, name, fields, pack=True):
        if not pack:
            raise ValueError("you think I've implemented padding, how cute")

        self._name = name
        self.fields = fields
        self._arch_offsets_cache = {}

    def _arch_offsets(self, arch):
        if arch in self._arch_offsets_cache:
            return self._arch_offsets_cache[arch]

        offsets = {}
        offset_so_far = 0
        for name, almost_ty in self.fields.iteritems():
            ty = almost_ty(arch)
            offsets[name] = (ty, offset_so_far)
            offset_so_far += ty.size / 8

        self._arch_offsets_cache[arch] = offsets
        return offsets

    def extract(self, state, addr):
        values = {name: ty.view(state, addr + offset)
                  for (name, (ty, offset))
                  in self._arch_offsets(state.arch).iteritems()}
        return SimStructValue(self, values=values)

    @property
    def name(self):
        return self._name

    @classmethod
    def from_c(cls, defn):
        if pycparser is None:
            raise ImportError("pycparser is needed to use SimStruct.from_c!")

        # if preprocess:
        #     defn = subprocess.Popen(['cpp'], stdin=subprocess.PIPE, stdout=subprocess.PIPE).communicate(input=defn)[0]

        node = pycparser.c_parser.CParser().parse(_C_STRUCT_PREAMBLE + defn)

        if not isinstance(node, pycparser.c_ast.FileAST) or \
           not isinstance(node.ext[-1], pycparser.c_ast.Decl) or \
           not isinstance(node.ext[-1].type, pycparser.c_ast.Struct):
            raise ValueError("invalid struct definition")

        struct = node.ext[-1].type
        fields = OrderedDict((decl.name, _decl_to_type(decl.type)) for decl in struct.decls)

        return cls(struct.name, fields)

    def _refine_dir(self):
        return self.fields.keys()

    def _refine(self, view, k):
        ty, offset = self._arch_offsets(view.state.arch)[k]
        return view._deeper(ty=ty, addr=view._addr + offset)

try:
    _example_struct = SimStruct.from_c("""
struct example {
  int foo;
  int bar;
  char *hello;
};
""")
except ImportError:
    _example_struct = None

ALL_TYPES = {
    'char': lambda _: SimTypeInt(8, True),
    'int8_t': lambda _: SimTypeInt(8, True),
    'uchar': lambda _: SimTypeInt(8, False),
    'uint8_t': lambda _: SimTypeInt(8, False),
    'byte': lambda _: SimTypeInt(8, False),

    'short': lambda _: SimTypeInt(16, True),
    'int16_t': lambda _: SimTypeInt(16, True),
    'ushort': lambda _: SimTypeInt(16, False),
    'uint16_t': lambda _: SimTypeInt(16, False),
    'word': lambda _: SimTypeInt(16, False),

    'int': lambda _: SimTypeInt(32, True),
    'int32_t': lambda _: SimTypeInt(32, True),
    'uint': lambda _: SimTypeInt(32, False),
    'uint32_t': lambda _: SimTypeInt(32, False),
    'dword': lambda _: SimTypeInt(32, False),

    'long': lambda _: SimTypeInt(64, True),
    'int64_t': lambda _: SimTypeInt(64, True),
    'ulong': lambda _: SimTypeInt(64, False),
    'uint64_t': lambda _: SimTypeInt(64, False),
    'qword': lambda _: SimTypeInt(64, False),

    'string': lambda arch: SimTypePointer(arch, SimTypeChar()),
    'example': lambda _: _example_struct,

}

def define_struct(defn):
    struct = SimStruct.from_c(defn)
    ALL_TYPES[struct.name] = lambda _: struct
    return struct

from .plugins.view import SimMemView
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
u"""Auth database

:copyright: Copyright (c) 2018-2019 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from __future__ import absolute_import, division, print_function
from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdc, pkdexc, pkdlog, pkdp
import contextlib
import sqlalchemy
import sqlalchemy.ext.declarative
import sqlalchemy.orm
# limit imports here
import sirepo.auth_role
import sirepo.srcontext
import sirepo.srdb
import sirepo.util

#: sqlite file located in sirepo_db_dir
_SQLITE3_BASENAME = 'auth.db'

_SRCONTEXT_SESSION_KEY = 'auth_db_session'

#: SQLAlchemy database engine
_engine = None

#: Keeps track of upgrades applied to the database
DbUpgrade = None

#: base for UserRegistration and *User models
UserDbBase = None

#: base for user models
UserRegistration = None

#: roles for each user
UserRole = None

def all_uids():
    return UserRegistration.search_all_for_column('uid')

def audit_proprietary_lib_files(uid, force=False, sim_types=None):
    """Add/removes proprietary files based on a user's roles

    For example, add the Flash tarball if user has the flash role.

    Args:
      uid (str): The uid of the user to audit
      force (bool): Overwrite existing lib files with the same name as new ones
      sim_types (set): Set of sim_types to audit (proprietary_sim_types if None)
    """
    import contextlib
    import py
    import pykern.pkconfig
    import pykern.pkio
    import sirepo.feature_config
    import sirepo.sim_data
    import sirepo.simulation_db
    import sirepo.util
    import subprocess

    def _add(proprietary_code_dir, sim_type, sim_data_class):
        p = proprietary_code_dir.join(sim_data_class.proprietary_code_tarball())
        sirepo.simulation_db.verify_app_directory(sim_type, uid=uid)
        with sirepo.simulation_db.tmp_dir(chdir=True, uid=uid) as t:
            d = t.join(p.basename)
            d.mksymlinkto(p, absolute=False)
            subprocess.check_output(
                [
                    'tar',
                    '--extract',
                    '--gunzip',
                    f'--file={d}',
                ],
                stderr=subprocess.STDOUT,
            )
            l = sirepo.simulation_db.simulation_lib_dir(sim_type, uid=uid)
            e = [f.basename for f in pykern.pkio.sorted_glob(l.join('*'))]
            for f in sim_data_class.proprietary_code_lib_file_basenames():
                if force or f not in e:
                    t.join(f).rename(l.join(f))

    s = sirepo.feature_config.cfg().proprietary_sim_types
    if sim_types:
        assert sim_types.issubset(s), \
            f'sim_types={sim_types} not a subset of proprietary_sim_types={s}'
        s = sim_types
    for t in s:
        c = sirepo.sim_data.get_class(t)
        if not c.proprietary_code_tarball():
            return
        d = sirepo.srdb.proprietary_code_dir(t)
        assert d.exists(), \
            f'{d} proprietary_code_dir must exist' \
            + ('; run: sirepo setup_dev' if pykern.pkconfig.channel_in('dev') else '')
        r = UserRole.has_role(
            uid,
            sirepo.auth_role.for_sim_type(t),
        )
        if r:
            _add(d, t, c)
            continue
        # SECURITY: User no longer has access so remove all artifacts
        pykern.pkio.unchecked_remove(sirepo.simulation_db.simulation_dir(t, uid=uid))

def db_filename():
    return sirepo.srdb.root().join(_SQLITE3_BASENAME)

def init():
    def _create_tables(engine):
        b = UserDbBase
        k = set(b.metadata.tables.keys())
        assert k.issubset(set(b.TABLES)), \
            f'sqlalchemy tables={k} not a subset of known tables={b.TABLES}'
        b.metadata.create_all(engine)

    global _engine, DbUpgrade, UserDbBase, UserRegistration, UserRole

    if _engine:
        return
    f = db_filename()
    _migrate_db_file(f)
    _engine = sqlalchemy.create_engine(
        'sqlite:///{}'.format(f),
        # We do our own thread locking so no need to have pysqlite warn us when
        # we access a single connection across threads
        connect_args={'check_same_thread': False},
    )

    @sqlalchemy.ext.declarative.as_declarative()
    class UserDbBase(object):
        STRING_ID = sqlalchemy.String(8)
        STRING_NAME =  sqlalchemy.String(100)
        TABLES = [
            # Order is important. SQLite doesn't allow for foreign key constraints to
            # be added after creation. So, the tables are ordered in the way
            # constraints should be carried out. For example, a user is deleted
            # from auth_email_user_t before user_registration_t.
            'auth_github_user_t',
            'auth_email_user_t',
            'jupyterhub_user_t',
            'user_role_t',
            'user_registration_t',
            'db_upgrade_t',
        ]

        def __init__(self, **kwargs):
            for k, v in kwargs.items():
                setattr(self, k, v)

        def delete(self):
            with sirepo.util.THREAD_LOCK:
                self._session().delete(self)
                self._session().commit()

        @classmethod
        def delete_user(cls, uid):
            """Delete user from all tables"""
            for t in cls.TABLES:
                m = cls._unchecked_model_from_tablename(t)
                # Exlicit None check because sqlalchemy overrides __bool__ to
                # raise TypeError
                if m is None or 'uid' not in m.columns:
                    continue
                cls.execute(sqlalchemy.delete(m).where(m.c.uid==uid))
            cls._session().commit()

        @classmethod
        def execute(cls, statement):
            cls._session().execute(
                statement.execution_options(synchronize_session='fetch')
            )

        @classmethod
        def delete_all(cls):
            with sirepo.util.THREAD_LOCK:
                cls._session().query(cls).delete()
                cls._session().commit()

        def save(self):
            with sirepo.util.THREAD_LOCK:
                self._session().add(self)
                self._session().commit()

        @classmethod
        def search_all_by(cls, **kwargs):
            with sirepo.util.THREAD_LOCK:
                return cls._session().query(cls).filter_by(**kwargs).all()

        @classmethod
        def search_by(cls, **kwargs):
            with sirepo.util.THREAD_LOCK:
                return cls._session().query(cls).filter_by(**kwargs).first()

        @classmethod
        def search_all_for_column(cls, column, **filter_by):
            with sirepo.util.THREAD_LOCK:
                return [
                    getattr(r, column) for r
                    in cls._session().query(cls).filter_by(**filter_by)
                ]

        @classmethod
        def delete_all_for_column_by_values(cls, column, values):
            with sirepo.util.THREAD_LOCK:
                cls.execute(sqlalchemy.delete(cls).where(
                    getattr(cls, column).in_(values),
                ))
                cls._session().commit()

        @classmethod
        def _session(cls):
            return sirepo.srcontext.get(_SRCONTEXT_SESSION_KEY)

        @classmethod
        def _unchecked_model_from_tablename(cls, tablename):
            for k, v in cls.metadata.tables.items():
                if k == tablename:
                    return v

    class DbUpgrade(UserDbBase):
        __tablename__ = 'db_upgrade_t'
        name = sqlalchemy.Column(UserDbBase.STRING_NAME, primary_key=True)
        created = sqlalchemy.Column(sqlalchemy.DateTime(), nullable=False)

    class UserRegistration(UserDbBase):
        __tablename__ = 'user_registration_t'
        uid = sqlalchemy.Column(UserDbBase.STRING_ID, primary_key=True)
        created = sqlalchemy.Column(sqlalchemy.DateTime(), nullable=False)
        display_name = sqlalchemy.Column(UserDbBase.STRING_NAME)

    class UserRole(UserDbBase):
        __tablename__ = 'user_role_t'
        uid = sqlalchemy.Column(UserDbBase.STRING_ID, primary_key=True)
        role = sqlalchemy.Column(UserDbBase.STRING_NAME, primary_key=True)

        @classmethod
        def all_roles(cls):
            with sirepo.util.THREAD_LOCK:
                return [
                    r[0] for r in cls._session().query(cls.role.distinct()).all()
                ]

        @classmethod
        def add_roles(cls, uid, roles):
            with sirepo.util.THREAD_LOCK:
                for r in roles:
                    UserRole(uid=uid, role=r).save()
                audit_proprietary_lib_files(uid)

        @classmethod
        def delete_roles(cls, uid, roles):
            with sirepo.util.THREAD_LOCK:
                cls.execute(sqlalchemy.delete(cls).where(
                    cls.uid == uid,
                ).where(
                    cls.role.in_(roles),
                ))
                cls._session().commit()
                audit_proprietary_lib_files(uid)

        @classmethod
        def get_roles(cls, uid):
            with sirepo.util.THREAD_LOCK:
                return UserRole.search_all_for_column(
                    'role',
                    uid=uid,
                )

        @classmethod
        def has_role(cls, uid, role):
            with sirepo.util.THREAD_LOCK:
                return bool(cls.search_by(uid=uid, role=role))

        @classmethod
        def uids_of_paid_users(cls):
            return [
                x[0] for x in cls._session().query(cls).with_entities(cls.uid).filter(
                    cls.role.in_(sirepo.auth_role.PAID_USER_ROLES),
                ).distinct().all()
            ]
    _create_tables(_engine)

def init_model(callback):
    with sirepo.util.THREAD_LOCK:
        callback(UserDbBase)
        UserDbBase.metadata.create_all(_engine)

@contextlib.contextmanager
def session():
    init()
    with sirepo.srcontext.create() as c:
        try:
            _create_session(c)
            yield
        finally:
            _destroy_session(c)

@contextlib.contextmanager
def session_and_lock():
    # TODO(e-carlin): Need locking across processes
    # git.radiasoft.org/sirepo/issues/3516
    with session():
        yield

def _create_session(context):
    s = context.get(_SRCONTEXT_SESSION_KEY)
    assert not s, \
        f'existing session={s}'
    context[_SRCONTEXT_SESSION_KEY] = sqlalchemy.orm.Session(bind=_engine)

def _destroy_session(context):
    context.pop(_SRCONTEXT_SESSION_KEY).rollback()

def _migrate_db_file(fn):
    o = fn.new(basename='user.db')
    if not o.exists():
        return
    assert not fn.exists(), \
        'db file collision: old={} and new={} both exist'.format(o, fn)
    try:
        # reduce the race condition between job_supervisor and sirepo starting
        x = o + '-migrating'
        if x.exists():
            # again, reduce race condition
            return
        o.rename(x)

        import sqlite3

        c = sqlite3.connect(str(x))
        c.row_factory = sqlite3.Row
        rows = c.execute('SELECT * FROM user_t')
        c2 = sqlite3.connect(str(fn))
        # sqlite3 saves the literal string as the schema
        # so we are formatting it just like SQLAlchemy would
        # format it.
        c2.execute(
            '''CREATE TABLE auth_github_user_t (
        oauth_id VARCHAR(100) NOT NULL,
        user_name VARCHAR(100) NOT NULL,
        uid VARCHAR(8),
        PRIMARY KEY (oauth_id),
        UNIQUE (user_name),
        UNIQUE (uid)
);'''
        )
        for r in rows:
            c2.execute(
                '''
                INSERT INTO auth_github_user_t
                (oauth_id, user_name, uid)
                VALUES (?, ?, ?)
                ''',
                (r['oauth_id'], r['user_name'], r['uid']),
            )
        c.close()
        c2.commit()
        c2.close()
    except sqlite3.OperationalError as e:
        if 'not such table' in e.message:
            return
        raise
    x.rename(o + '-migrated')
    pkdlog('migrated user.db to auth.db')
<commit_msg>WIP radiasoft/ops#645: ENOENT when trying to rename flash.tar.gz<commit_after># -*- coding: utf-8 -*-
u"""Auth database

:copyright: Copyright (c) 2018-2019 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from __future__ import absolute_import, division, print_function
from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdc, pkdexc, pkdlog, pkdp
import contextlib
import sqlalchemy
import sqlalchemy.ext.declarative
import sqlalchemy.orm
# limit imports here
import sirepo.auth_role
import sirepo.srcontext
import sirepo.srdb
import sirepo.util

#: sqlite file located in sirepo_db_dir
_SQLITE3_BASENAME = 'auth.db'

_SRCONTEXT_SESSION_KEY = 'auth_db_session'

#: SQLAlchemy database engine
_engine = None

#: Keeps track of upgrades applied to the database
DbUpgrade = None

#: base for UserRegistration and *User models
UserDbBase = None

#: base for user models
UserRegistration = None

#: roles for each user
UserRole = None

def all_uids():
    return UserRegistration.search_all_for_column('uid')

def audit_proprietary_lib_files(uid, force=False, sim_types=None):
    """Add/removes proprietary files based on a user's roles

    For example, add the Flash tarball if user has the flash role.

    Args:
      uid (str): The uid of the user to audit
      force (bool): Overwrite existing lib files with the same name as new ones
      sim_types (set): Set of sim_types to audit (proprietary_sim_types if None)
    """
    import pykern.pkconfig
    import pykern.pkio
    import sirepo.feature_config
    import sirepo.sim_data
    import sirepo.simulation_db
    import sirepo.util
    import subprocess

    def _add(proprietary_code_dir, sim_type, sim_data_class):
        p = proprietary_code_dir.join(sim_data_class.proprietary_code_tarball())
        sirepo.simulation_db.verify_app_directory(sim_type, uid=uid)
        with sirepo.simulation_db.tmp_dir(chdir=True, uid=uid) as t:
            d = t.join(p.basename)
            d.mksymlinkto(p, absolute=False)
            subprocess.check_output(
                [
                    'tar',
                    '--extract',
                    '--gunzip',
                    f'--file={d}',
                ],
                stderr=subprocess.STDOUT,
            )
            l = sirepo.simulation_db.simulation_lib_dir(sim_type, uid=uid)
            # See git.radiasoft.org/ops/issues/645
            assert l.check(), \
                'app dir for sim_type={sim_type} but no lib_dir={l}'
            e = [f.basename for f in pykern.pkio.sorted_glob(l.join('*'))]
            for f in sim_data_class.proprietary_code_lib_file_basenames():
                if force or f not in e:
                    t.join(f).rename(l.join(f))

    s = sirepo.feature_config.cfg().proprietary_sim_types
    if sim_types:
        assert sim_types.issubset(s), \
            f'sim_types={sim_types} not a subset of proprietary_sim_types={s}'
        s = sim_types
    for t in s:
        c = sirepo.sim_data.get_class(t)
        if not c.proprietary_code_tarball():
            return
        d = sirepo.srdb.proprietary_code_dir(t)
        assert d.exists(), \
            f'{d} proprietary_code_dir must exist' \
            + ('; run: sirepo setup_dev' if pykern.pkconfig.channel_in('dev') else '')
        r = UserRole.has_role(
            uid,
            sirepo.auth_role.for_sim_type(t),
        )
        if r:
            _add(d, t, c)
            continue
        # SECURITY: User no longer has access so remove all artifacts
        pykern.pkio.unchecked_remove(sirepo.simulation_db.simulation_dir(t, uid=uid))

def db_filename():
    return sirepo.srdb.root().join(_SQLITE3_BASENAME)

def init():
    def _create_tables(engine):
        b = UserDbBase
        k = set(b.metadata.tables.keys())
        assert k.issubset(set(b.TABLES)), \
            f'sqlalchemy tables={k} not a subset of known tables={b.TABLES}'
        b.metadata.create_all(engine)

    global _engine, DbUpgrade, UserDbBase, UserRegistration, UserRole

    if _engine:
        return
    f = db_filename()
    _migrate_db_file(f)
    _engine = sqlalchemy.create_engine(
        'sqlite:///{}'.format(f),
        # We do our own thread locking so no need to have pysqlite warn us when
        # we access a single connection across threads
        connect_args={'check_same_thread': False},
    )

    @sqlalchemy.ext.declarative.as_declarative()
    class UserDbBase(object):
        STRING_ID = sqlalchemy.String(8)
        STRING_NAME =  sqlalchemy.String(100)
        TABLES = [
            # Order is important. SQLite doesn't allow for foreign key constraints to
            # be added after creation. So, the tables are ordered in the way
            # constraints should be carried out. For example, a user is deleted
            # from auth_email_user_t before user_registration_t.
            'auth_github_user_t',
            'auth_email_user_t',
            'jupyterhub_user_t',
            'user_role_t',
            'user_registration_t',
            'db_upgrade_t',
        ]

        def __init__(self, **kwargs):
            for k, v in kwargs.items():
                setattr(self, k, v)

        def delete(self):
            with sirepo.util.THREAD_LOCK:
                self._session().delete(self)
                self._session().commit()

        @classmethod
        def delete_user(cls, uid):
            """Delete user from all tables"""
            for t in cls.TABLES:
                m = cls._unchecked_model_from_tablename(t)
                # Exlicit None check because sqlalchemy overrides __bool__ to
                # raise TypeError
                if m is None or 'uid' not in m.columns:
                    continue
                cls.execute(sqlalchemy.delete(m).where(m.c.uid==uid))
            cls._session().commit()

        @classmethod
        def execute(cls, statement):
            cls._session().execute(
                statement.execution_options(synchronize_session='fetch')
            )

        @classmethod
        def delete_all(cls):
            with sirepo.util.THREAD_LOCK:
                cls._session().query(cls).delete()
                cls._session().commit()

        def save(self):
            with sirepo.util.THREAD_LOCK:
                self._session().add(self)
                self._session().commit()

        @classmethod
        def search_all_by(cls, **kwargs):
            with sirepo.util.THREAD_LOCK:
                return cls._session().query(cls).filter_by(**kwargs).all()

        @classmethod
        def search_by(cls, **kwargs):
            with sirepo.util.THREAD_LOCK:
                return cls._session().query(cls).filter_by(**kwargs).first()

        @classmethod
        def search_all_for_column(cls, column, **filter_by):
            with sirepo.util.THREAD_LOCK:
                return [
                    getattr(r, column) for r
                    in cls._session().query(cls).filter_by(**filter_by)
                ]

        @classmethod
        def delete_all_for_column_by_values(cls, column, values):
            with sirepo.util.THREAD_LOCK:
                cls.execute(sqlalchemy.delete(cls).where(
                    getattr(cls, column).in_(values),
                ))
                cls._session().commit()

        @classmethod
        def _session(cls):
            return sirepo.srcontext.get(_SRCONTEXT_SESSION_KEY)

        @classmethod
        def _unchecked_model_from_tablename(cls, tablename):
            for k, v in cls.metadata.tables.items():
                if k == tablename:
                    return v

    class DbUpgrade(UserDbBase):
        __tablename__ = 'db_upgrade_t'
        name = sqlalchemy.Column(UserDbBase.STRING_NAME, primary_key=True)
        created = sqlalchemy.Column(sqlalchemy.DateTime(), nullable=False)

    class UserRegistration(UserDbBase):
        __tablename__ = 'user_registration_t'
        uid = sqlalchemy.Column(UserDbBase.STRING_ID, primary_key=True)
        created = sqlalchemy.Column(sqlalchemy.DateTime(), nullable=False)
        display_name = sqlalchemy.Column(UserDbBase.STRING_NAME)

    class UserRole(UserDbBase):
        __tablename__ = 'user_role_t'
        uid = sqlalchemy.Column(UserDbBase.STRING_ID, primary_key=True)
        role = sqlalchemy.Column(UserDbBase.STRING_NAME, primary_key=True)

        @classmethod
        def all_roles(cls):
            with sirepo.util.THREAD_LOCK:
                return [
                    r[0] for r in cls._session().query(cls.role.distinct()).all()
                ]

        @classmethod
        def add_roles(cls, uid, roles):
            with sirepo.util.THREAD_LOCK:
                for r in roles:
                    UserRole(uid=uid, role=r).save()
                audit_proprietary_lib_files(uid)

        @classmethod
        def delete_roles(cls, uid, roles):
            with sirepo.util.THREAD_LOCK:
                cls.execute(sqlalchemy.delete(cls).where(
                    cls.uid == uid,
                ).where(
                    cls.role.in_(roles),
                ))
                cls._session().commit()
                audit_proprietary_lib_files(uid)

        @classmethod
        def get_roles(cls, uid):
            with sirepo.util.THREAD_LOCK:
                return UserRole.search_all_for_column(
                    'role',
                    uid=uid,
                )

        @classmethod
        def has_role(cls, uid, role):
            with sirepo.util.THREAD_LOCK:
                return bool(cls.search_by(uid=uid, role=role))

        @classmethod
        def uids_of_paid_users(cls):
            return [
                x[0] for x in cls._session().query(cls).with_entities(cls.uid).filter(
                    cls.role.in_(sirepo.auth_role.PAID_USER_ROLES),
                ).distinct().all()
            ]
    _create_tables(_engine)

def init_model(callback):
    with sirepo.util.THREAD_LOCK:
        callback(UserDbBase)
        UserDbBase.metadata.create_all(_engine)

@contextlib.contextmanager
def session():
    init()
    with sirepo.srcontext.create() as c:
        try:
            _create_session(c)
            yield
        finally:
            _destroy_session(c)

@contextlib.contextmanager
def session_and_lock():
    # TODO(e-carlin): Need locking across processes
    # git.radiasoft.org/sirepo/issues/3516
    with session():
        yield

def _create_session(context):
    s = context.get(_SRCONTEXT_SESSION_KEY)
    assert not s, \
        f'existing session={s}'
    context[_SRCONTEXT_SESSION_KEY] = sqlalchemy.orm.Session(bind=_engine)

def _destroy_session(context):
    context.pop(_SRCONTEXT_SESSION_KEY).rollback()

def _migrate_db_file(fn):
    o = fn.new(basename='user.db')
    if not o.exists():
        return
    assert not fn.exists(), \
        'db file collision: old={} and new={} both exist'.format(o, fn)
    try:
        # reduce the race condition between job_supervisor and sirepo starting
        x = o + '-migrating'
        if x.exists():
            # again, reduce race condition
            return
        o.rename(x)

        import sqlite3

        c = sqlite3.connect(str(x))
        c.row_factory = sqlite3.Row
        rows = c.execute('SELECT * FROM user_t')
        c2 = sqlite3.connect(str(fn))
        # sqlite3 saves the literal string as the schema
        # so we are formatting it just like SQLAlchemy would
        # format it.
        c2.execute(
            '''CREATE TABLE auth_github_user_t (
        oauth_id VARCHAR(100) NOT NULL,
        user_name VARCHAR(100) NOT NULL,
        uid VARCHAR(8),
        PRIMARY KEY (oauth_id),
        UNIQUE (user_name),
        UNIQUE (uid)
);'''
        )
        for r in rows:
            c2.execute(
                '''
                INSERT INTO auth_github_user_t
                (oauth_id, user_name, uid)
                VALUES (?, ?, ?)
                ''',
                (r['oauth_id'], r['user_name'], r['uid']),
            )
        c.close()
        c2.commit()
        c2.close()
    except sqlite3.OperationalError as e:
        if 'not such table' in e.message:
            return
        raise
    x.rename(o + '-migrated')
    pkdlog('migrated user.db to auth.db')
<|endoftext|>