nxtfilemgr

#!/usr/bin/env python
#
# nxt_filemgr program -- Updated from nxt_filer
# Based on: nxt_filer program -- Simple GUI to manage files on a LEGO Mindstorms NXT
# Copyright (C) 2006  Douglas P Lau
# Copyright (C) 2010  rhn
# Copyright (C) 2017  TC Wan
#
# This program 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 2 of the License, or
# (at your option) any later version.
#
# This program 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.

VERSION="1.02"
CONNECT_STRING = "Connect to NXT"
PROTOCOL_VER = (1, 124)

FILENAME_MINWIDTH = 100
FILENAME_MAXWIDTH = 300
FILENAME_WIDTH = 200
FILESIZE_MINWIDTH = 40
FILESIZE_MAXWIDTH = 120
FILESIZE_WIDTH = 80
FILELIST_HEIGHT = 300

### WIN_WIDTH = 400
### WIN_HEIGHT = 300
WIN_WIDTH = FILENAME_MAXWIDTH + FILESIZE_MAXWIDTH
WIN_HEIGHT = FILELIST_HEIGHT

FILELIST_MINWIDTH = 1.0 * WIN_WIDTH
FILELIST_MAXWIDTH = 1.0 * WIN_WIDTH

# Use Gtk+ 3.x
import gi
gi.require_version("Gtk", "3.0")
from gi.repository import GObject
from gi.repository import Gtk
from gi.repository import GLib

import os.path
import urllib
import sys

import nxt.locator
from nxt.brick import FileFinder, FileReader, FileWriter
from nxt.utils import parse_command_line_arguments

from usb.core import USBError
from nxt.locator import BrickNotFoundError
from nxt.error import FileNotFound, ModuleNotFound, FileExistsError, SysProtError
from nxt.telegram import OPCODES, Telegram


DEBUG = False

def play_soundfile(b, fname):
    if DEBUG:
        print("Playing soundfile %s" % fname)
    b.play_sound_file(0,fname)

def run_program(b, fname):
    if DEBUG:
        print("Running %s" % fname)
    b.start_program(fname)

def delete_file(b, fname):
    b.delete(fname)
    if DEBUG:
        print("Deleted %s" % fname)

def read_file(b, fname):
    with FileReader(b, fname) as r:
        with open(fname, 'wb') as f:
            for data in r:
                f.write(data)

def file_exists_dialog(fname):
    fileexists_str = "Cannot add %s!" % fname
    dialog = Gtk.MessageDialog(win, 0, Gtk.MessageType.INFO,
        Gtk.ButtonsType.OK, "File Exsts")
    dialog.format_secondary_text(fileexists_str)
    dialog.run()
    dialog.destroy()
    if DEBUG:
        print("File Exists box closed.")

def system_error_dialog():
    dialog = Gtk.MessageDialog(win, 0, Gtk.MessageType.INFO,
        Gtk.ButtonsType.OK, "System Error")
    dialog.format_secondary_text("Insufficient Contiguous Space!\n\nPlease delete files to create contiguous space.")
    dialog.run()
    dialog.destroy()
    if DEBUG:
        print("System Error box closed.")

def usb_error_dialog():
    dialog = Gtk.MessageDialog(win, 0, Gtk.MessageType.INFO,
        Gtk.ButtonsType.OK, "USB Error")
    dialog.format_secondary_text("Can't connect to NXT!\n\nIs the NXT powered off?")
    dialog.run()
    dialog.destroy()
    if DEBUG:
        print("USB Error box closed.")

def write_file(b, fname, data):
    try:
        w = FileWriter(b, fname, len(data))
        if DEBUG:
            print("Pushing %s (%d bytes) ..." % (fname, w.size)),
        sys.stdout.flush()
        w.write(data)
        w.close()
        if DEBUG:
            print("wrote %d bytes" % len(data))
    except FileExistsError:
        file_exists_dialog(fname)
    except SysProtError:
        system_error_dialog()
    except USBError:
        raise

def write_files(b, names):
    for fname in names.split('\r\n'):
        if fname:
            if DEBUG:
                print("File:", fname)
            bname = os.path.basename(fname)
            url = urllib.urlopen(fname)
            try:
                data = url.read()
            finally:
                url.close()
            if DEBUG:
                print("name %s, size: %d " % (bname, len(data)))
            write_file(b, bname, data)

class NXTListing(Gtk.ListStore):

    def __init__(self):
        'Create an empty file list'
        Gtk.ListStore.__init__(self, str, str)
        self.set_sort_column_id(0, Gtk.SortType(Gtk.SortType.ASCENDING))
###        self.populate(brick, '*.*')

    def populate(self, brick, pattern):
        f = FileFinder(brick, pattern)
        for (fname_b, size) in f:
            fname = fname_b.decode('utf-8')
            self.append((fname, str(size)))


class NXT_FileMgr(Gtk.Window):
 ### FIXME
 ###   TARGETS = Gtk.target_list_add_uri_targets()

    def __init__(self):
        Gtk.Window.__init__(self)
        self.set_default_size(WIN_WIDTH, WIN_HEIGHT)
        self.set_resizable(False)
        self.set_border_width(6)
        self.brick = None
        self.selected_file = None
        self.brick_info_str = CONNECT_STRING
        self.nxt_filelist = NXTListing()

        h = Gtk.Grid()
        h.set_orientation(Gtk.Orientation.VERTICAL)
        self.brick_button = Gtk.Button(label=self.brick_info_str)
        self.brick_button.connect("clicked", self.reload_filelist)
        h.add(self.brick_button)
        h.add(self.make_file_panel())
        h.add(self.make_button_panel())
        self.add(h)
        self.connect('destroy', self.quit)
        self.reload_filelist(self.nxt_filelist)

    def make_file_view(self):
        tv = Gtk.TreeView()
        tv.set_headers_visible(True)
        tv.set_property('fixed_height_mode', True)
        r = Gtk.CellRendererText()
        c = Gtk.TreeViewColumn('File name', r, text=0)
        c.set_fixed_width(FILENAME_WIDTH)
        c.set_min_width(FILENAME_MINWIDTH)
        c.set_max_width(FILENAME_MAXWIDTH)
        c.set_resizable(True)
        c.set_sizing(Gtk.TreeViewColumnSizing.FIXED)
        tv.append_column(c)
        r = Gtk.CellRendererText()
        c = Gtk.TreeViewColumn('Bytes', r, text=1)
        c.set_resizable(True)
        c.set_fixed_width(FILESIZE_WIDTH)
        c.set_min_width(FILESIZE_MINWIDTH)
        c.set_max_width(FILESIZE_MAXWIDTH)
        c.set_sizing(Gtk.TreeViewColumnSizing.FIXED)
        tv.append_column(c)

### FIXME
#        tv.enable_model_drag_source(Gtk.gdk.BUTTON1_MASK, self.TARGETS,
#        Gtk.gdk.ACTION_DEFAULT | Gtk.gdk.ACTION_MOVE)
#        tv.enable_model_drag_dest(self.TARGETS, Gtk.gdk.ACTION_COPY)
#        tv.connect("drag_data_get", self.drag_data_get_data)

        tv.connect("drag_data_received", self.drag_data_received_data)
        tv.connect("row_activated", self.row_activated_action)
        return tv

    def make_file_panel(self):
        v = Gtk.Box(Gtk.Orientation.VERTICAL, 0)
        tv = self.make_file_view()
        tv.set_model(self.nxt_filelist)
        select = tv.get_selection()
        select.connect("changed", self.on_tree_selection_changed)

        s = Gtk.ScrolledWindow()
        s.set_min_content_width(FILELIST_MINWIDTH)
        s.set_max_content_width(min(FILENAME_MAXWIDTH+FILESIZE_MAXWIDTH, FILELIST_MAXWIDTH))
        s.set_min_content_height(FILELIST_HEIGHT)
        s.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        s.add(tv)
        s.set_border_width(2)
        v.pack_start(s, True, False, 0)
        return v

    def make_button_panel(self):
        vb = Gtk.Box(Gtk.Orientation.HORIZONTAL, 0)

        self.about_button = Gtk.Button(label="?")
        self.about_button.connect("clicked", self.about)
        vb.pack_start(self.about_button, True, True, 10)

        self.delete_button = Gtk.Button(label="Remove")
        self.delete_button.connect("clicked", self.remove_file)
        vb.pack_start(self.delete_button, True, True, 10)

        self.add_button = Gtk.Button(label="Add")
        self.add_button.connect("clicked", self.add_file)
        vb.pack_start(self.add_button, True, True, 10)

        self.exec_button = Gtk.Button(label="Execute")
        self.exec_button.connect("clicked", self.execute_file)
        vb.pack_start(self.exec_button, True, True, 10)
        return vb

    def do_housekeeping(self):
        self.brick.sock.close()
        self.brick = None
        self.selected_file = None

    def	reload_filelist(self, widget):
        if DEBUG:
            print("Reloading filelist...")
        if self.brick != None:
            self.do_housekeeping()
        try:
            self.brick = nxt.locator.find_one_brick(keyword_arguments.get('host',None))
        except BrickNotFoundError:
            print("Brick not found!")
            sys.stdout.flush()
        except USBError:
            usb_error_dialog()
            self.do_housekeeping()
        if self.brick != None:
            self.update_view()

    def update_view(self):
        self.get_brick_info()
        self.brick_button.set_label(self.brick_info_str)
        self.nxt_filelist.clear()
        if self.brick != None:
            self.nxt_filelist.populate(self.brick, '*.*')

    def get_brick_info(self):
        if self.brick != None:
            print("Reading from NXT..."),
            prot_version, fw_version = self.brick.get_firmware_version()
            print("Protocol version: %s.%s" % prot_version)
            if prot_version == PROTOCOL_VER:
                brick_name, brick_hostid, brick_signal_strength, brick_user_flash = self.brick.get_device_info()
                connection_type = str(self.brick.sock)
                if DEBUG:
                    print("Brick Info: ", connection_type, brick_name, brick_hostid, brick_signal_strength, fw_version, brick_user_flash)
                self.brick_info_str = "%s ==> %s [%s]\tFirmware: v%s.%s\tFreespace: %s" % (connection_type, brick_name, brick_hostid, fw_version[0], fw_version[1], brick_user_flash)
            else:
                print("Invalid Protocol version! Closing connection.")
                self.do_housekeeping()
        else:
            self.brick_info_str = CONNECT_STRING
        sys.stdout.flush()

    def choose_files(self):
        dialog = Gtk.FileChooserNative();
        dialog.new("Add File", self, Gtk.FileChooserAction.OPEN, None, None)
        dialog.set_transient_for(self)
        dialog.set_select_multiple(False)
        res = dialog.run()
        if res == Gtk.ResponseType.ACCEPT:
            ### Handle single file download for now
            uri = dialog.get_uri()
            if DEBUG:
                print("You selected", uri)
            try:
                write_files(self.brick, uri)
            except USBError:
                'Cannot recover, close connection'
                usb_error_dialog()
                self.do_housekeeping()
        else:
            if DEBUG:
                print("Add File Cancelled")
        dialog.destroy()

    def on_tree_selection_changed(self, selection):
        model, treeiter = selection.get_selected()
        if treeiter != None:
            if DEBUG:
                print("You selected %s (%s bytes)" % (model[treeiter][0], model[treeiter][1]))
            self.selected_file = model[treeiter][0]
        else:
            self.selected_file = None

    def drag_data_get_data(self, treeview, context, selection, target_id,
        etime):
        treeselection = treeview.get_selection()
        model, iter = treeselection.get_selected()
        data = model.get_value(iter, 0)
        print(data)
        selection.set(selection.target, 8, data)

    def drag_data_received_data(self, treeview, context, x, y, selection,
        info, etime):
        if context.action == Gtk.gdk.ACTION_COPY:
            write_files(self.brick, selection.data)
        # FIXME: update file listing after writing files
        # FIXME: finish context

    def row_activated_action(self, treeview, context, selection):
        self.execute_file(selection)

    def about(self, widget):
        about_str ="\tVersion %s\n\n\n\tBased on nxt_filer\n\n\t(c) 2006-2017, by Douglas P. Lau, rhn and TC Wan\n\n\
\tReleased under GPL v2 or later." % VERSION
        dialog = Gtk.MessageDialog(self, 0, Gtk.MessageType.INFO,
            Gtk.ButtonsType.OK, "NXT File Manager")
        dialog.format_secondary_text(about_str)
        dialog.run()
        dialog.destroy()
        if DEBUG:
            print("About box closed.")

    def	execute_file(self, widget):
        if self.selected_file != None:
            name = ''
            ext = ''
            if DEBUG:
                print("Attempting to execute file...%s" % self.selected_file)
            try:
                name, ext = self.selected_file.split('.')
                if ext == 'rso':
                    play_soundfile(self.brick,self.selected_file)
                elif ext == 'rxe':
                    run_program(self.brick,self.selected_file)
                    self.do_housekeeping()
                    self.update_view()
                else:
                    print("Can't execute '*.%s' files" % ext)
            except ValueError:
                print("No file extension (unknown file type)")
            except USBError:
                'Cannot recover, close connection'
                usb_error_dialog()
                self.do_housekeeping()
                self.update_view()

    def	remove_file(self, widget):
        if self.selected_file != None:
            warning_str = "Do you really want to remove %s?" % self.selected_file
            if DEBUG:
                print("Removing file...%s" % self.selected_file)
            dialog = Gtk.MessageDialog(self, 0, Gtk.MessageType.WARNING,
                Gtk.ButtonsType.OK_CANCEL, "Remove File")
            dialog.format_secondary_text(warning_str)
            response = dialog.run()
            dialog.destroy()
            if response == Gtk.ResponseType.OK:
                if DEBUG:
                    print("Proceed to Delete File")
                try:
                    delete_file(self.brick,self.selected_file)
                except USBError:
                    'Cannot recover, close connection'
                    usb_error_dialog()
                    self.do_housekeeping()
                self.selected_file = None
                self.update_view()
            elif response == Gtk.ResponseType.CANCEL:
                if DEBUG:
                    print("CANCEL Delete File")

    def	add_file(self, widget):
        if self.brick != None:
            if DEBUG:
                print("Adding file...")
            self.choose_files()
            self.selected_file = None
            self.update_view()

    def quit(self, w):
        'Quit the program'
        if self.brick != None:
            self.do_housekeeping()
        Gtk.main_quit()

if __name__ == '__main__':
    # FIXME: add dialog with progress bar when scanning bluetooth devices
    arguments, keyword_arguments = parse_command_line_arguments(sys.argv)

    if '-v' in arguments:
        print("Enable Verbose Mode")
        DEBUG = True

    if '--help' in arguments:
        print("nxt_filemgr -- Simple GUI to manage files on a LEGO Mindstorms NXT\n\
Usage:	nxt_filemgr [--host <macaddress>]")
        exit(0)

    win = NXT_FileMgr()
    win.show_all()
    Gtk.main()