uuu/uuu.cpp

   1 /*
   2 * Copyright 2018 NXP.
   3 *
   4 * Redistribution and use in source and binary forms, with or without modification,
   5 * are permitted provided that the following conditions are met:
   6 *
   7 * Redistributions of source code must retain the above copyright notice, this
   8 * list of conditions and the following disclaimer.
   9 *
  10 * Redistributions in binary form must reproduce the above copyright notice, this
  11 * list of conditions and the following disclaimer in the documentation and/or
  12 * other materials provided with the distribution.
  13 *
  14 * Neither the name of the NXP Semiconductor nor the names of its
  15 * contributors may be used to endorse or promote products derived from this
  16 * software without specific prior written permission.
  17 *
  18 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  19 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  20 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  21 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
  22 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  23 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  24 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  25 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  26 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  27 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  28 * POSSIBILITY OF SUCH DAMAGE.
  29 *
  30 */
  31
  32 #include <iostream>
  33 #include <stdio.h>
  34 #include <thread>
  35 #include <atomic>
  36 #include <iomanip>
  37 #include <map>
  38 #include <mutex>
  39 #include <vector>
  40 #include <sstream>
  41 #include <fstream>
  42 #include <stdarg.h>
  43 #include <time.h>
  44 #include <string.h>
  45 #include <signal.h>
  46 #include "buildincmd.h"
  47 #include <string>
  48 #include <streambuf>
  49
  50 #include "../libuuu/libuuu.h"
  51
  52 const char * g_vt_yellow = "\x1B[93m";
  53 const char * g_vt_default = "\x1B[0m";
  54 const char * g_vt_green = "\x1B[92m";
  55 const char * g_vt_red = "\x1B[91m";
  56 const char * g_vt_kcyn = "\x1B[36m";
  57 const char * g_vt_boldwhite = "\x1B[97m";
  58
  59 void clean_vt_color() noexcept
  60 {
  61         g_vt_yellow = "";
  62         g_vt_default = g_vt_yellow;
  63         g_vt_green = g_vt_yellow;
  64         g_vt_red = g_vt_yellow;
  65         g_vt_kcyn = g_vt_yellow;
  66         g_vt_boldwhite = g_vt_yellow;
  67 }
  68
  69 using namespace std;
  70
  71 int get_console_width();
  72 void print_oneline(string str);
  73 int auto_complete(int argc, char**argv);
  74 void print_autocomplete_help();
  75
  76 char g_sample_cmd_list[] = {
  77 #include "uuu.clst"
  78 };
  79
  80 vector<string> g_usb_path_filter;
  81
  82 int g_verbose = 0;
  83 static bool g_start_usb_transfer;
  84
  85 bmap_mode g_bmap_mode = bmap_mode::Default;
  86
  87 class AutoCursor
  88 {
  89 public:
  90         ~AutoCursor()
  91         {
  92                 printf("\x1b[?25h\n\n\n");
  93         }
  94 };
  95
  96 void ctrl_c_handle(int)
  97 {
  98         do {
  99                 AutoCursor a;
 100         } while(0);
 101
 102         exit(1);
 103 }
 104
 105 class string_ex : public std::string
 106 {
 107 public:
 108         int format(const char *fmt, ...)
 109         {
 110                 va_list args;
 111                 va_start(args, fmt);
 112                 size_t len = std::vsnprintf(NULL, 0, fmt, args);
 113                 va_end(args);
 114
 115                 this->resize(len);
 116
 117                 va_start(args, fmt);
 118                 std::vsnprintf((char*)c_str(), len + 1, fmt, args);
 119                 va_end(args);
 120
 121                 return 0;
 122         }
 123 };
 124
 125 #ifdef _WIN32
 126 #include <conio.h>
 127 #else
 128 #include <termios.h>
 129 #include <unistd.h>
 130 #endif
 131
 132 int ask_passwd(char* prompt, char user[MAX_USER_LEN], char passwd[MAX_USER_LEN])
 133 {
 134         cout << endl << prompt << " Required Login"<<endl;
 135         cout << "Username:";
 136         cin.getline(user, 128);
 137         cout << "Password:";
 138         int i = 0;
 139
 140 #ifdef _WIN32
 141         while ((passwd[i] = _getch()) != '\r') {
 142                 if (passwd[i] == '\b') {
 143                         if (i != 0) {
 144                                 cout << "\b \b";
 145                                 i--;
 146                         }
 147                 }
 148                 else {
 149                         cout << '*';
 150                         i++;
 151                 }
 152         }
 153 #else
 154         struct termios old, tty;
 155         tcgetattr(STDIN_FILENO, &tty);
 156         old = tty;
 157         tty.c_lflag &= ~ECHO;
 158         tcsetattr(STDIN_FILENO, TCSANOW, &tty);
 159
 160         string pd;
 161         getline(cin, pd);
 162
 163         tcsetattr(STDIN_FILENO, TCSANOW, &old);
 164         if(pd.size() > MAX_USER_LEN -1)
 165                 return -1;
 166         memcpy(passwd, pd.data(), pd.size());
 167         i=pd.size();
 168
 169 #endif
 170         passwd[i] = 0;
 171         cout << endl;
 172         return 0;
 173 }
 174
 175 void print_help(bool detail = false)
 176 {
 177         const char help[] =
 178                 "uuu [-d -m -v -V -bmap -no-bmap] <" "bootloader|cmdlists|cmd" ">\n\n"
 179                 "    bootloader  download bootloader to board by usb\n"
 180                 "    cmdlist     run all commands in cmdlist file\n"
 181                 "                If it is path, search uuu.auto in dir\n"
 182                 "                If it is zip, search uuu.auto in zip\n"
 183                 "    cmd         Run one command, use -H see detail\n"
 184                 "                example: SDPS: boot -f flash.bin\n"
 185                 "    -d          Daemon mode, wait for forever.\n"
 186                 "    -v -V       verbose mode, -V enable libusb error\\warning info\n"
 187                 "    -dry        Dry run mode, check if script or cmd correct \n"
 188                 "    -bmap       Try using .bmap files even if flash commands do not specify them\n"
 189                 "    -no-bmap    Ignore .bmap files even if flash commands specify them\n"
 190                 "    -m          USBPATH Only monitor these paths.\n"
 191                 "                    -m 1:2 -m 1:3\n\n"
 192                 "    -t          Timeout second for wait known usb device appeared\n"
 193                 "    -T          Timeout second for wait next known usb device appeared at stage switch\n"
 194                 "    -e          set environment variable key=value\n"
 195                 "    -pp         usb polling period in milliseconds\n"
 196                 "    -dm         disable small memory\n"
 197                 "uuu -s          Enter shell mode. uuu.inputlog record all input commands\n"
 198                 "                you can use \"uuu uuu.inputlog\" next time to run all commands\n\n"
 199                 "uuu -udev       linux: show udev rule to avoid sudo each time \n"
 200                 "uuu -lsusb      List connected know devices\n"
 201                 "uuu -IgSerNum   Set windows registry to ignore USB serial number for known uuu devices\n"
 202                 "uuu -h          show general help\n"
 203                 "uuu -H          show general help and detailed help for commands\n\n";
 204         printf("%s", help);
 205         printf("uuu [-d -m -v -bmap -no-bmap] -b[run] ");
 206         g_BuildScripts.ShowCmds();
 207         printf(" arg...\n");
 208         printf("\tRun Built-in scripts\n");
 209         g_BuildScripts.ShowAll();
 210         printf("\nuuu -bshow ");
 211         g_BuildScripts.ShowCmds();
 212         printf("\n");
 213         printf("\tShow built-in script\n");
 214         printf("\n");
 215
 216         print_autocomplete_help();
 217
 218         if (detail == false)
 219                 return;
 220
 221         size_t start = 0, pos = 0;
 222         string str= g_sample_cmd_list;
 223
 224         bool bprint = false;
 225         while ((pos = str.find('\n',pos)) != str.npos)
 226         {
 227                 string s = str.substr(start, pos - start);
 228                 if (s.substr(0, 6) == "# ----")
 229                         bprint = true;
 230
 231                 if (bprint)
 232                 {
 233                         if (s[0] == '#')
 234                         {
 235                                 printf("%s\n", &(s[1]));
 236                         }
 237                 }
 238                 pos += 1;
 239                 start = pos;
 240         }
 241 }
 242 void print_version()
 243 {
 244         printf("uuu (Universal Update Utility) for nxp imx chips -- %s\n\n", uuu_get_version_string());
 245 }
 246
 247 int print_cfg(const char *pro, const char * chip, const char * /*compatible*/, uint16_t vid, uint16_t pid, uint16_t bcdmin, uint16_t bcdmax, void * /*p*/)
 248 {
 249         const char *ext;
 250         if (strlen(chip) >= 7)
 251                 ext = "";
 252         else
 253                 ext = "\t";
 254
 255         if (bcdmin == 0 && bcdmax == 0xFFFF)
 256                 printf("\t%s\t %s\t%s 0x%04x\t 0x%04x\n", pro, chip, ext, vid, pid);
 257         else
 258                 printf("\t%s\t %s\t%s 0x%04x\t 0x%04x\t [0x%04x..0x%04x]\n", pro, chip, ext, vid, pid, bcdmin, bcdmax);
 259         return 0;
 260 }
 261
 262 int print_udev_rule(const char * /*pro*/, const char * /*chip*/, const char * /*compatible*/,
 263         uint16_t vid, uint16_t pid, uint16_t /*bcdmin*/, uint16_t /*bcdmax*/, void * /*p*/)
 264 {
 265         printf("SUBSYSTEM==\"usb\", ATTRS{idVendor}==\"%04x\", ATTRS{idProduct}==\"%04x\", TAG+=\"uaccess\"\n",
 266                         vid, pid);
 267         return 0;
 268 }
 269
 270 int polling_usb(std::atomic<int>& bexit);
 271
 272 int g_overall_status;
 273 int g_overall_okay;
 274 int g_overall_failure;
 275 char g_wait[] = "|/-\\";
 276 int g_wait_index;
 277
 278
 279 string build_process_bar(size_t width, size_t pos, size_t total)
 280 {
 281         string str;
 282         str.resize(width, ' ');
 283         str[0] = '[';
 284         str[width - 1] = ']';
 285
 286         if (total == 0)
 287         {
 288                 if (pos == 0)
 289                         return str;
 290
 291                 string_ex loc;
 292                 size_t s = pos / (1024 * 1024);
 293                 loc.format("%dM", s);
 294                 str.replace(1, loc.size(), loc);
 295                 return str;
 296         }
 297
 298         size_t i;
 299
 300         if (pos > total)
 301                 pos = total;
 302
 303         for (i = 1; i < (width-2) * pos / total; i++)
 304         {
 305                 str[i] = '=';
 306         }
 307
 308         if (i > 1)
 309                 str[i] = '>';
 310
 311         if (pos == total)
 312                 str[str.size() - 2] = '=';
 313
 314         string_ex per;
 315         per.format("%d%%", pos * 100 / total);
 316
 317         size_t start = (width - per.size()) / 2;
 318         str.replace(start, per.size(), per);
 319         str.insert(start, g_vt_yellow);
 320         str.insert(start + per.size() + strlen(g_vt_yellow), g_vt_default);
 321         return str;
 322 }
 323
 324 void print_auto_scroll(string str, size_t len, size_t start)
 325 {
 326         if (str.size() <= len)
 327         {
 328                 str.resize(len, ' ');
 329                 cout << str;
 330                 return;
 331         }
 332
 333         if(str.size())
 334                 start = start % str.size();
 335         else
 336                 start = 0;
 337
 338         string s = str.substr(start, len);
 339         s.resize(len, ' ');
 340         cout << s;
 341 }
 342 class ShowNotify
 343 {
 344 public:
 345         string m_cmd;
 346         string m_dev;
 347         size_t m_trans_pos = 0;
 348         int m_status = 0;
 349         size_t m_cmd_total = 0;
 350         size_t m_cmd_index = 0;
 351         string m_last_err;
 352         int m_done = 0;
 353         size_t m_start_pos = 0;
 354         size_t  m_trans_size = 0;
 355         clock_t m_start_time;
 356         uint64_t m_cmd_start_time;
 357         uint64_t m_cmd_end_time;
 358         bool m_IsEmptyLine = false;
 359
 360         ShowNotify() : m_start_time{clock()} {}
 361
 362         bool update(uuu_notify nt)
 363         {
 364                 if (nt.type == uuu_notify::NOTIFY_DEV_ATTACH)
 365                 {
 366                         m_dev = nt.str;
 367                         m_done = 0;
 368                         m_status = 0;
 369                 }
 370                 if (nt.type == uuu_notify::NOTIFY_CMD_START)
 371                 {
 372                         m_start_pos = 0;
 373                         m_cmd = nt.str;
 374                         m_cmd_start_time = nt.timestamp;
 375                 }
 376                 if (nt.type == uuu_notify::NOTIFY_DECOMPRESS_START)
 377                 {
 378                         m_start_pos = 0;
 379                         m_cmd = nt.str;
 380                         m_cmd_start_time = nt.timestamp;
 381                         m_dev = "Prep";
 382                 }
 383                 if (nt.type == uuu_notify::NOTIFY_DOWNLOAD_START)
 384                 {
 385                         m_start_pos = 0;
 386                         m_cmd = nt.str;
 387                         m_cmd_start_time = nt.timestamp;
 388                         m_dev = "Prep";
 389                 }
 390                 if (nt.type == uuu_notify::NOTIFY_DOWNLOAD_END)
 391                 {
 392                         m_IsEmptyLine = true;
 393                 }
 394                 if (nt.type == uuu_notify::NOTIFY_TRANS_SIZE || nt.type == uuu_notify::NOTIFY_DECOMPRESS_SIZE)
 395                 {
 396                         m_trans_size = nt.total;
 397                         return false;
 398                 }
 399                 if (nt.type == uuu_notify::NOTIFY_CMD_TOTAL)
 400                 {
 401                         m_cmd_total = nt.total;
 402                         return false;
 403                 }
 404                 if (nt.type == uuu_notify::NOTIFY_CMD_INDEX)
 405                 {
 406                         m_cmd_index = nt.index;
 407                         return false;
 408                 }
 409                 if (nt.type == uuu_notify::NOTIFY_DONE)
 410                 {
 411                         if (m_status)
 412                                 g_overall_failure++;
 413                         else
 414                                 g_overall_okay++;
 415
 416                         m_done = 1;
 417                 }
 418                 if (nt.type == uuu_notify::NOTIFY_CMD_END)
 419                 {
 420                         m_cmd_end_time = nt.timestamp;
 421                         if(nt.status)
 422                         {
 423                                 g_overall_status = nt.status;
 424                                 m_last_err = uuu_get_last_err_string();
 425                         }
 426                         m_status |= nt.status;
 427                         if (m_status)
 428                                 g_overall_failure++;
 429                 }
 430                 if (nt.type == uuu_notify::NOTIFY_TRANS_POS || nt.type == uuu_notify::NOTIFY_DECOMPRESS_POS)
 431                 {
 432                         if (m_trans_size == 0) {
 433
 434                                 m_trans_pos = nt.index;
 435                                 return true;
 436                         }
 437
 438                         if ((nt.index - m_trans_pos) < (m_trans_size / 100)
 439                                 && nt.index != m_trans_size)
 440                                 return false;
 441
 442                         m_trans_pos = nt.index;
 443                 }
 444
 445                 return true;
 446         }
 447         void print_verbose(uuu_notify*nt)
 448         {
 449                 if (this->m_dev == "Prep" && g_start_usb_transfer)
 450                         return;
 451
 452                 if (nt->type == uuu_notify::NOTIFY_DEV_ATTACH)
 453                 {
 454                         cout << "New USB Device Attached at " << nt->str << endl;
 455                 }
 456                 if (nt->type == uuu_notify::NOTIFY_CMD_START)
 457                 {
 458                         cout << m_dev << ">" << "Start Cmd:" << nt->str << endl;
 459                 }
 460                 if (nt->type == uuu_notify::NOTIFY_CMD_END)
 461                 {
 462                         double diff = m_cmd_end_time - m_cmd_start_time;
 463                         diff /= 1000;
 464                         if (nt->status)
 465                         {
 466                                 cout << m_dev << ">" << g_vt_red <<"Fail " << uuu_get_last_err_string() << "("<< std::setprecision(4) << diff << "s)" <<  g_vt_default << endl;
 467                         }
 468                         else
 469                         {
 470                                 cout << m_dev << ">" << g_vt_green << "Okay ("<< std::setprecision(4) << diff << "s)" << g_vt_default << endl;
 471                         }
 472                 }
 473
 474                 if (nt->type == uuu_notify::NOTIFY_TRANS_POS || nt->type == uuu_notify::NOTIFY_DECOMPRESS_POS)
 475                 {
 476                         if (m_trans_size)
 477                                 cout << g_vt_yellow << "\r" << m_trans_pos * 100 / m_trans_size <<"%" << g_vt_default;
 478                         else
 479                                 cout << "\r" << m_trans_pos;
 480
 481                         cout.flush();
 482                 }
 483
 484                 if (nt->type == uuu_notify::NOTIFY_CMD_INFO)
 485                         cout << nt->str;
 486
 487                 if (nt->type == uuu_notify::NOTIFY_WAIT_FOR)
 488                         cout << "\r" << nt->str << " "<< g_wait[((g_wait_index++) & 0x3)];
 489
 490                 if (nt->type == uuu_notify::NOTIFY_DECOMPRESS_START)
 491                         cout << "Decompress file:" << nt->str << endl;
 492
 493                 if (nt->type == uuu_notify::NOTIFY_DOWNLOAD_START)
 494                         cout << "Download file:" << nt->str << endl;
 495
 496         }
 497         void print(int verbose = 0, uuu_notify*nt=NULL)
 498         {
 499                 verbose ? print_verbose(nt) : print_simple();
 500         }
 501         string get_print_dev_string()
 502         {
 503                 string str;
 504                 str = m_dev;
 505                 str.resize(8, ' ');
 506
 507                 string_ex s;
 508                 s.format("%2d/%2d", m_cmd_index+1, m_cmd_total);
 509
 510                 str += s;
 511                 return str;
 512         }
 513         void print_simple()
 514         {
 515                 int width = get_console_width();
 516                 int info, bar;
 517                 info = 14;
 518                 bar = 40;
 519
 520                 if (m_IsEmptyLine)
 521                 {
 522                         string str(width, ' ');
 523                         cout << str;
 524                         return;
 525                 }
 526                 if (width <= bar + info + 3)
 527                 {
 528                         string_ex str;
 529
 530                         str += get_print_dev_string();
 531
 532                         str += g_wait[(g_wait_index++) & 0x3];
 533
 534                         print_oneline(str);
 535                         return ;
 536                 }
 537                 else
 538                 {
 539                         string_ex str;
 540                         str += get_print_dev_string();
 541
 542                         str.resize(info, ' ');
 543                         cout << str;
 544
 545                         if (m_done || m_status)
 546                         {
 547                                 string str;
 548                                 str.resize(bar, ' ');
 549                                 str[0] = '[';
 550                                 str[str.size() - 1] = ']';
 551                                 string err;
 552                                 if (m_status)
 553                                 {
 554                                         err = uuu_get_last_err_string();
 555                                         err.resize(bar - 2, ' ');
 556                                         str.replace(1, err.size(), err);
 557                                         str.insert(1, g_vt_red);
 558                                         str.insert(1 + strlen(g_vt_red) + err.size(), g_vt_default);
 559                                 }
 560                                 else
 561                                 {
 562                                         str.replace(1, 4, "Done");
 563                                         str.insert(1, g_vt_green);
 564                                         str.insert(1 + strlen(g_vt_green) + strlen("Done"), g_vt_default);
 565                                 }
 566                                 cout << str;
 567                         } else {
 568                                 cout << build_process_bar(bar, m_trans_pos, m_trans_size);
 569                         }
 570                         cout << " ";
 571                         print_auto_scroll(m_cmd, width - bar - info-1, m_start_pos);
 572
 573 if (clock() - m_start_time > CLOCKS_PER_SEC / 4)
 574 {
 575         m_start_pos++;
 576         m_start_time = clock();
 577 }
 578 cout << endl;
 579
 580 return;
 581                 }
 582         }
 583 };
 584
 585 static map<string, ShowNotify> g_map_path_nt;
 586 mutex g_callback_mutex;
 587
 588 void print_oneline(string str)
 589 {
 590         size_t w = get_console_width();
 591         if (w <= 3)
 592                 return;
 593
 594         if (str.size() >= w)
 595         {
 596                 str.resize(w - 1);
 597                 str[str.size() - 1] = '.';
 598                 str[str.size() - 2] = '.';
 599                 str[str.size() - 3] = '.';
 600         }
 601         else
 602         {
 603                 str.resize(w, ' ');
 604         }
 605         cout << str << endl;
 606
 607 }
 608
 609 ShowNotify Summary(map<uint64_t, ShowNotify> *np)
 610 {
 611         ShowNotify sn;
 612         for (auto it = np->begin(); it != np->end(); it++)
 613         {
 614                 if (it->second.m_dev == "Prep")
 615                 {
 616                         sn.m_trans_size += it->second.m_trans_size;
 617                         sn.m_trans_pos += it->second.m_trans_pos;
 618                 }
 619                 else
 620                 {
 621                         if (it->second.m_trans_pos || it->second.m_cmd_index)
 622                                 g_start_usb_transfer = true; // Hidden HTTP download when USB start transfer
 623                 }
 624         }
 625
 626         if(g_start_usb_transfer)
 627                 sn.m_IsEmptyLine = true; // Hidden HTTP download when USB start transfer
 628
 629         sn.m_dev = "Prep";
 630         sn.m_cmd = "Http Download\\Uncompress";
 631         return sn;
 632 }
 633
 634 int progress(uuu_notify nt, void *p)
 635 {
 636         map<uint64_t, ShowNotify> *np = (map<uint64_t, ShowNotify>*)p;
 637         map<string, ShowNotify>::iterator it;
 638
 639         std::lock_guard<std::mutex> lock(g_callback_mutex);
 640
 641         if ((*np)[nt.id].update(nt))
 642         {
 643                 if (!(*np)[nt.id].m_dev.empty())
 644                         if ((*np)[nt.id].m_dev != "Prep")
 645                                 g_map_path_nt[(*np)[nt.id].m_dev] = (*np)[nt.id];
 646
 647                 if (g_verbose)
 648                 {
 649                         if((*np)[nt.id].m_dev == "Prep")
 650                                 Summary(np).print(g_verbose, &nt);
 651                         else
 652                                 (*np)[nt.id].print(g_verbose, &nt);
 653                 }
 654                 else
 655                 {
 656                         string_ex str;
 657                         str.format("\rSuccess %d    Failure %d    ", g_overall_okay, g_overall_failure);
 658
 659                         if (g_map_path_nt.empty())
 660                                 str += "Wait for Known USB Device Appear...";
 661
 662                         if (!g_usb_path_filter.empty())
 663                         {
 664                                 str += " at path ";
 665                                 for (size_t i = 0; i < g_usb_path_filter.size(); i++)
 666                                         str += g_usb_path_filter[i] + " ";
 667                         }
 668
 669                         print_oneline(str);
 670                         print_oneline("");
 671                         if ((*np)[nt.id].m_dev == "Prep" && !g_start_usb_transfer)
 672                         {
 673                                 Summary(np).print();
 674                         }else
 675                                 print_oneline("");
 676
 677                         for (it = g_map_path_nt.begin(); it != g_map_path_nt.end(); it++)
 678                                 it->second.print();
 679
 680                         for (size_t i = 0; i < g_map_path_nt.size() + 3; i++)
 681                                 cout << "\x1B[1F";
 682
 683                 }
 684
 685                 //(*np)[nt.id] = g_map_path_nt[(*np)[nt.id].m_dev];
 686         }
 687
 688         if (nt.type == uuu_notify::NOTIFY_THREAD_EXIT)
 689         {
 690                 if(np->find(nt.id) != np->end())
 691                         np->erase(nt.id);
 692         }
 693         return 0;
 694 }
 695 #ifdef _MSC_VER
 696
 697 #define DEFINE_CONSOLEV2_PROPERTIES
 698 #include <windows.h>
 699 #include <stdlib.h>
 700 #include <stdio.h>
 701
 702 bool enable_vt_mode()
 703 {
 704         // Set output mode to handle virtual terminal sequences
 705         HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);
 706         if (hOut == INVALID_HANDLE_VALUE)
 707         {
 708                 clean_vt_color();
 709                 return false;
 710         }
 711
 712         DWORD dwMode = 0;
 713         if (!GetConsoleMode(hOut, &dwMode))
 714         {
 715                 clean_vt_color();
 716                 return false;
 717         }
 718
 719         dwMode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
 720         if (!SetConsoleMode(hOut, dwMode))
 721         {
 722                 clean_vt_color();
 723                 return false;
 724         }
 725         return true;
 726 }
 727
 728 int get_console_width()
 729 {
 730         CONSOLE_SCREEN_BUFFER_INFO sbInfo;
 731         GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &sbInfo);
 732         return sbInfo.dwSize.X;
 733 }
 734 #else
 735 #include <sys/ioctl.h>
 736 bool enable_vt_mode() { return true; }
 737 int get_console_width()
 738 {
 739         struct winsize w;
 740         ioctl(0, TIOCGWINSZ, &w);
 741         return w.ws_col;
 742 }
 743 #endif
 744 void print_usb_filter()
 745 {
 746         if (!g_usb_path_filter.empty())
 747         {
 748                 cout << " at path ";
 749                 for (size_t i = 0; i < g_usb_path_filter.size(); i++)
 750                         cout << g_usb_path_filter[i] << " ";
 751         }
 752 }
 753
 754 int runshell(int shell)
 755 {
 756         int uboot_cmd = 0;
 757         string prompt = "U>";
 758
 759         if (shell)
 760         {
 761                 cout << "Please input command: " << endl;
 762                 string cmd;
 763                 ofstream log("uuu.inputlog", ofstream::binary);
 764                 log << "uuu_version "
 765                         << ((uuu_get_version() & 0xFF000000) >> 24)
 766                         << "."
 767                         << ((uuu_get_version() & 0xFFF000) >> 12)
 768                         << "."
 769                         << ((uuu_get_version() & 0xFFF))
 770                         << endl;
 771                 while (1)
 772                 {
 773                         cout << prompt;
 774                         getline(cin, cmd);
 775
 776                         if (cmd == "uboot")
 777                         {
 778                                 uboot_cmd = 1;
 779                                 prompt = "=>";
 780                                 cout << "Enter into u-boot cmd mode" << endl;
 781                                 cout << "Okay" << endl;
 782                         }
 783                         else if (cmd == "exit" && uboot_cmd == 1)
 784                         {
 785                                 uboot_cmd = 0;
 786                                 prompt = "U>";
 787                                 cout << "Exit u-boot cmd mode" << endl;
 788                                 cout << "Okay" << endl;
 789                         }else if (cmd == "help" || cmd == "?")
 790                         {
 791                                 print_help();
 792                         }
 793                         else if (cmd == "q" || cmd == "quit")
 794                         {
 795                                 return 0;
 796                         }
 797                         else
 798                         {
 799                                 log << cmd << endl;
 800                                 log.flush();
 801
 802                                 if (uboot_cmd)
 803                                         cmd = "fb: ucmd " + cmd;
 804
 805                                 int ret = uuu_run_cmd(cmd.c_str(), 0);
 806                                 if (ret)
 807                                         cout << uuu_get_last_err_string() << endl;
 808                                 else
 809                                         cout << "Okay" << endl;
 810                         }
 811                 }
 812                 return 0;
 813         }
 814
 815         return -1;
 816 }
 817
 818 void print_udev()
 819 {
 820         uuu_for_each_cfg(print_udev_rule, NULL);
 821         fprintf(stderr, "\n1: put above udev run into /etc/udev/rules.d/70-uuu.rules\n");
 822         fprintf(stderr, "\tsudo sh -c \"uuu -udev >> /etc/udev/rules.d/70-uuu.rules\"\n");
 823         fprintf(stderr, "2: update udev rule\n");
 824         fprintf(stderr, "\tsudo udevadm control --reload\n");
 825 }
 826
 827 int print_usb_device(const char *path, const char *chip, const char *pro, uint16_t vid, uint16_t pid, uint16_t bcd, void * /*p*/)
 828 {
 829         printf("\t%s\t %s\t %s\t 0x%04X\t0x%04X\t 0x%04X\n", path, chip, pro, vid, pid, bcd);
 830         return 0;
 831 }
 832
 833 void print_lsusb()
 834 {
 835         cout << "Connected Known USB Devices\n";
 836         printf("\tPath\t Chip\t Pro\t Vid\t Pid\t BcdVersion\n");
 837         printf("\t==================================================\n");
 838
 839         uuu_for_each_devices(print_usb_device, NULL);
 840 }
 841
 842 #ifdef WIN32
 843
 844 int ignore_serial_number(const char *pro, const char *chip, const char */*comp*/, uint16_t vid, uint16_t pid, uint16_t /*bcdlow*/, uint16_t /*bcdhigh*/, void */*p*/)
 845 {
 846         printf("\t %s\t %s\t 0x%04X\t0x%04X\n", chip, pro, vid, pid);
 847
 848         char sub[128];
 849         snprintf(sub, 128, "IgnoreHWSerNum%04x%04x", vid, pid);
 850         const BYTE value = 1;
 851
 852         LSTATUS ret = RegSetKeyValueA(HKEY_LOCAL_MACHINE,
 853                                                                   "SYSTEM\\CurrentControlSet\\Control\\UsbFlags",
 854                                                                         sub, REG_BINARY, &value, 1);
 855         if(ret == ERROR_SUCCESS)
 856                 return 0;
 857
 858         printf("Set key failure, try run as administrator permission\n");
 859         return -1;
 860 }
 861 #endif
 862
 863 int set_ignore_serial_number()
 864 {
 865 #ifndef WIN32
 866         printf("Only windows system need set ignore serial number registry");
 867         return -1;
 868 #else
 869         printf("Set window registry to ignore usb hardware serial number for known uuu device:\n");
 870         return uuu_for_each_cfg(ignore_serial_number, NULL);
 871 #endif
 872 }
 873
 874 int main(int argc, char **argv)
 875 {
 876         if (auto_complete(argc, argv) == 0)
 877                 return 0;
 878
 879         if (argc >= 2)
 880         {
 881                 string s = argv[1];
 882                 if(s == "-udev")
 883                 {
 884                         print_udev();
 885                         return 0;
 886                 }
 887                 if (s == "-bshow")
 888                 {
 889                         if (2 == argc || g_BuildScripts.find(argv[2]) == g_BuildScripts.end())
 890                         {
 891                                 fprintf(stderr, "Error, must be have script name: ");
 892                                 g_BuildScripts.ShowCmds(stderr);
 893                                 fprintf(stderr,"\n");
 894                                 return -1;
 895                         }
 896                         else
 897                         {
 898                                 string str = g_BuildScripts[argv[2]].m_text;
 899                                 while (str.size() > 0 && (str[0] == '\n' || str[0] == ' '))
 900                                         str = str.erase(0,1);
 901
 902                                 printf("%s", str.c_str());
 903                                 return 0;
 904                         }
 905                 }
 906         }
 907
 908         AutoCursor a;
 909
 910         print_version();
 911
 912         if (!enable_vt_mode())
 913         {
 914                 cout << "Your console don't support VT mode, fail back to verbose mode" << endl;
 915                 g_verbose = 1;
 916         }
 917
 918         if (argc == 1)
 919         {
 920                 print_help();
 921                 return 0;
 922         }
 923
 924         int deamon = 0;
 925         int shell = 0;
 926         string filename;
 927         string cmd;
 928         int ret;
 929         int dryrun  = 0;
 930
 931         string cmd_script;
 932
 933         for (int i = 1; i < argc; i++)
 934         {
 935                 string s = argv[i];
 936                 if (!s.empty() && s[0] == '-')
 937                 {
 938                         if (s == "-d")
 939                         {
 940                                 deamon = 1;
 941                                 uuu_set_small_mem(0);
 942
 943                         }
 944                         else if (s == "-dm")
 945                         {
 946                                 uuu_set_small_mem(0);
 947                         }
 948                         else if (s == "-s")
 949                         {
 950                                 shell = 1;
 951                                 g_verbose = 1;
 952                         }
 953                         else if (s == "-v")
 954                         {
 955                                 g_verbose = 1;
 956                         }
 957                         else if (s == "-V")
 958                         {
 959                                 g_verbose = 1;
 960                                 uuu_set_debug_level(2);
 961                         }else if (s == "-dry")
 962                         {
 963                                 dryrun = 1;
 964                                 g_verbose = 1;
 965                         }
 966                         else if (s == "-h")
 967                         {
 968                                 print_help(false);
 969                                 return 0;
 970                         }
 971                         else if (s == "-H")
 972                         {
 973                                 print_help(true);
 974                                 return 0;
 975                         }
 976                         else if (s == "-m")
 977                         {
 978                                 i++;
 979                                 uuu_add_usbpath_filter(argv[i]);
 980                                 g_usb_path_filter.push_back(argv[i]);
 981                         }
 982                         else if (s == "-t")
 983                         {
 984                                 i++;
 985                                 uuu_set_wait_timeout(atoll(argv[i]));
 986                         }
 987                         else if (s == "-T")
 988                         {
 989                                 i++;
 990                                 uuu_set_wait_next_timeout(atoll(argv[i]));
 991                         }
 992                         else if (s == "-pp")
 993                         {
 994                                 i++;
 995                                 uuu_set_poll_period(atoll(argv[i]));
 996                         }
 997                         else if (s == "-lsusb")
 998                         {
 999                                 print_lsusb();
1000                                 return 0;
1001                         }
1002                         else if (s == "-IgSerNum")
1003                         {
1004                                 return set_ignore_serial_number();
1005                         }
1006                         else if (s == "-bmap")
1007                         {
1008                                 g_bmap_mode = bmap_mode::Force;
1009                         }
1010                         else if (s == "-no-bmap")
1011                         {
1012                                 g_bmap_mode = bmap_mode::Ignore;
1013                         }
1014                         else if (s == "-e")
1015                         {
1016 #ifndef WIN32
1017         #define _putenv putenv
1018 #endif
1019                                 i++;
1020                                 if (_putenv(argv[i]))
1021                                 {
1022                                         printf("error, failed to set '%s', environment parameter must have the form key=value\n", argv[i]);
1023                                         return -1;
1024                                 }
1025                         }
1026                         else if (s == "-b" || s == "-brun")
1027                         {
1028                                 if (i + 1 == argc)
1029                                 {
1030                                         printf("error, must be have script name: ");
1031                                         g_BuildScripts.ShowCmds();
1032                                         printf("\n");
1033                                         return -1;
1034                                 }
1035
1036                                 vector<string> args;
1037                                 for (int j = i + 2; j < argc; j++)
1038                                 {
1039                                         string s = argv[j];
1040                                         if (s.find(' ') != string::npos)
1041                                         {
1042                                                 s.insert(s.begin(), '"');
1043                                                 s.insert(s.end(), '"');
1044                                         }
1045                                         args.push_back(s);
1046                                 }
1047
1048                                 // if script name is not build-in, try to look for a file
1049                                 if (g_BuildScripts.find(argv[i + 1]) == g_BuildScripts.end()) {
1050                                         const string tmpCmdFileName{argv[i + 1]};
1051
1052                                         std::ifstream t(tmpCmdFileName);
1053                                         std::string fileContents((std::istreambuf_iterator<char>(t)),
1054                                                 std::istreambuf_iterator<char>());
1055
1056                                         if (fileContents.empty()) {
1057                                                 printf("%s is not built-in script or fail load external script file", tmpCmdFileName.c_str());
1058                                                 return -1;
1059                                         }
1060
1061                                         BuiltInScriptRawData tmpCmd{
1062                                                 tmpCmdFileName.c_str(),
1063                                                 fileContents.c_str(),
1064                                                 "Script loaded from file"
1065                                         };
1066
1067                                         g_BuildScripts.emplace(tmpCmdFileName, &tmpCmd);
1068
1069                                         cmd_script = g_BuildScripts[tmpCmdFileName].replace_script_args(args);
1070                                 }
1071                                 else {
1072                                         cmd_script = g_BuildScripts[argv[i + 1]].replace_script_args(args);
1073                                 }
1074                                 break;
1075                         }
1076                         else
1077                         {
1078                                 cout << "Unknown option: " << s.c_str();
1079                                 return -1;
1080                         }
1081                 }else if (!s.empty() && s[s.size() - 1] == ':')
1082                 {
1083                         for (int j = i; j < argc; j++)
1084                         {
1085                                 s = argv[j];
1086                                 if (s.find(' ') != string::npos && s[s.size() - 1] != ':')
1087                                 {
1088                                         s.insert(s.begin(), '"');
1089                                         s.insert(s.end(), '"');
1090                                 }
1091                                 cmd.append(s);
1092                                 if(j != (argc -1)) /* Don't add space at last arg */
1093                                         cmd.append(" ");
1094                         }
1095                         break;
1096                 }
1097                 else
1098                 {
1099                         filename = s;
1100                         break;
1101                 }
1102         }
1103
1104         signal(SIGINT, ctrl_c_handle);
1105
1106         uuu_set_askpasswd(ask_passwd);
1107
1108         if (deamon && shell)
1109         {
1110                 printf("Error: -d -s Can't apply at the same time\n");
1111                 return -1;
1112         }
1113
1114         if (deamon && dryrun)
1115         {
1116                 printf("Error: -d -dry Can't apply at the same time\n");
1117                 return -1;
1118         }
1119
1120         if (shell && dryrun)
1121         {
1122                 printf("Error: -dry -s Can't apply at the same time\n");
1123                 return -1;
1124         }
1125
1126         if (g_verbose)
1127         {
1128                 printf("%sBuild in config:%s\n", g_vt_boldwhite, g_vt_default);
1129                 printf("\tPctl\t Chip\t\t Vid\t Pid\t BcdVersion\n");
1130                 printf("\t==================================================\n");
1131                 uuu_for_each_cfg(print_cfg, NULL);
1132
1133                 if (!cmd_script.empty())
1134                         printf("\n%sRun built-in script:%s\n %s\n\n", g_vt_boldwhite, g_vt_default, cmd_script.c_str());
1135
1136                 if (!shell)
1137                         cout << "Wait for Known USB Device Appear...";
1138
1139                 print_usb_filter();
1140
1141                 printf("\n");
1142         }
1143         else {
1144                 cout << "Wait for Known USB Device Appear...";
1145                 print_usb_filter();
1146                 cout << "\r";
1147                 cout << "\x1b[?25l";
1148                 cout.flush();
1149         }
1150
1151         map<uint64_t, ShowNotify> nt_session;
1152
1153         uuu_register_notify_callback(progress, &nt_session);
1154
1155
1156         if (!cmd.empty())
1157         {
1158                 ret = uuu_run_cmd(cmd.c_str(), dryrun);
1159
1160                 for (size_t i = 0; i < g_map_path_nt.size()+3; i++)
1161                         printf("\n");
1162                 if(ret)
1163                         printf("\nError: %s\n", uuu_get_last_err_string());
1164                 else
1165                         printf("Okay\n");
1166
1167                 runshell(shell);
1168                 return ret;
1169         }
1170
1171         if (!cmd_script.empty())
1172                 ret = uuu_run_cmd_script(cmd_script.c_str(), dryrun);
1173         else
1174                 ret = uuu_auto_detect_file(filename.c_str());
1175
1176         if (ret)
1177         {
1178                 ret = runshell(shell);
1179                 if(ret)
1180                         cout << g_vt_red << "\nError: " << g_vt_default <<  uuu_get_last_err_string();
1181                 return ret;
1182         }
1183
1184         if (uuu_wait_uuu_finish(deamon, dryrun))
1185         {
1186                 cout << g_vt_red << "\nError: " << g_vt_default << uuu_get_last_err_string();
1187                 return -1;
1188         }
1189
1190         runshell(shell);
1191
1192         /*Wait for the other thread exit, after send out CMD_DONE*/
1193         std::this_thread::sleep_for(std::chrono::milliseconds(100));
1194         if(!g_verbose)
1195                 printf("\n\n\n");
1196         return g_overall_status;
1197 }