/*
 * Copyleft (c) 2002, 2012
 * The NSFOCUS INFORMATION TECHNOLOGY CO.,LTD.
 * -----------------------------------------------------------------------
 * Author   : NSFocus Security Team <security@nsfocus.com>
 *          : http://www.nsfocus.com
 * Maintain : scz <scz@nsfocus.com>
 * Version  : 1.20
 * Compile  : For x86/EWindows XP SP1 & VC 7
 *          : cl enumshare.c /nologo /Os /G6 /Gs65536 /W3 /D "WIN32" /D "NDEBUG" /D "_CONSOLE" /link /RELEASE
 *          :
 * Create   : 2003-11-20 10:21
 * Modify   : 2003-11-21 15:28
 * -----------------------------------------------------------------------
 * The only thing they can't take from us are our minds. !H
 */

/*
 * 如果同时指定了target、username、password，就试着注销并重新建立SMB会话，
 * 最终注销SMB会话。否则本程序不影响SMB会话。
 *
 * You can also use the WNetEnumResource function to retrieve resource
 * information. However, WNetEnumResource does not enumerate hidden shares
 * or users connected to a share.
 */

/************************************************************************
 *                                                                      *
 *                               Head File                              *
 *                                                                      *
 ************************************************************************/

/*
 * #define _WIN32_WINNT 0x0501
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <windows.h>
#include <winnetwk.h>
#include <lm.h>

/************************************************************************
 *                                                                      *
 *                               Macro                                  *
 *                                                                      *
 ************************************************************************/

#pragma comment( linker, "/INCREMENTAL:NO"    )
#pragma comment( linker, "/subsystem:console" )
#pragma comment( lib,    "kernel32.lib"       )
#pragma comment( lib,    "mpr.lib"            )
#pragma comment( lib,    "netapi32.lib"       )

#define VERSION             "1.20"

/************************************************************************
 *                                                                      *
 *                            Function Prototype                        *
 *                                                                      *
 ************************************************************************/

static void enumshare          ( unsigned char *target );
static void PrintWin32ErrorCUI ( char *message, DWORD dwMessageId );
static BOOL SessionBegin
(
    unsigned char *username,
    unsigned char *password,
    unsigned char *resource
);
static BOOL SessionEnd         ( unsigned char *resource );
static void usage              ( char *arg );

/************************************************************************
 *                                                                      *
 *                            Static Global Var                         *
 *                                                                      *
 ************************************************************************/

/************************************************************************/

static void enumshare ( unsigned char *target )
{
    unsigned char  *ansiserver    = NULL,
                   *ansishare     = NULL,
                   *ansiremark    = NULL,
                   *ansibuf_0     = NULL,
                   *ansibuf_1     = NULL;
    WCHAR          *unicodeserver = NULL,
                   *share         = NULL,
                   *remark        = NULL;
    unsigned int    len           = 0;
    int             i             = 0,
                    j             = 0;
    LPBYTE          buf           = NULL;
    PSHARE_INFO_1   shareinfo     = NULL;
    NET_API_STATUS  error;
    DWORD           entriesread   = 0,
                    totalentries  = 0,
                    resume_handle = 0,
                    count,
                    type;
    unsigned char * shi1_type[]   =
    {
        "磁盘驱动器",
        "打印机",
        "通信设备",
        "9x/Samba进程间通信",
        "管理共享",
        "NT/2K/XP/2003进程间通信",
        "未知类型"
    };

    if ( NULL != target )
    {
        /*
         * "\\<target>"，包括结尾的\0
         */
        j             = 2 + strlen( target ) + 1;
        ansiserver    = ( unsigned char * )HeapAlloc
                                           (
                                               GetProcessHeap(),
                                               HEAP_ZERO_MEMORY,
                                               j
                                           );
        if ( NULL == ansiserver )
        {
            PrintWin32ErrorCUI( "HeapAlloc() failed [0]", ERROR_NOT_ENOUGH_MEMORY );
            goto enumshare_exit;
        }
        sprintf( ansiserver, "\\\\%s", target );
        j             = strlen( ansiserver ) + 1;
        i             = MultiByteToWideChar
                        (
                            CP_ACP,
                            0,
                            ansiserver,
                            j,
                            NULL,
                            0
                        );
        if ( 0 == i )
        {
            PrintWin32ErrorCUI( "MultiByteToWideChar() failed [0]", GetLastError() );
            goto enumshare_exit;
        }
        unicodeserver = ( WCHAR * )HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY, i * sizeof( WCHAR ) );
        if ( NULL == unicodeserver )
        {
            PrintWin32ErrorCUI( "HeapAlloc() failed [1]", ERROR_NOT_ENOUGH_MEMORY );
            goto enumshare_exit;
        }
        if ( 0 == MultiByteToWideChar
                  (
                      CP_ACP,
                      0,
                      ansiserver,
                      j,
                      unicodeserver,
                      i
                  ) )
        {
            PrintWin32ErrorCUI( "MultiByteToWideChar() failed [1]", GetLastError() );
            goto enumshare_exit;
        }
    }
    error     = NetShareEnum
                (
                    ( LPSTR )unicodeserver,
                    1,
                    &buf,
                    MAX_PREFERRED_LENGTH,
                    &entriesread,
                    &totalentries,
                    &resume_handle
                );
    if ( NERR_Success != error )
    {
        PrintWin32ErrorCUI( "NetShareEnum() failed", error );
        goto enumshare_exit;
    }
    printf
    (
        "read entries  = %u\n"
        "total entries = %u\n\n",
        entriesread,
        totalentries
    );
    shareinfo = ( PSHARE_INFO_1 )buf;
    for ( count = 0; count < entriesread; count++, shareinfo++ )
    {
        switch ( shareinfo->shi1_type )
        {
        case STYPE_DISKTREE:
            type = 0;
            break;
        case STYPE_PRINTQ:
            type = 1;
            break;
        case STYPE_DEVICE:
            type = 2;
            break;
        case STYPE_IPC:
            type = 3;
            break;
        case STYPE_SPECIAL:
            type = 4;
            break;
        case 0x80000003:
            type = 5;
            break;
        default:
            type = 6;
            break;
        }  /* end of switch */
        share  = ( WCHAR * )shareinfo->shi1_netname;
        remark = ( WCHAR * )shareinfo->shi1_remark;
        len    = ( wcslen( share ) + 1 ) * 2;
        if ( 0 == wcscmp( share, L"" ) )
        {
            ansishare = "(No sharename)";
        }
        else
        {
            ansibuf_0 = ( unsigned char * )HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY, len );
            if ( NULL == ansibuf_0 )
            {
                ansishare = "(No memory for ansibuf_0)";
            }
            else
            {
                i = WideCharToMultiByte
                    (
                        CP_ACP,
                        0,
                        share,
                        ( int )( len / 2 ),
                        ansibuf_0,
                        len,
                        NULL,
                        NULL
                    );
                if ( 0 == i )
                {
                    ansishare = "(WideCharToMultiByte failed 0)";
                }
                else
                {
                    ansishare = ansibuf_0;
                }
            }
        }
        len    = ( wcslen( remark ) + 1 ) * 2;
        if ( 0 == wcscmp( remark, L"" ) )
        {
            ansiremark = "(没有注释)";
        }
        else
        {
            ansibuf_1 = ( unsigned char * )HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY, len );
            if ( NULL == ansibuf_1 )
            {
                ansiremark = "(No memory for ansibuf_1)";
            }
            else
            {
                i = WideCharToMultiByte
                    (
                        CP_ACP,
                        0,
                        remark,
                        ( int )( len / 2 ),
                        ansibuf_1,
                        len,
                        NULL,
                        NULL
                    );
                if ( 0 == i )
                {
                    ansiremark = "(WideCharToMultiByte failed 1)";
                }
                else
                {
                    ansiremark = ansibuf_1;
                }
            }
        }
        printf
        (
            "%-15s %-30s %s\n",
            ansishare,
            ansiremark,
            shi1_type[type]
        );
        if ( NULL != ansibuf_0 )
        {
            HeapFree( GetProcessHeap(), 0, ansibuf_0 );
            ansibuf_0 = NULL;
        }
        if ( NULL != ansibuf_1 )
        {
            HeapFree( GetProcessHeap(), 0, ansibuf_1 );
            ansibuf_1 = NULL;
        }
    }  /* end of for */

enumshare_exit:

    if ( NULL != buf )
    {
        NetApiBufferFree( buf );
        buf = NULL;
    }
    if ( NULL != unicodeserver )
    {
        HeapFree( GetProcessHeap(), 0, unicodeserver );
        unicodeserver = NULL;
    }
    if ( NULL != ansiserver )
    {
        HeapFree( GetProcessHeap(), 0, ansiserver );
        ansiserver = NULL;
    }
    return;
}  /* end of enumshare */

static void PrintWin32ErrorCUI ( char *message, DWORD dwMessageId )
{
    char *errMsg;

    FormatMessage
    (
        FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
        NULL,
        dwMessageId,
        MAKELANGID( LANG_NEUTRAL, SUBLANG_DEFAULT ),
        ( LPTSTR )&errMsg,
        0,
        NULL
    );
    fprintf( stderr, "%s: %s", message, errMsg );
    LocalFree( errMsg );
    return;
}  /* end of PrintWin32ErrorCUI */

static BOOL SessionBegin
(
    unsigned char *username,
    unsigned char *password,
    unsigned char *resource
)
{
    BOOL        ret = FALSE;
    DWORD       error;
    NETRESOURCE netresource;

    SessionEnd( resource );
    ZeroMemory( &netresource, sizeof( netresource ) );
    netresource.dwType       = RESOURCETYPE_ANY;
    netresource.lpLocalName  = NULL;
    netresource.lpRemoteName = resource;
    netresource.lpProvider   = NULL;
    error = WNetAddConnection2
            (
                &netresource,
                password,
                username,
                0
            );
    if ( NO_ERROR != error )
    {
        PrintWin32ErrorCUI( "WNetAddConnection2() failed", error );
        goto SessionBegin_exit;
    }
    ret = TRUE;

SessionBegin_exit:

    return( ret );
}  /* end of SessionBegin */

static BOOL SessionEnd ( unsigned char *resource )
{
    BOOL  ret = FALSE;
    DWORD error;

    error = WNetCancelConnection2
            (
                resource,
                CONNECT_UPDATE_PROFILE,
                TRUE
            );
    if ( NO_ERROR != error )
    {
        goto SessionEnd_exit;
    }
    ret   = TRUE;

SessionEnd_exit:

    return( ret );
}  /* end of SessionEnd */

static void usage ( char *arg )
{
    fprintf
    (
        stderr,
        "Usage: %s [-h] [-v] [-t target] [-u username] [-p password]\n",
        arg
    );
    exit( EXIT_FAILURE );
}  /* end of usage */

int __cdecl main ( int argc, char * argv[] )
{
    int            c,
                   ret           = EXIT_FAILURE;
    /*
     * If this string is NULL, the account name is looked up on the local
     * system.
     */
    unsigned char *target        = NULL,
                  *username      = NULL,
                  *password      = NULL,
                  *resource      = NULL;
    unsigned int   resourcelen   = 0;

    if ( 1 == argc )
    {
        usage( argv[0] );
    }
    for ( c = 1; c < argc; c++ )
    {
        if ( ( ( argv[c][0] != '-' ) && ( argv[c][0] != '/' ) ) || ( strlen( argv[c] ) < 2 ) )
        {
            usage( argv[0] );
        }
        else
        {
            switch ( tolower( argv[c][1] ) )
            {
            case 't':
                if ( ( c + 1 ) >= argc )
                {
                    usage( argv[0] );
                }
                target      = argv[++c];
                break;
            case 'u':
                if ( ( c + 1 ) >= argc )
                {
                    usage( argv[0] );
                }
                username    = argv[++c];
                break;
            case 'p':
                if ( ( c + 1 ) >= argc )
                {
                    usage( argv[0] );
                }
                password    = argv[++c];
                break;
            case 'v':
                fprintf( stderr, "%s ver "VERSION"\n", argv[0] );
                return( EXIT_SUCCESS );
            case 'h':
            case '?':
            default:
                usage( argv[0] );
                break;
            }  /* end of switch */
        }
    }  /* end of for */
    if ( FALSE == SetConsoleCtrlHandler( NULL, FALSE ) )
    {
        PrintWin32ErrorCUI( "SetConsoleCtrlHandler() failed", GetLastError() );
        goto main_exit;
    }
    if ( NULL != target && NULL != username && NULL != password )
    {
        resourcelen = 2 + strlen( target ) + 1 + 4 + 1;
        resource    = ( unsigned char * )HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY, resourcelen );
        if ( NULL == resource )
        {
            PrintWin32ErrorCUI( "HeapAlloc() failed", ERROR_NOT_ENOUGH_MEMORY );
            goto main_exit;
        }
        sprintf( resource, "\\\\%s\\IPC$", target );
        SessionBegin( username, password, resource );
    }
    else
    {
        printf( "\n[Assuming one session already existed or target is null.]\n\n" );
    }
    enumshare( target );
    ret = EXIT_SUCCESS;

main_exit:

    if ( NULL != resource )
    {
        SessionEnd( resource );
        HeapFree( GetProcessHeap(), 0, resource );
        resource = NULL;
    }
    return( ret );
}  /* end of main */

/************************************************************************/