A small AI-Generated Samba client written in Python.

tinysamba.py

import os
import uuid
import sys

from smbprotocol.connection import Connection
from smbprotocol.session import Session
from smbprotocol.tree import TreeConnect
from smbprotocol.open import (
    Open,
    ImpersonationLevel,
    DirectoryAccessMask,
    FilePipePrinterAccessMask,
    FileAttributes,
    ShareAccess,
    CreateDisposition,
    CreateOptions,
)
from smbprotocol.file_info import FileDirectoryInformation, FileInformationClass
from smbprotocol.exceptions import SMBResponseException


SMB_PORT = 445


class SMBSession:
    def __init__(self, host, port=SMB_PORT):
        self.host = host
        self.port = port
        self.connection = None
        self.session = None

    def connect(self, username="", password=""):
        self.connection = Connection(uuid.uuid4(), self.host, self.port)
        self.connection.connect()
        self.session = Session(
            self.connection,
            username=username,
            password=password,
            require_encryption=False,
        )
        self.session.connect()
        return True

    def list_directory(self, share, path=""):
        tree = TreeConnect(self.session, f"\\\\{self.host}\\{share}")
        tree.connect()

        share_path = path.replace("/", "\\") if path else ""
        directory = Open(tree, share_path)
        directory.create(
            ImpersonationLevel.Impersonation,
            DirectoryAccessMask.FILE_LIST_DIRECTORY
            | DirectoryAccessMask.FILE_READ_ATTRIBUTES,
            FileAttributes.FILE_ATTRIBUTE_DIRECTORY,
            ShareAccess.FILE_SHARE_READ | ShareAccess.FILE_SHARE_WRITE,
            CreateDisposition.FILE_OPEN,
            CreateOptions.FILE_DIRECTORY_FILE,
        )

        entries = directory.query_directory(
            "*", FileInformationClass.FILE_DIRECTORY_INFORMATION
        )
        results = []
        for entry in entries:
            name = entry["file_name"].get_value().decode("utf-16-le").rstrip("\x00")
            if name in (".", ".."):
                continue
            is_dir = bool(
                entry["file_attributes"].get_value()
                & FileAttributes.FILE_ATTRIBUTE_DIRECTORY
            )
            size = entry["end_of_file"].get_value()
            results.append((name, is_dir, size))

        directory.close()
        tree.disconnect()
        return results

    def upload_file(self, share, local_path, remote_name=None):
        if not os.path.isfile(local_path):
            raise FileNotFoundError(f"Local file not found: {local_path}")

        if remote_name is None:
            remote_name = os.path.basename(local_path)

        file_size = os.path.getsize(local_path)
        tree = TreeConnect(self.session, f"\\\\{self.host}\\{share}")
        tree.connect()

        file = Open(tree, remote_name)
        file.create(
            ImpersonationLevel.Impersonation,
            FilePipePrinterAccessMask.GENERIC_WRITE,
            FileAttributes.FILE_ATTRIBUTE_NORMAL,
            ShareAccess.FILE_SHARE_READ | ShareAccess.FILE_SHARE_WRITE,
            CreateDisposition.FILE_OVERWRITE_IF,
            CreateOptions.FILE_NON_DIRECTORY_FILE,
        )

        chunk_size = self.connection.max_write_size
        offset = 0
        with open(local_path, "rb") as f:
            while True:
                data = f.read(chunk_size)
                if not data:
                    break
                file.write(data, offset)
                offset += len(data)
                pct = offset * 100 // file_size
                print(
                    f"\r  Uploading... {pct}% ({offset:,}/{file_size:,} bytes)", end=""
                )

        file.close()
        tree.disconnect()
        print()
        return remote_name

    def close(self):
        if self.connection:
            self.connection.disconnect()
            self.connection = None
            self.session = None


def main():
    if len(sys.argv) < 5:
        print("Usage:")
        print(
            "  python smb_client.py <host> <user> <pass> <share> [path]         # list directory"
        )
        print(
            "  python smb_client.py <host> <user> <pass> <share> upload <local> [remote_name]  # upload file"
        )
        sys.exit(1)

    host = sys.argv[1]
    username = sys.argv[2]
    password = sys.argv[3]
    share = sys.argv[4]

    session = SMBSession(host)

    try:
        print(f"Connecting to {host}:{SMB_PORT}...")
        session.connect(username, password)
        print("Connected!")

        if len(sys.argv) > 5 and sys.argv[5] == "upload":
            local_path = sys.argv[6]
            remote_name = sys.argv[7] if len(sys.argv) > 7 else None
            print(f"\nUploading {local_path} to \\\\{host}\\{share}\\...")
            session.upload_file(share, local_path, remote_name)
            print("Upload complete!")
        else:
            path = sys.argv[5] if len(sys.argv) > 5 else ""
            print(f"\nContents of \\\\{host}\\{share}\\{path}:")
            entries = session.list_directory(share, path)
            for name, is_dir, size in entries:
                if is_dir:
                    print(f"  [DIR]  {name}")
                else:
                    print(f"  [FILE] {name:50s} {size:,} bytes")

    except SMBResponseException as e:
        print(f"\nSMB error: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"\nError: {e}")
        sys.exit(1)
    finally:
        session.close()


if __name__ == "__main__":
    main()