class DemoClass
{
  uint                Index;
  ConnectionClass     Conn;
  String              Output_Folder,
                      Input_Folder;

  DemoClass(ConnectionClass Conn, const String &input_folder, const String &output_folder) :
    Conn = Conn,
    Input_Folder = input_folder,
    Output_Folder = output_folder
  {

  }

  void Recurse()
  {
    Recurse(Input_Folder);
  }

  void Recurse(const String &path)
  {
    DirRootClass droot(Conn);
    DirectoryClass root();
    if (droot.GetListing(path, root))
    {
      foreach (DirectoryClass child in root)
      {
        String child_path = String::Format("{0}\\{1}", path, child.Name());
        if (child.Extension().Compare("evtx") == 0)
        {
          if (child.LogicalSize())
          {
            Console.WriteLine("Processing '{0}'", child_path);
            String error_message;
            if (Extract(child.Name(), child_path, Conn, error_message))
            {
              Console.WriteLine("Successfully extracted '{0}' to '{1}'.", child.Name(), error_message);
            }
            else
            {
              Console.WriteLine("Unable to extract '{0}'.", child_path);
            }
          }
        }
        else if (child.IsFolder())
        {
          Recurse(child_path);
        }
      }
    }
  }

  bool Extract(const String &name, const String &path, ConnectionClass Conn, String &error_message)
  {
    bool retval;
    HostFileClass input(Conn);
    if (input.Open(path))
    {
      String output_path = String::Format("{0}\\{1} - {2}", Output_Folder, String::FormatInt(Index++, int::DECIMAL, 0, 6), name);
      LocalFileClass output();
      if (output.Open(output_path, WRITE))
      {
        output.WriteBuffer(input);
        error_message = output_path;
        retval = true;
      }
      else
      {
        error_message = String::Format("Can't open output_path at {0}.", output_path);
      }
    }
    else
    {
      error_message = String::Format("Can't open file-data for '{0}'.", path);
    }
    return retval;
  }
}

class MainClass {

  static const String Input_Folder = "C:\\Windows\\System32\\winevt\\Logs";

  SafeClass Safe;
  RoleClass RoleRoot,
            Role;
  NetworkClass SweepNet;
  String NetText,
         ClientReturnAddress,
         StatusBarName,
         OutputFolder;
  int NumConnections,
     ConnectOptions;
  uint HostIndex;


  MainClass() :
    Safe(),
    RoleRoot(),
    Role(),
    SweepNet(),
    NumConnections = 1,
    ConnectOptions = ConnectionClass::CLIENTNODESAFE,
    StatusBarName = "Enterprise Example - Using Snapshot Data"
  {
  }

  /**
  Entry point of the Enscript
  **/
  void Main() {
    SystemClass::ClearConsole();
    if (SystemClass::FolderDialog(OutputFolder, "Choose base output folder"))
    {
    if (Safe.Logon(null) && ShowDiag() == SystemClass::OK) {
        Sweep();
        SystemClass::Message(0, "Success", String::Format("{0}: Completed Successfully!", StatusBarName));
      }
    }
  }

  /**
  This method contains the logic we want to apply to each node on the network
  **/
  void Process(ConnectionClass conn, SnapshotClass snap) {
    Console.WriteLine("Connected To Servlet On {0} Snapshot State = {1}", snap.Name(), SnapshotClass::States::SourceText(snap.State()));
    String host_output_folder = String::Format("{0}\\{1} - {2}", OutputFolder, String::FormatInt(HostIndex++, int::DECIMAL, 0, 6), snap.Name());
    if (LocalMachine.CreateFolder(host_output_folder))
    {
      DemoClass demo(conn, Input_Folder, host_output_folder);
      demo.Recurse();
    }
  }

  /**
   Display dialogs
  **/
  int ShowDiag() {
    RoleRoot = Safe.RoleRoot();
    DialogClass diag();
    new NetTextDialogClass(diag, this);
    return diag.Wizard();
  }

  /**
    Code that gets connection and snapshot
  **/
  void ReadNetwork(BatchClass batch, SnapshotClass root) {
    String message,
           name;
    DateClass d();
    do {
      ConnectionClass conn;
      SnapshotClass ss(null);
      message = "";
      BatchClass::ConnectionTypes reply = batch.GetConnection(conn, ss, name, message, 0);
      if (reply == BatchClass::BATCHCONNECT) { //successfully connected to remote node
        Process(conn, ss);
        SystemClass::StatusInc(1);
        root.Insert(ss);
      }
      else if (reply == BatchClass::BATCHERROR) { //could not connect to remote node. ss object will have the state of the node
        d.Now();
        Console.WriteLine("Could Not Connect To {0} SAFE Error Message: ", name, message);
        SystemClass::StatusInc(1);
        root.Insert(ss);
      }
      else if (reply == BatchClass::BATCHWAIT)
        SystemClass::Sleep(100);
      else if (reply == BatchClass::BATCHFATAL) {
        String err = SystemClass::LastError();
        Console.WriteLine("The SAFE is not responding: {0}. This Enscript will terminate.", err);
        return;
      }
    } while (reply != BatchClass::BATCHDONE);
  }

  /** Code that creates a batchclass
  **/
  void Sweep() {
    DateClass now;
    SnapshotClass newSnaps = new SnapshotClass(null, "Snapshot");
    BatchClass batch(Safe, Role, NumConnections, ConnectionClass::SNAPALL);
    if (batch.Add(SweepNet)) {
      batch.SetMode(ConnectionClass::Options::Convert(ConnectOptions), ClientReturnAddress);
      if (batch.Start()) {
        uint machines = batch.TotalMachines();
        Console.WriteLine("Scanning {0} using {1}", Plural("node", machines), Plural("connection", batch.ConnectionsUsed()));
        SystemClass::StatusRange(StatusBarName, machines);
        uint start;
        now.Now();
        start = now.GetUnix();
        ReadNetwork(batch, newSnaps);
        now.Now();
        Console.WriteLine("Scan completed in {0} seconds", (now.GetUnix() - start));
      }
      else {
        SystemClass::Message(0, "BatchClass error", SystemClass::LastError());
      }
    }
    else {
      SystemClass::Message(0, "BatchClass Error", "Unable to add any IPs to the sweep");
    }
  }

  String Plural(const String &str, uint n) {
    return String::Format("{0} {1}{2}", n, str, n == 1 ? "" : "s");
  }

  /**
   Turn a string of text into networkclass objects
  **/
  bool ParseText(String t) {
    SweepNet.Close();
    bool ret = false;
    while (t) {
      ret = true;
      int    end  = t.Find("\n");
      String line = end < 0 ? t : t.SubString(0, end);
      int    dash = line.Find("-");
      if (dash >= 0) {
        IPClass ip1(ExtractIP(line.SubString(0, dash))),
                ip2(ExtractIP(line.SubString(dash+1, -1)));
        if (ip1 && ip2) {
          NetworkClass n(SweepNet, "IP Range", NodeClass::SELECTED);
          n.SetStart(ip1);
          n.SetStop(ip2);
        }
        else
          NetworkClass n(SweepNet, line, NodeClass::SELECTED);
      }
      else if (line != "")  {
        NetworkClass n(SweepNet, line, NodeClass::SELECTED);
      }
      if (end >= 0)
        t.Delete(0, end+1);
      else
        break;
    }
    return ret;
  }

  /**
   Check for IPs in nettext
  **/
  String ExtractIP(const String &s) {
    String ret = s;
    ret.Trim(" ", String::TRIMSTART | String::TRIMEND);
    return ret.IsValidIPAddress() ? ret : "";
  }
}

/**
 Dialog to choose a role and enter nodes to sweep
**/
class NetTextDialogClass: DialogClass {

  MainClass Data;
  StaticTextClass SafeTextEdit;
  TreeEditClass Tree;
  StaticTextClass Help;
  StringEditClass NetTextEdit;

  NetTextDialogClass(DialogClass diag, MainClass d) :
    DialogClass(diag, String::Format("{0} Options", d.StatusBarName)),
    Data = d,
    SafeTextEdit(this, "", START, 15, 200, 100, 0),
    Tree(this, "Choose The Role You Want To Assume", NEXT, START, 200, 100, 0, d.RoleRoot, 0),
    Help(this, "Enter IP addresses or machine names on separate\n"
                 "lines. Enter ranges on separate lines and delimit\n"
                 "the start and stop address with a dash (\"-\").\n\n"
                 "Example:\n\n"
                 "\tlocalhost\n"
                 "\t192.168.5.5\n"
                 "\t192.168.0.16-192.168.0.64\n"
                 "\t192.168.1.1-192.168.3.255\n"
                 "\tfd00:0:1000:20:0:0:0:100\n",
                 START, NEXT, 200, 100, REQUIRED),
    NetTextEdit(this, "", NEXT, SAME, 200, 100, AUTOVSCROLL | MULTILINE | WANTRETURN, d.NetText, 9999, 0)
  {

  }

  virtual void Setup() {
    DialogClass::Setup();
    SafeTextEdit.SetText("SAFE:\t\t\t\t" + Data.Safe.Name() +
                         "\nUser:\t\t\t\t" + Data.Safe.UserName() +
                          "\n\nTotal Connections:\t\t" + Data.Safe.TotalConnections() +
                          "\nActive Connections:\t\t" + Data.Safe.ActiveConnections() +
                          "\nConnections To Use:\t\t" + Data.NumConnections +
                          "\n\nRemediation Allowed:\t\t" + (Data.Safe.RemediationAllowed() ? "Yes" : "No") +
                          "\nSnapshot Allowed:\t\t" + (Data.Safe.SnapshotAllowed() ? "Yes" : "No") +
                          "\n\nSAFE Version:\t\t\t" + Data.Safe.Version()
                          );
  }

  virtual void CheckControls() {
    DialogClass::CheckControls();
    EnableClose(Tree.GetValue().Parent());
  }

  virtual bool CanClose() {
    Output();
    bool ret = false;
    if (DialogClass::CanClose()) {
      Data.Role = RoleClass::TypeCast(Tree.GetValue());
      ret = Data.ParseText(Data.NetText);
      if (!ret)
        ErrorMessage("Please Enter a value in the IP List Text Area.");
    }
    return ret;
  }
}