#!/usr/bin/env python
# -*- mode: python; tab-width: 4 -*-
# $Id: smbcp,v 1.8 2002/08/03 05:45:32 miketeo Exp $
#
# Copyright (C) 2001 Michael Teo <michaelteo@bigfoot.com>
# smbcp - file transfer program between local machine and remote SMB machines
#
# This software is provided 'as-is', without any express or implied warranty. 
# In no event will the author be held liable for any damages arising from the 
# use of this software.
#
# Permission is granted to anyone to use this software for any purpose, 
# including commercial applications, and to alter it and redistribute it 
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not 
#    claim that you wrote the original software. If you use this software 
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
#
# 2. Altered source versions must be plainly marked as such, and must not be 
#    misrepresented as being the original software.
#
# 3. This notice cannot be removed or altered from any source distribution.
#

import os, os.path, sys, stat, string, re, getopt, getpass
import smb, nmb

VERSION = '0.1.4'

USER_HOST_PATH_PATTERN = re.compile('(\w+)@([\w\.]+):/*([^/]+)(/.*)')
HOST_PATH_PATTERN = re.compile(r"([\w\.]+):/*([^/]+)(/*.*)")

def parse_path(path):
    m = USER_HOST_PATH_PATTERN.match(path)
    if m:
        return m.groups()
    m = HOST_PATH_PATTERN.match(path)
    if m:
        return '', m.group(1), m.group(2), m.group(3)
    return None



def print_usage_screen():
    print 'Usage: smbcp [options] src ... [<user>@]<nbname>:/share/path'
    print '       smbcp [options] [<user>@]<nbname>:/share/path dest'
    print
    print 'Options:'
    print '  -n                    show what would have been transferred without doing it'
    print '  -r                    recurse into directories'
    print '  -B broadcast addr     the address to use for broadcasts'
    print '  -I remote IP addr     connect to this IP address for remote host'
    print '  --version             show version and exit'
    print
    


def remote2local_copy(src, src_service, src_path, dest_path, is_dest_dir, password):
    global recursive, dry_run

    try:
        info = src.list_path(src_service, src_path, password = password)
        is_src_dir = info[0].is_directory()
    except smb.SessionError, ex:
        print ex[0], ':', smb.strerror(ex[1], ex[2])[1]
        return

    if is_src_dir:
        if not is_dest_dir:
            print 'Warning: Destination is not directory'
            return

        if recursive:
            if not dry_run:
                try:
                    os.mkdir(os.path.join(dest_path, os.path.basename(src_path)))
                except:
                    pass
        
            for f in src.list_path(src_service, src_path + '/*', password = password):
                name = f.get_longname()
                if name == '.' or name == '..':
                    continue

                if f.is_directory():
                    remote2local_copy(src, src_service, src_path + '/' + name, os.path.join(dest_path, os.path.basename(src_path)), is_dest_dir, password)
                else:
                    print src_path + '/' + name,
                    print '=>', os.path.join(dest_path, os.path.basename(src_path), name)
                    if not dry_run:
                        fh = open(os.path.join(dest_path, os.path.basename(src_path), name), 'wb')
                        try:
                            src.retr_file(src_service, src_path + '/' + name, fh.write, password = password)
                        finally:
                            fh.close()
                        
        else:
            print 'Warning: source is a directory but -r is not set.'
    else:
        if is_dest_dir:
            print src_path, '=>', os.path.join(dest_path, os.path.basename(src_path))
        else:
            print src_path, '=>', dest_path
            
        if not dry_run:
            if is_dest_dir:
                fh = open(os.path.join(dest_path, os.path.basename(src_path)), 'wb')
            else:
                fh = open(dest_path, 'wb')
            try:
                src.retr_file(src_service, src_path, fh.write, password = password)
            finally:
                fh.close()
                
    

def local2remote_copy(remote, dest_service, dest_path, src_path, is_dest_dir, password):
    global recursive, dry_run

    if os.path.isdir(src_path):
        if not is_dest_dir:
            print 'Warning: Destination is not directory'
            return
        
        if recursive:
            for f in os.listdir(src_path):
                if f != '.' and f != '..':
                    try:
                        remote.mkdir(dest_service, dest_path + '/' + os.path.basename(src_path), password)
                    except smb.SessionError:
                        pass
                    local2remote_copy(remote, dest_service, dest_path + '/' + os.path.basename(src_path), os.path.join(src_path, f), is_dest_dir, password)
        else:
            print 'Warning:', src_path, 'is a directory but -r is not set.'
    else:
        print src_path, '=>',
        if is_dest_dir:
            print dest_path + '/' + os.path.basename(src_path)
        else:
            print dest_path
            
        if not os.path.exists(src_path):
            print 'Warning:', src_path, 'does not exist'
        elif not dry_run:
            fh = open(src_path, 'rb')
            try:
                if is_dest_dir:
                    remote.stor_file(dest_service, dest_path + '/' + os.path.basename(src_path), fh.read, password = password)
                else:
                    remote.stor_file(dest_service, dest_path, fh.read, password = password)
            finally:
                fh.close()



recursive = 0
dry_run = 0

def main():
    global recursive, dry_run

    remote_ip_addr = None
    netbios = nmb.NetBIOS()
    try:
        optlist, args = getopt.getopt(sys.argv[1:], 'rnB:I:', [ 'version' ])
        for k, v in optlist:
            if k == '-r':
                recursive = 1
            elif k == '-n':
                dry_run = 1
            elif k == '-B':
                netbios.set_broadcastaddr(v)
            elif k == '-I':
                remote_ip_addr = v
            elif k == '--version':
                print 'smbcp version', VERSION, '(smb:', smb.CVS_REVISION, 'and nmb:', nmb.CVS_REVISION, ')'
                return
    except getopt.error:
        print_usage_screen()
        return
    
    if len(args) < 2:
        print_usage_screen()
        return

    dest = args[-1]
    sources = args[:-1]
    
    dest_info = parse_path(dest)
    if dest_info:
        user, dest_name, dest_service, dest_path = dest_info

        if not remote_ip_addr:
            try:
                addrs = netbios.gethostbyname(dest_name)
                if not addrs:
                    print 'Warning: No such host,', dest_name
                    return
            except nmb.NetBIOSTimeout:
                print 'Warning: Cannot resolve', dest_name
                return
            except nmb.NetBIOSError, ex:
                print 'Warning: Error resolving', dest_name
                print ex[0], ':', nmb.strerror(ex[1], ex[2])[1]
                return

            print 'Connecting to remote,', dest_name, 'using', addrs[0].get_ip()
            remote = smb.SMB(dest_name, addrs[0].get_ip())
        else:
            print 'Connecting to remote,', dest_name, 'using', remote_ip_addr
            remote = smb.SMB(dest_name, remote_ip_addr)
            
        if remote.is_login_required():
            if not user:
                user = getpass.getuser()
            password = getpass.getpass('Password for ' + user + '@' + dest_name + ': ')
            remote.login(user, password)
            password = None
        else:
            password = getpass.getpass('Password for ' + dest_service + ': ')
                
        is_dest_dir = 0
        try:
            remote.check_dir(dest_service, dest_path, password)
            is_dest_dir = 1
        except smb.SessionError:
            pass

        if len(sources) > 1 and not is_dest_dir:
            print 'Warning: Destination not directory'
            return

        if dest_path and dest_path[-1] == '/':
            dest_path = dest_path[:-1]

        for source in sources:
            if parse_path(source):
                print 'Warning: You have specified a remote=>remote copy. We can only do local=>remote or remote=>local copy.'
            else:
                source = os.path.normpath(source)
                local2remote_copy(remote, dest_service, dest_path, source, is_dest_dir, password)
    else:
        is_dest_dir = os.path.isdir(dest)
        if len(sources) > 1 and not is_dest_dir:
            print 'Warning: Destination not directory'
            return

        passwords = { }
        for source in sources:
            src_info = parse_path(source)
            if not src_info:
                print 'Warning: You have specified a local=>local copy. We can only do local=>remote or remote=>local copy.'

            user, src_name, src_service, src_path = src_info

            if not remote_ip_addr:
                try:
                    addrs = netbios.gethostbyname(src_name)
                    if not addrs:
                        print 'Warning: No such host,', dest_name
                        return
                except nmb.NetBIOSTimeout:
                    print 'Warning: Cannot resolve', src_name
                    return
                except nmb.NetBIOSError, ex:
                    print 'Warning: Error resolving', src_name
                    print ex[0], ':', nmb.strerror(ex[1], ex[2])[1]
                    return

                print 'Connecting to remote,', src_name, 'using', addrs[0].get_ip()
                src = smb.SMB(src_name, addrs[0].get_ip())
            else:
                print 'Connecting to remote,', src_name, 'using', remote_ip_addr
                src = smb.SMB(src_name, remote_ip_addr)

            if src.is_login_required():
                if not user:
                    user = getpass.getuser()
                if passwords.has_key(( user, src_name )):
                    password = passwords[( user, src_name )]
                else:
                    password = getpass.getpass('Password for ' + user + '@' + src_name + ': ')
                    passwords[( user, src_name )] = password
                src.login(user, password)
                password = None
            else:
                password = getpass.getpass('Password for ' + src_service + ': ')
                 
            if src_path and src_path[-1] == '/':
                src_path = src_path[:-1]

            remote2local_copy(src, src_service, src_path, dest, is_dest_dir, password)
            
    

if __name__ == '__main__':
    try:
        main()
    except nmb.NetBIOSError, ex:
        _, err_msg = nmb.strerror(ex[1], ex[2])
        print ex[0], ':', err_msg
    except smb.SessionError, ex:
        _, err_msg = smb.strerror(ex[1], ex[2])
        print ex[0], ':', err_msg
    except KeyboardInterrupt:
        pass