#!/usr/bin/env python
#
# This file is part of SALI
#
# SALI is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# SALI is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with SALI.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright 2010-2013 SURFsara

import sys
import argparse
import os
import re
import subprocess
import tempfile
import atexit
import datetime
    
try:
    from sali import general
except ImportError:
    sys.path.append( '../' )

    from sali import general

class ExitException( Exception ):
    pass

def ask( question, choices=list(), default=None, showlist=False ):

    choices.sort()

    if showlist or len( choices ) > 3:
        for item in choices:
            print( '%2d %s' % ( choices.index( item ) + 1, item ) )
        str_question = '%s: ' % question
    else:
        str_question = '%s: %s' % ( question, ' '.join( choices ) )

    if len( choices ) > 3 and not showlist:
        e_default = '? for options'
        str_question = '%s: ' % question
    else:
        e_default = ''

    if default:
        str_default = default
    else:
        str_default = ''

    answer = raw_input( '%s [%s] ' % ( str_question, str_default ) )

    if answer == '?':
        if len( choices ) > 3:
            return ask( question, choices, default, showlist=True )
        else:
            return ask( question, choices, default, showlist=False )
    elif not answer and default:
        return default
    elif showlist and answer.isdigit() or len( choices ) > 3 and answer.isdigit():
        return choices[ int( answer ) - 1 ]
    elif not answer:
        return ask( question, choices, default, showlist=False )
        
    return answer

class ServerClient( object ):

    def __init__( self, args ): 
        self.cfg = general.Config( args.config ) 

    def run( self ):
        pass

class GetImage( ServerClient ):

    exclude_tmp = None

    def __init__( self, args ):
        ServerClient.__init__( self, args )

        p = subprocess.Popen( [ 'which', 'rsync' ], stdout=subprocess.PIPE, stderr=subprocess.PIPE )
        std, ste = p.communicate()

        if p.returncode > 0:
            raise ExitException( 'Could not located rsync' )

        self.rsync = std.strip()

        if not args.host:
            self.hostname = ask( 'Golden image hostname' )
        else:
            self.hostname = args.host

        if not args.image:
            self.imagename = ask( 'Imagename', choices=os.listdir( self.cfg.get( 'update', 'imagedir' ) ) )
        else:
            self.imagename = args.image

        if args.exclude:
            self.excludefile = args.exclude
        else:
            self.excludefile = None

        self.no_tarball = os.path.join( self.cfg.get_default( 'maindir' ), 'notorrent' )

        if not os.path.exists( self.no_tarball ):
            os.makedirs( self.no_tarball, 0777 )

        fo = open( os.path.join( self.no_tarball, self.imagename ), 'w' )
        fo.write( "Busy" )
        fo.close()

        atexit.register( self.exit )

    def get_image( self, excludefile ):

        if self.cfg.getboolean( 'rsync', 'quiet' ):
            opts = '-aHS'
            quiet = True
        else:
            opts = '-aHSv'
            quiet = False

        command = [ 
            self.rsync, 
            opts,
            '--numeric-ids',
            '--delete',
            '--delete-excluded',
            '--exclude-from=%s' % excludefile,
            'rsync://%s:%d/root/' % ( 
                self.hostname,
                self.cfg.getint( 'rsync', 'port' ) ) ,
            os.path.join( 
                self.cfg.get( 'update', 'imagedir' ), 
                self.imagename
            ),
        ]

        p = subprocess.Popen( command , stdout=subprocess.PIPE, stderr=subprocess.PIPE )

        if not quiet:
            print( '\n### Fetching image from host %s\n' % self.hostname )
            for line in p.stdout:
                print( line.strip() )

        std, ste = p.communicate()

        if p.returncode > 0:
            print( ste )
            raise ExitException( 'An error has occured during rsync transfer' )

        if not quiet:
            print( '\n### Done with imaging' )

        ccommand = [
            self.rsync,
            'rsync://%s:%d/__cloning_completed__' % ( 
                self.hostname,
                self.cfg.getint( 'rsync', 'port' ) ) ,
            '/',
        ]
        
        p = subprocess.Popen( ccommand, stdout=subprocess.PIPE, stderr=subprocess.PIPE )
        std, ste = p.communicate()

        if p.returncode <> 5 and p.returncode > 0:
            print( ste )
            raise ExitException( 'Could not send terminate signal to golden client' )

    def exclude_file( self ):
        fd, fname = tempfile.mkstemp()
        fto = os.fdopen( fd, 'w' )

        mfs_file = os.path.join( self.cfg.get( 'update', 'imagedir' ), '%s/etc/systemimager/mounted_filesystems' % self.imagename )

        mfo = open( mfs_file )
        fto.write( '## Rules from the mounted_filesystems file\n' )
        for line in mfo.readlines():
            splitted = line.strip().split( ' ' )
            if splitted[ 4 ] not in [ 'ext2','ext3','ext4','reiserfs','jfs','xfs','vfat','fat' ]:
                fto.write( '%s/*\n' % splitted[2] )
        mfo.close()

        fto.write( '## Rules from the getimage.exclude\n' )
        gex = open( '/etc/sali/getimage.exclude' )
        fto.write( gex.read() )
        gex.close()

        if self.excludefile:
            fto.write( '## External rules\n' )
            exf = open( self.excludefile )
            fto.write( exf.read() )
            exf.close()

        fto.close()
        return fname

    def get_mountedfilesystem( self ):

        if self.cfg.getboolean( 'rsync', 'quiet' ):
            opts = '-a'
            quiet = True
        else:
            opts = '-av'
            quiet = False

        command = [ 
            self.rsync, 
            opts, 
            'rsync://%s:%d/root/etc/systemimager/mounted_filesystems' % ( 
                self.hostname,
                self.cfg.getint( 'rsync', 'port' ) ) ,
            os.path.join( 
                self.cfg.get( 'update', 'imagedir' ), 
                '%s/etc/systemimager/mounted_filesystems' % self.imagename
            ),
        ]

        p = subprocess.Popen( command , stdout=subprocess.PIPE, stderr=subprocess.PIPE )

        if not quiet:
            print( '\n### Fetching mounted_filesystems from host %s\n' % self.hostname )
            for line in p.stdout:
                print( line.strip() )

        std, ste = p.communicate()

        if p.returncode > 0:
            print( ste )
            raise ExitException( 'An error has occured during rsync transfer' )

        if not quiet:
            print( '\n### Done' )
        
    def run( self ):

        imagepath = os.path.join( self.cfg.get( 'update', 'imagedir' ), self.imagename )

        if not os.path.exists( imagepath ):
            os.makedirs( imagepath, 0777 )
            print( 'A new imagedirectory \'%s\' has been created in \'%s\'' % (
                self.imagename,
                self.cfg.get( 'update', 'imagedir' ) )
            )

        if not os.path.exists( os.path.join( imagepath, 'etc/systemimager' ) ):
            os.makedirs( os.path.join( imagepath, 'etc/systemimager' ), 0777 )
            print( 'A new directory for systemimager has been created \'%s\'' % (
                os.path.join( imagepath, 'etc/systemimager' ) )
            )

        self.get_mountedfilesystem()
        self.exclude_tmp = self.exclude_file()
        self.get_image( self.exclude_tmp )
        os.utime( imagepath, None )

        self.exit()

    def exit( self ):

        if self.exclude_tmp and os.path.exists( self.exclude_tmp ):
            os.remove( self.exclude_tmp )

        if os.path.exists( os.path.join( self.no_tarball, self.imagename ) ):
            os.remove( os.path.join( self.no_tarball, self.imagename ) )

class Rsync( ServerClient ):

    def run( self ):
        ## if we can't find the location, STOP
        if not os.path.isdir( self.cfg.get( 'rsync', 'stubdir' ) ):
            raise ExitException( 'Cannot find rsync stubdir' )

        if '00header.conf' not in os.listdir( self.cfg.get( 'rsync', 'stubdir' ) ):
            raise ExitException( '00header.conf must be available!' )

        fo = open( self.cfg.get( 'rsync', 'configfile' ), 'w' )
        fo.write( '### Do not edit this file, check the documentation of SALI ###\n' )
        str_date = datetime.datetime.strftime( datetime.datetime.now(), '%m/%d/%Y %H:%M' )
        fo.write( '###                    %s                    ###\n\n' % str_date )
        for file in os.listdir( self.cfg.get( 'rsync', 'stubdir' ) ):
            fi = open( os.path.join( self.cfg.get( 'rsync', 'stubdir' ), file ) )
            fo.write( '### Begin %s\n' % file )
            fo.write( fi.read() )
            fo.write( '### End %s\n\n' % file )
            fi.close()

        fo.write( '### Images from %s\n' % self.cfg.get( 'update', 'imagedir' ) )
        for image in os.listdir( self.cfg.get( 'update', 'imagedir' ) ):
            fo.write( '[%s]\n' % image )
            fo.write( '   path=%s\n\n' % os.path.join( self.cfg.get( 'update', 'imagedir' ), image ) )
        fo.close()

OPTIONS = {
    'rsync': Rsync,
    'getimage' : GetImage,
}

if __name__ == '__main__':

    usage = '''%(prog)s <getimage|rsync> [options]

getimage:
  Fetch an image from a goldenclient, see --help for the options

rsync:
  Updates the rsynd.conf file by reading the stub files and images

options:
  Use --help
'''

    parser = argparse.ArgumentParser( 
        description='The server client tool for fetching images and creating the rsynd file',
        conflict_handler='resolve',
        usage=usage,
    )

    parser.add_argument( 
        'option', 
        metavar='OPTION', 
        type=str, 
        choices=OPTIONS.keys(),
        help='Specify which operation must be done' 
    )

    parser.add_argument(
        '-c', '--config',
        metavar='CONFIG',
        type=str,
        default='/etc/sali/sali.cfg',
        help='Configuration file of sali'
    )

    parser.add_argument(
        '-h', '--host',
        metavar='HOSTNAME',
        type=str,
        help='Specify which host that must be imaged',
    )

    parser.add_argument(
        '-i', '--image',
        metavar='IMAGENAME',
        type=str,
        help='The name of the image, if the image already exists the image will be updated',
    )

    parser.add_argument(
        '-e', '--exclude',
        metavar='EXCLUDE',
        type=str,
        help='Specify the exlude file (rsync style)',
    )

    args = parser.parse_args()

    try:
        tr = OPTIONS[ args.option ]( args )
        tr.run()
    except ExitException, err:
        print( err )
        sys.exit( 1 )
    except KeyboardInterrupt:
        print( '\n\nCaught a keyboardinterrupt' )
        sys.exit( 0 )

