<commit_before>#!/usr/bin/python
"""
Copyright 2016 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import gzip
import json
import logging
import math
import os
import struct
import time

#Globals
options = None

########################################################################################################################
#   Pcap processing
########################################################################################################################
class Pcap():
  def __init__(self):
    self.start_seconds = None
    self.start_time = None
    self.local_ethernet_mac = None
    self.slices = {'in': [], 'out': [], 'in_dup': []}
    self.bytes = {'in': 0, 'out': 0, 'in_dup': 0}
    self.streams = {}
    return

  def SaveStats(self, out):
    file_name, ext = os.path.splitext(out)
    if ext.lower() == '.gz':
      f = gzip.open(out, 'wb')
    else:
      f = open(out, 'wb')
    try:
      result = {"bytes": self.bytes}
      json.dump(result, f)
      logging.info('Result stats written to {0}'.format(out))
    except:
      logging.critical('Error writing result stats to {0}'.format(out))
    f.close()

  def SaveDetails(self, out):
    file_name, ext = os.path.splitext(out)
    if ext.lower() == '.gz':
      f = gzip.open(out, 'wb')
    else:
      f = open(out, 'wb')
    try:
      json.dump(self.slices, f)
      logging.info('Result details written to {0}'.format(out))
    except:
      logging.critical('Error writing result details to {0}'.format(out))
    f.close()

  def Print(self):
    global options
    if options.json:
      print(json.dumps(self.bytes, indent=2))
    else:
      print "Bytes Out: {0:d}".format(self.bytes['out'])
      print "Bytes In:  {0:d}".format(self.bytes['in'])
      print "Duplicate Bytes In:  {0:d}".format(self.bytes['in_dup'])

  def Process(self, pcap):
    f = None
    self.__init__() #Reset state if called multiple times
    try:
      file_name, ext = os.path.splitext(pcap)
      if ext.lower() == '.gz':
        f = gzip.open(pcap, 'rb')
      else:
        f = open(pcap, 'rb')
      bytes = f.read(24)
      # File header:
      # Magic Number - 4 bytes - 0xa1b2c3d4
      # Major Version - 2 bytes
      # Minor version - 2 bytes
      # Tz offset - 4 bytes (always 0)
      # Timestamp accuracy - 4 bytes (always 0)
      # Snapshot length - 4 bytes
      # Link layer header type - 4 bytes
      #
      # unpack constants:
      # L - unsigned long (4 byte)
      # H - unsigned short (2 byte)
      # B - unsigned char (1 byte int)
      file_header = struct.unpack("=LHHLLLL", bytes)

      # ignore byte order reversals for now
      if file_header[0] == 0xa1b2c3d4:
        ok = True
        self.linktype = file_header[6]
        if self.linktype == 1:
          self.linklen = 12
        elif self.linktype == 113:
          self.linklen = 14
        else:
          logging.critical("Unknown link layer header type: {0:d}".format(self.linktype))
          ok = False

        # Packet header:
        # Time stamp (seconds) - 4 bytes
        # Time stamp (microseconds value) - 4 bytes
        # Captured data length - 4 bytes
        # Original length - 4 bytes
        while ok:
          bytes = f.read(16)
          if not bytes or len(bytes) < 16:
            break
          (seconds, useconds, captured_length, packet_length) = struct.unpack("=LLLL", bytes)
          if self.start_seconds is None:
            self.start_seconds = seconds
          seconds -= self.start_seconds
          if packet_length and captured_length <= packet_length:
            packet_time = float(seconds) + float(useconds) / 1000000.0
            packet_info = {}
            packet_info['time'] = packet_time
            packet_info['length'] = packet_length
            packet_info['captured_length'] = captured_length
            packet_info['valid'] = False
            if captured_length:
              packet_data = f.read(captured_length)
            else:
              packet_data = None
            if len(packet_data) >= self.linklen:
              try:
                self.ProcessPacket(packet_data, packet_info)
              except Exception as e:
                print(e)
      else:
        logging.critical("Invalid pcap file " + pcap)
    except:
      logging.critical("Error processing pcap " + pcap)

    if f is not None:
      f.close()

    return

  def ProcessPacket(self, packet_data, packet_info):
    if self.linktype == 1:
      # Ethernet:
      # dst1: 2 bytes
      # dst2: 2 bytes
      # dst3: 2 bytes
      # src1: 2 bytes
      # src2: 2 bytes
      # src3: 2 bytes
      ethernet_header = struct.unpack("!HHHHHH", packet_data[0:self.linklen])
      # Ignore broadcast traffic
      packet_info['ethernet_dst'] = [ethernet_header[0], ethernet_header[1], ethernet_header[2]]
      packet_info['ethernet_src'] = [ethernet_header[3], ethernet_header[4], ethernet_header[5]]
      dst = packet_info['ethernet_dst']
      if dst[0] != 0xFFFF or dst[1] != 0xFFFF or  dst[2] != 0xFFFF:
        packet_info['valid'] = True

    elif self.linktype == 113:
      # Linux cooked capture
      # Packet Type: 2 bytes
      # aprhrd type: 2 bytes
      # Address length: 2 bytes
      # Address part 1: 4 bytes
      # Address part 2: 4 bytes
      cooked_header = struct.unpack("!HHHLL", packet_data[0:self.linklen])
      if cooked_header[0] == 0:
        packet_info['valid'] = True
        packet_info['direction'] = 'in'
      if cooked_header[0] == 4:
        packet_info['valid'] = True
        packet_info['direction'] = 'out'

    protocol = struct.unpack("!H", packet_data[self.linklen:self.linklen + 2])[0]
    if packet_info['valid'] and protocol == 0x800: # Only handle IPv4 for now
      self.ProcessIPv4Packet(packet_data[self.linklen + 2:], packet_info)

    if packet_info['valid'] and self.start_time:
      if self.local_ethernet_mac is not None:
        local = self.local_ethernet_mac
        src = packet_info['ethernet_src']
        if src[0] == local[0] and src[1] == local[1] and src[2] == local[2]:
          packet_info['direction'] = 'out'
        else:
          packet_info['direction'] = 'in'
      if 'direction' in packet_info:
        self.ProcessPacketInfo(packet_info)

    return

  def ProcessIPv4Packet(self, ip_packet, packet_info):
    # IP Header:
    # Version/len: 1 Byte (4 bits each)
    # dscp/ecn: 1 Byte
    # Total Length: 2 Bytes
    # Identification: 2 Bytes
    # Flags/Fragment: 2 Bytes
    # TTL: 1 Byte
    # Protocol: 1 Byte
    # Header Checksum: 2 Bytes
    # Source Address: 4 Bytes
    # Dest Address: 4 Bytes
    if len(ip_packet) > 20:
      ip_header = struct.unpack("!BBHHHBBHLL", ip_packet[0:20])
      header_length = (ip_header[0] & 0x0F) * 4
      total_length = ip_header[2]
      payload_length = total_length - header_length
      packet_info['ip_payload_length'] = payload_length
      packet_info['ip_protocol'] = ip_header[6]
      packet_info['ip_src'] = ip_header[8]
      addr = struct.unpack("BBBB", ip_packet[12:16])
      packet_info['ip_src_str'] = '{0:d}.{1:d}.{2:d}.{3:d}'.format(addr[0], addr[1], addr[2], addr[3])
      packet_info['ip_dst'] = ip_header[9]
      addr = struct.unpack("BBBB", ip_packet[16:20])
      packet_info['ip_dst_str'] = '{0:d}.{1:d}.{2:d}.{3:d}'.format(addr[0], addr[1], addr[2], addr[3])
      if payload_length > 0:
        payload = ip_packet[header_length:]
        if packet_info['ip_protocol'] == 6:
          self.ProcessTCPPacket(payload, packet_info)
        elif packet_info['ip_protocol'] == 17:
          self.ProcessUDPPacket(payload, packet_info)
        else:
          packet_info['valid'] = False
    else:
      packet_info['valid'] = False

  def ProcessTCPPacket(self, payload, packet_info):
    # TCP Packet Header
    # Source Port: 2 bytes
    # Dest Port: 2 bytes
    # Sequence number: 4 bytes
    # Ack number: 4 bytes
    # Header len: 1 byte (masked)
    if len(payload) > 8:
      tcp_header = struct.unpack("!HHLLB", payload[0:13])
      header_length = (tcp_header[4] >> 4 & 0x0F) * 4
      packet_info['tcp_payload_length'] = packet_info['ip_payload_length'] - header_length
      packet_info['src_port'] = tcp_header[0]
      packet_info['dst_port'] = tcp_header[1]
      packet_info['tcp_sequence'] = tcp_header[2]
      packet_info['stream_id'] = '{0}:{1:d}->{2}:{3:d}'.format(packet_info['ip_src_str'], packet_info['src_port'],
                                                               packet_info['ip_dst_str'], packet_info['dst_port'])
      # Use traffic <> to port 80 or 443 as a starting point if timing hasn't already started.
      # Assume the server is the one on 80/443 for determining the local machine MAC for ethernet direction.
      if self.start_time is None and\
          (packet_info['src_port'] == 80 or packet_info['dst_port'] == 80 or
           packet_info['src_port'] == 443 or packet_info['dst_port'] == 443):
        self.start_time = packet_info['time']
        if 'ethernet_dst' in packet_info and self.local_ethernet_mac is None and\
            (packet_info['src_port'] == 80 or packet_info['src_port'] == 443):
          self.local_ethernet_mac = packet_info['ethernet_dst']
        elif 'ethernet_src' in packet_info and self.local_ethernet_mac is None and \
            (packet_info['dst_port'] == 80 or packet_info['dst_port'] == 443):
          self.local_ethernet_mac = packet_info['ethernet_src']
    else:
      packet_info['valid'] = False

  def ProcessUDPPacket(self, payload, packet_info):
    # UDP Packet header:
    # Source Port: 2 bytes
    # Dest Port: 2 bytes
    # Length (including header): 2 bytes
    # Checksum: 2 bytes
    if len(payload) > 8:
      udp_header = struct.unpack("!HHHH", payload[0:8])
      packet_info['src_port'] = udp_header[0]
      packet_info['dst_port'] = udp_header[1]
      if packet_info['dst_port'] == 53:
        self.ProcessDNSRequest(payload[8:], packet_info)
    else:
      packet_info['valid'] = False

  def ProcessDNSRequest(self, payload, packet_info):
    if 'ethernet_src' in packet_info and self.local_ethernet_mac is None:
      self.local_ethernet_mac = packet_info['ethernet_src']
    if self.start_time is None:
      self.start_time = packet_info['time']

  def ProcessPacketInfo(self, packet_info):
    elapsed = packet_info['time'] - self.start_time
    bucket = int(math.floor(elapsed * 10))

    # Make sure the time slice lists in both directions are the same size and big enough to include the current bucket
    for direction in ['in', 'out', 'in_dup']:
      length = len(self.slices[direction])
      if length <= bucket:
        need = bucket - length + 1
        self.slices[direction] += [0] * need

    # Update the actual accounting
    bytes = packet_info['length']
    direction = packet_info['direction']
    self.bytes[direction] += bytes
    self.slices[direction][bucket] += bytes

    # If it is a tcp stream, keep track of the sequence numbers and see if any of the data overlaps with previous
    # ranges on the same connection.
    if direction == 'in' and\
            'stream_id' in packet_info and\
            'tcp_sequence' in packet_info and\
            'tcp_payload_length' in packet_info and\
            packet_info['tcp_payload_length'] > 0:
      stream = packet_info['stream_id']
      data_len = packet_info['tcp_payload_length']
      stream_start = packet_info['tcp_sequence']
      stream_end = stream_start + data_len
      if stream not in self.streams:
        self.streams[stream] = []

      # Loop through all of the existing packets on the stream to see if the data is duplicate (a spurious retransmit)
      duplicate_bytes = 0
      for start, end in self.streams[stream]:
        overlap = max(0, min(end, stream_end) - max(start, stream_start))
        if overlap > duplicate_bytes:
          duplicate_bytes = overlap

      # If the entire payload is duplicate then the whole packet is duplicate
      if duplicate_bytes >= data_len:
        duplicate_bytes = packet_info['length']

      if duplicate_bytes > 0:
        self.bytes['in_dup'] += duplicate_bytes
        self.slices['in_dup'][bucket] += duplicate_bytes

      # Keep track of the current packet byte range
      self.streams[stream].append([stream_start, stream_end])

########################################################################################################################
#   Main Entry Point
########################################################################################################################
def main():
  global options
  import argparse
  parser = argparse.ArgumentParser(description='WebPageTest pcap parser.',
                                   prog='pcap-parser')
  parser.add_argument('-v', '--verbose', action='count',
                      help="Increase verbosity (specify multiple times for more). -vvvv for full debug output.")
  parser.add_argument('-i', '--input', help="Input pcap file.")
  parser.add_argument('-s', '--stats', help="Output bandwidth information file.")
  parser.add_argument('-d', '--details', help="Output bandwidth details file (time sliced bandwidth data).")
  parser.add_argument('-j', '--json', action='store_true', default=False, help="Set output format to JSON")
  options = parser.parse_args()

  # Set up logging
  log_level = logging.CRITICAL
  if options.verbose == 1:
    log_level = logging.ERROR
  elif options.verbose == 2:
    log_level = logging.WARNING
  elif options.verbose == 3:
    log_level = logging.INFO
  elif options.verbose >= 4:
    log_level = logging.DEBUG
  logging.basicConfig(level=log_level, format="%(asctime)s.%(msecs)03d - %(message)s", datefmt="%H:%M:%S")

  if not options.input:
    parser.error("Input trace file is not specified.")

  start = time.time()
  pcap = Pcap()
  pcap.Process(options.input)

  if options.stats:
    pcap.SaveStats(options.stats)
  if options.details:
    pcap.SaveDetails(options.details)
  pcap.Print()

  end = time.time()
  elapsed = end - start
  logging.debug("Elapsed Time: {0:0.4f}".format(elapsed))

if '__main__' == __name__:
  main()
<commit_msg>Fixed detection of the start of activity for Opera Mini<commit_after>#!/usr/bin/python
"""
Copyright 2016 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import gzip
import json
import logging
import math
import os
import struct
import time

#Globals
options = None

########################################################################################################################
#   Pcap processing
########################################################################################################################
class Pcap():
  def __init__(self):
    self.start_seconds = None
    self.start_time = None
    self.local_ethernet_mac = None
    self.slices = {'in': [], 'out': [], 'in_dup': []}
    self.bytes = {'in': 0, 'out': 0, 'in_dup': 0}
    self.streams = {}
    return

  def SaveStats(self, out):
    file_name, ext = os.path.splitext(out)
    if ext.lower() == '.gz':
      f = gzip.open(out, 'wb')
    else:
      f = open(out, 'wb')
    try:
      result = {"bytes": self.bytes}
      json.dump(result, f)
      logging.info('Result stats written to {0}'.format(out))
    except:
      logging.critical('Error writing result stats to {0}'.format(out))
    f.close()

  def SaveDetails(self, out):
    file_name, ext = os.path.splitext(out)
    if ext.lower() == '.gz':
      f = gzip.open(out, 'wb')
    else:
      f = open(out, 'wb')
    try:
      json.dump(self.slices, f)
      logging.info('Result details written to {0}'.format(out))
    except:
      logging.critical('Error writing result details to {0}'.format(out))
    f.close()

  def Print(self):
    global options
    if options.json:
      print(json.dumps(self.bytes, indent=2))
    else:
      print "Bytes Out: {0:d}".format(self.bytes['out'])
      print "Bytes In:  {0:d}".format(self.bytes['in'])
      print "Duplicate Bytes In:  {0:d}".format(self.bytes['in_dup'])

  def Process(self, pcap):
    f = None
    self.__init__() #Reset state if called multiple times
    try:
      file_name, ext = os.path.splitext(pcap)
      if ext.lower() == '.gz':
        f = gzip.open(pcap, 'rb')
      else:
        f = open(pcap, 'rb')
      bytes = f.read(24)
      # File header:
      # Magic Number - 4 bytes - 0xa1b2c3d4
      # Major Version - 2 bytes
      # Minor version - 2 bytes
      # Tz offset - 4 bytes (always 0)
      # Timestamp accuracy - 4 bytes (always 0)
      # Snapshot length - 4 bytes
      # Link layer header type - 4 bytes
      #
      # unpack constants:
      # L - unsigned long (4 byte)
      # H - unsigned short (2 byte)
      # B - unsigned char (1 byte int)
      file_header = struct.unpack("=LHHLLLL", bytes)

      # ignore byte order reversals for now
      if file_header[0] == 0xa1b2c3d4:
        ok = True
        self.linktype = file_header[6]
        if self.linktype == 1:
          self.linklen = 12
        elif self.linktype == 113:
          self.linklen = 14
        else:
          logging.critical("Unknown link layer header type: {0:d}".format(self.linktype))
          ok = False

        # Packet header:
        # Time stamp (seconds) - 4 bytes
        # Time stamp (microseconds value) - 4 bytes
        # Captured data length - 4 bytes
        # Original length - 4 bytes
        while ok:
          bytes = f.read(16)
          if not bytes or len(bytes) < 16:
            break
          (seconds, useconds, captured_length, packet_length) = struct.unpack("=LLLL", bytes)
          if self.start_seconds is None:
            self.start_seconds = seconds
          seconds -= self.start_seconds
          if packet_length and captured_length <= packet_length:
            packet_time = float(seconds) + float(useconds) / 1000000.0
            packet_info = {}
            packet_info['time'] = packet_time
            packet_info['length'] = packet_length
            packet_info['captured_length'] = captured_length
            packet_info['valid'] = False
            if captured_length:
              packet_data = f.read(captured_length)
            else:
              packet_data = None
            if len(packet_data) >= self.linklen:
              try:
                self.ProcessPacket(packet_data, packet_info)
              except Exception as e:
                print(e)
      else:
        logging.critical("Invalid pcap file " + pcap)
    except:
      logging.critical("Error processing pcap " + pcap)

    if f is not None:
      f.close()

    return

  def ProcessPacket(self, packet_data, packet_info):
    if self.linktype == 1:
      # Ethernet:
      # dst1: 2 bytes
      # dst2: 2 bytes
      # dst3: 2 bytes
      # src1: 2 bytes
      # src2: 2 bytes
      # src3: 2 bytes
      ethernet_header = struct.unpack("!HHHHHH", packet_data[0:self.linklen])
      # Ignore broadcast traffic
      packet_info['ethernet_dst'] = [ethernet_header[0], ethernet_header[1], ethernet_header[2]]
      packet_info['ethernet_src'] = [ethernet_header[3], ethernet_header[4], ethernet_header[5]]
      dst = packet_info['ethernet_dst']
      if dst[0] != 0xFFFF or dst[1] != 0xFFFF or  dst[2] != 0xFFFF:
        packet_info['valid'] = True

    elif self.linktype == 113:
      # Linux cooked capture
      # Packet Type: 2 bytes
      # aprhrd type: 2 bytes
      # Address length: 2 bytes
      # Address part 1: 4 bytes
      # Address part 2: 4 bytes
      cooked_header = struct.unpack("!HHHLL", packet_data[0:self.linklen])
      if cooked_header[0] == 0:
        packet_info['valid'] = True
        packet_info['direction'] = 'in'
      if cooked_header[0] == 4:
        packet_info['valid'] = True
        packet_info['direction'] = 'out'

    protocol = struct.unpack("!H", packet_data[self.linklen:self.linklen + 2])[0]
    if packet_info['valid'] and protocol == 0x800: # Only handle IPv4 for now
      self.ProcessIPv4Packet(packet_data[self.linklen + 2:], packet_info)

    if packet_info['valid'] and self.start_time:
      if self.local_ethernet_mac is not None:
        local = self.local_ethernet_mac
        src = packet_info['ethernet_src']
        if src[0] == local[0] and src[1] == local[1] and src[2] == local[2]:
          packet_info['direction'] = 'out'
        else:
          packet_info['direction'] = 'in'
      if 'direction' in packet_info:
        self.ProcessPacketInfo(packet_info)

    return

  def ProcessIPv4Packet(self, ip_packet, packet_info):
    # IP Header:
    # Version/len: 1 Byte (4 bits each)
    # dscp/ecn: 1 Byte
    # Total Length: 2 Bytes
    # Identification: 2 Bytes
    # Flags/Fragment: 2 Bytes
    # TTL: 1 Byte
    # Protocol: 1 Byte
    # Header Checksum: 2 Bytes
    # Source Address: 4 Bytes
    # Dest Address: 4 Bytes
    if len(ip_packet) > 20:
      ip_header = struct.unpack("!BBHHHBBHLL", ip_packet[0:20])
      header_length = (ip_header[0] & 0x0F) * 4
      total_length = ip_header[2]
      payload_length = total_length - header_length
      packet_info['ip_payload_length'] = payload_length
      packet_info['ip_protocol'] = ip_header[6]
      packet_info['ip_src'] = ip_header[8]
      addr = struct.unpack("BBBB", ip_packet[12:16])
      packet_info['ip_src_str'] = '{0:d}.{1:d}.{2:d}.{3:d}'.format(addr[0], addr[1], addr[2], addr[3])
      packet_info['ip_dst'] = ip_header[9]
      addr = struct.unpack("BBBB", ip_packet[16:20])
      packet_info['ip_dst_str'] = '{0:d}.{1:d}.{2:d}.{3:d}'.format(addr[0], addr[1], addr[2], addr[3])
      if payload_length > 0:
        payload = ip_packet[header_length:]
        if packet_info['ip_protocol'] == 6:
          self.ProcessTCPPacket(payload, packet_info)
        elif packet_info['ip_protocol'] == 17:
          self.ProcessUDPPacket(payload, packet_info)
        else:
          packet_info['valid'] = False
    else:
      packet_info['valid'] = False

  def ProcessTCPPacket(self, payload, packet_info):
    # TCP Packet Header
    # Source Port: 2 bytes
    # Dest Port: 2 bytes
    # Sequence number: 4 bytes
    # Ack number: 4 bytes
    # Header len: 1 byte (masked)
    if len(payload) > 8:
      tcp_header = struct.unpack("!HHLLB", payload[0:13])
      header_length = (tcp_header[4] >> 4 & 0x0F) * 4
      packet_info['tcp_payload_length'] = packet_info['ip_payload_length'] - header_length
      packet_info['src_port'] = tcp_header[0]
      packet_info['dst_port'] = tcp_header[1]
      packet_info['tcp_sequence'] = tcp_header[2]
      packet_info['stream_id'] = '{0}:{1:d}->{2}:{3:d}'.format(packet_info['ip_src_str'], packet_info['src_port'],
                                                               packet_info['ip_dst_str'], packet_info['dst_port'])
      # If DNS didn't trigger a start yet and we see outbound TCP traffic, use that to identify the starting point.
      # Outbound can be explicit (if we have a cooked capture like android) or implicit if dest port is 80, 443, 1080.
      if self.start_time is None:
        is_outbound = False
        if 'direction' in packet_info and packet_info['direction'] == 'out':
          is_outbound = True
        elif packet_info['dst_port'] == 80 or packet_info['dst_port'] == 443 or packet_info['dst_port'] == 1080:
          is_outbound = True
        if is_outbound:
          self.start_time = packet_info['time']
          if 'ethernet_src' in packet_info and self.local_ethernet_mac is None:
            self.local_ethernet_mac = packet_info['ethernet_src']
    else:
      packet_info['valid'] = False

  def ProcessUDPPacket(self, payload, packet_info):
    # UDP Packet header:
    # Source Port: 2 bytes
    # Dest Port: 2 bytes
    # Length (including header): 2 bytes
    # Checksum: 2 bytes
    if len(payload) > 8:
      udp_header = struct.unpack("!HHHH", payload[0:8])
      packet_info['src_port'] = udp_header[0]
      packet_info['dst_port'] = udp_header[1]
      if packet_info['dst_port'] == 53:
        self.ProcessDNSRequest(payload[8:], packet_info)
    else:
      packet_info['valid'] = False

  def ProcessDNSRequest(self, payload, packet_info):
    if 'ethernet_src' in packet_info and self.local_ethernet_mac is None:
      self.local_ethernet_mac = packet_info['ethernet_src']
    if self.start_time is None:
      self.start_time = packet_info['time']

  def ProcessPacketInfo(self, packet_info):
    elapsed = packet_info['time'] - self.start_time
    bucket = int(math.floor(elapsed * 10))

    # Make sure the time slice lists in both directions are the same size and big enough to include the current bucket
    for direction in ['in', 'out', 'in_dup']:
      length = len(self.slices[direction])
      if length <= bucket:
        need = bucket - length + 1
        self.slices[direction] += [0] * need

    # Update the actual accounting
    bytes = packet_info['length']
    direction = packet_info['direction']
    self.bytes[direction] += bytes
    self.slices[direction][bucket] += bytes

    # If it is a tcp stream, keep track of the sequence numbers and see if any of the data overlaps with previous
    # ranges on the same connection.
    if direction == 'in' and\
            'stream_id' in packet_info and\
            'tcp_sequence' in packet_info and\
            'tcp_payload_length' in packet_info and\
            packet_info['tcp_payload_length'] > 0:
      stream = packet_info['stream_id']
      data_len = packet_info['tcp_payload_length']
      stream_start = packet_info['tcp_sequence']
      stream_end = stream_start + data_len
      if stream not in self.streams:
        self.streams[stream] = []

      # Loop through all of the existing packets on the stream to see if the data is duplicate (a spurious retransmit)
      duplicate_bytes = 0
      for start, end in self.streams[stream]:
        overlap = max(0, min(end, stream_end) - max(start, stream_start))
        if overlap > duplicate_bytes:
          duplicate_bytes = overlap

      # If the entire payload is duplicate then the whole packet is duplicate
      if duplicate_bytes >= data_len:
        duplicate_bytes = packet_info['length']

      if duplicate_bytes > 0:
        self.bytes['in_dup'] += duplicate_bytes
        self.slices['in_dup'][bucket] += duplicate_bytes

      # Keep track of the current packet byte range
      self.streams[stream].append([stream_start, stream_end])

########################################################################################################################
#   Main Entry Point
########################################################################################################################
def main():
  global options
  import argparse
  parser = argparse.ArgumentParser(description='WebPageTest pcap parser.',
                                   prog='pcap-parser')
  parser.add_argument('-v', '--verbose', action='count',
                      help="Increase verbosity (specify multiple times for more). -vvvv for full debug output.")
  parser.add_argument('-i', '--input', help="Input pcap file.")
  parser.add_argument('-s', '--stats', help="Output bandwidth information file.")
  parser.add_argument('-d', '--details', help="Output bandwidth details file (time sliced bandwidth data).")
  parser.add_argument('-j', '--json', action='store_true', default=False, help="Set output format to JSON")
  options = parser.parse_args()

  # Set up logging
  log_level = logging.CRITICAL
  if options.verbose == 1:
    log_level = logging.ERROR
  elif options.verbose == 2:
    log_level = logging.WARNING
  elif options.verbose == 3:
    log_level = logging.INFO
  elif options.verbose >= 4:
    log_level = logging.DEBUG
  logging.basicConfig(level=log_level, format="%(asctime)s.%(msecs)03d - %(message)s", datefmt="%H:%M:%S")

  if not options.input:
    parser.error("Input trace file is not specified.")

  start = time.time()
  pcap = Pcap()
  pcap.Process(options.input)

  if options.stats:
    pcap.SaveStats(options.stats)
  if options.details:
    pcap.SaveDetails(options.details)
  pcap.Print()

  end = time.time()
  elapsed = end - start
  logging.debug("Elapsed Time: {0:0.4f}".format(elapsed))

if '__main__' == __name__:
  main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from openerp import http
from openerp.http import request
from openerp import SUPERUSER_ID

from openerp.addons.survey.controllers.main import WebsiteSurvey
import json
import logging

_logger = logging.getLogger(__name__)

class StrengthItem():
    name = ''
    value = 0

    def __init__(self, name, value):
        self.name = name
        self.value = value

class StrengthsResults(http.Controller):

    def _check_bad_cases(self, cr, uid, request, survey_obj, survey, user_input_obj, context=None):
        # In case of bad survey, redirect to surveys list
        if survey_obj.exists(cr, SUPERUSER_ID, survey.id, context=context) == []:
            return werkzeug.utils.redirect("/survey/")

        # In case of auth required, block public user
        if survey.auth_required and uid == request.website.user_id.id:
            return request.website.render("survey.auth_required", {'survey': survey})

        # In case of non open surveys
        if survey.stage_id.closed:
            return request.website.render("survey.notopen")

        # If there is no pages
        if not survey.page_ids:
            return request.website.render("survey.nopages")

        # Everything seems to be ok
        return None

# Survey wrapper for strengths survey response
    @http.route(['/mystrengths/<string:token>'],
                type='http', auth='public', website=True)
    def survey_wrapper(self, token, prev=None, **post):
        '''wraps a response'''
        cr, uid, context = request.cr, request.uid, request.context
        survey_obj = request.registry['survey.survey']
        user_input_obj = request.registry['survey.user_input']

        survey_id = survey_obj.search(cr, SUPERUSER_ID, [('title', '=', 'Strengths')], context=context)[0]
        survey = survey_obj.browse(cr, SUPERUSER_ID, [survey_id], context=context)[0]

        # Controls if the survey can be displayed
        errpage = self._check_bad_cases(cr, uid, request, survey_obj, survey, user_input_obj, context=context)
        if errpage:
            return errpage

        # see if response for this survey already exists.
        try:
            user_input_id = user_input_obj.search(cr, SUPERUSER_ID, [('token', '=', token), ('survey_id', '=', survey_id)], context=context)[0]
        except IndexError:  # No response currently exists for this assignment
            print('response token was invalid')
            return request.website.render("website.403")
        else:
            user_input = user_input_obj.browse(cr, SUPERUSER_ID, [user_input_id], context=context)[0]

        # Identify what category each question is in
        dic_cat = dict()
        for page in survey.page_ids:
            question_ids = []
            for question in page.question_ids:
                question_ids.append(question.id)
                dic_cat[question.id] = question.categ_id.name

        # create result dictionary including each unique category
        strength_names = set(dic_cat.values())
        res = dict()
        for item in strength_names:
            res[item] = 0.0

        n = res

        # get responses
        #user_input_line_obj = request.registry['survey.user_input_line']
        #user_input_line_ids = user_input_line_obj.search(cr, SUPERUSER_ID, [('survey_id', '=', survey_id)], context=context)
        #user_input_line = user_input_line_obj.browse(cr, SUPERUSER_ID, user_input_line_ids, context=context)

        # analyze survey results, using dictionary of question ids and strength_names
        #for line in user_input_line:
        for line in user_input.user_input_line_ids:
            cat = dic_cat[line.question_id.id]
            testing = False
            if cat == "Learning": testing = True
            if testing: print('mark=' + str(line.quizz_mark) + ", res[cat] pre=" + str(res[cat]) + ", n[cat] pre=" + str(n[cat]))
            res[cat] = (res[cat]*(n[cat]) + line.quizz_mark) / (n[cat] + 1.0)
            n[cat] = n[cat] + 1.0
            if testing: print('mark=' + str(line.quizz_mark) + ", res[cat] post=" + str(res[cat]) + ", n[cat] post=" + str(n[cat]))

        strengths_obj = request.registry['mqo.snippet.strengths']
        # strength_ids = strengths_obj.search(cr, SUPERUSER_ID, [], context=context)
        # strengths = strengths_obj.browse(cr, SUPERUSER_ID, strength_ids, context=context)
        strengths = strengths_obj.search_read(cr, SUPERUSER_ID, [], ['id', 'name', 'title', 'text'])
        for strength in strengths:
            strength["value"] = res[strength["name"]]

        strengths = sorted(strengths, key=lambda strength: strength["value"], reverse=True)
        for i, strength in enumerate(strengths):
            strength["seq_num"] = i + 1

        vals = {'strengths': strengths}
        return request.website.render("mqo_strengths.results", vals)

<commit_msg>Bugfixing averaging algorithm more.<commit_after># -*- coding: utf-8 -*-
from openerp import http
from openerp.http import request
from openerp import SUPERUSER_ID

from openerp.addons.survey.controllers.main import WebsiteSurvey
import json
import logging

_logger = logging.getLogger(__name__)

class StrengthItem():
    name = ''
    value = 0

    def __init__(self, name, value):
        self.name = name
        self.value = value

class StrengthsResults(http.Controller):

    def _check_bad_cases(self, cr, uid, request, survey_obj, survey, user_input_obj, context=None):
        # In case of bad survey, redirect to surveys list
        if survey_obj.exists(cr, SUPERUSER_ID, survey.id, context=context) == []:
            return werkzeug.utils.redirect("/survey/")

        # In case of auth required, block public user
        if survey.auth_required and uid == request.website.user_id.id:
            return request.website.render("survey.auth_required", {'survey': survey})

        # In case of non open surveys
        if survey.stage_id.closed:
            return request.website.render("survey.notopen")

        # If there is no pages
        if not survey.page_ids:
            return request.website.render("survey.nopages")

        # Everything seems to be ok
        return None

# Survey wrapper for strengths survey response
    @http.route(['/mystrengths/<string:token>'],
                type='http', auth='public', website=True)
    def survey_wrapper(self, token, prev=None, **post):
        '''wraps a response'''
        cr, uid, context = request.cr, request.uid, request.context
        survey_obj = request.registry['survey.survey']
        user_input_obj = request.registry['survey.user_input']

        survey_id = survey_obj.search(cr, SUPERUSER_ID, [('title', '=', 'Strengths')], context=context)[0]
        survey = survey_obj.browse(cr, SUPERUSER_ID, [survey_id], context=context)[0]

        # Controls if the survey can be displayed
        errpage = self._check_bad_cases(cr, uid, request, survey_obj, survey, user_input_obj, context=context)
        if errpage:
            return errpage

        # see if response for this survey already exists.
        try:
            user_input_id = user_input_obj.search(cr, SUPERUSER_ID, [('token', '=', token), ('survey_id', '=', survey_id)], context=context)[0]
        except IndexError:  # No response currently exists for this assignment
            print('response token was invalid')
            return request.website.render("website.403")
        else:
            user_input = user_input_obj.browse(cr, SUPERUSER_ID, [user_input_id], context=context)[0]

        # Identify what category each question is in
        dic_cat = dict()
        for page in survey.page_ids:
            question_ids = []
            for question in page.question_ids:
                question_ids.append(question.id)
                dic_cat[question.id] = question.categ_id.name

        # create result dictionary including each unique category
        strength_names = set(dic_cat.values())
        res = dict()
        n = dict()
        for item in strength_names:
            res[item] = 0.0
            n[item] = 0.0

        # get responses
        #user_input_line_obj = request.registry['survey.user_input_line']
        #user_input_line_ids = user_input_line_obj.search(cr, SUPERUSER_ID, [('survey_id', '=', survey_id)], context=context)
        #user_input_line = user_input_line_obj.browse(cr, SUPERUSER_ID, user_input_line_ids, context=context)

        # analyze survey results, using dictionary of question ids and strength_names
        #for line in user_input_line:
        for line in user_input.user_input_line_ids:
            cat = dic_cat[line.question_id.id]
            testing = False
            if cat == "Learning": testing = True
            if testing: print('mark=' + str(line.quizz_mark) + ", res[cat] pre=" + str(res[cat]) + ", n[cat] pre=" + str(n[cat]))
            res[cat] = (res[cat]*(n[cat]) + line.quizz_mark) / (n[cat] + 1.0)
            n[cat] = n[cat] + 1.0
            if testing: print('mark=' + str(line.quizz_mark) + ", res[cat] post=" + str(res[cat]) + ", n[cat] post=" + str(n[cat]))

        strengths_obj = request.registry['mqo.snippet.strengths']
        # strength_ids = strengths_obj.search(cr, SUPERUSER_ID, [], context=context)
        # strengths = strengths_obj.browse(cr, SUPERUSER_ID, strength_ids, context=context)
        strengths = strengths_obj.search_read(cr, SUPERUSER_ID, [], ['id', 'name', 'title', 'text'])
        for strength in strengths:
            strength["value"] = res[strength["name"]]

        strengths = sorted(strengths, key=lambda strength: strength["value"], reverse=True)
        for i, strength in enumerate(strengths):
            strength["seq_num"] = i + 1

        vals = {'strengths': strengths}
        return request.website.render("mqo_strengths.results", vals)

<|endoftext|>


<commit_before>from datetime import datetime
from couchdbkit import ResourceNotFound
from django.contrib import messages
from django.core.exceptions import SuspiciousOperation
from django.core.urlresolvers import reverse
from django.http import HttpResponseRedirect, HttpResponseBadRequest, Http404
from django.utils.decorators import method_decorator
import json
from corehq import toggles
from corehq.apps.app_manager.dbaccessors import get_apps_in_domain
from corehq.apps.app_manager.models import Application
from corehq.apps.app_manager.templatetags.xforms_extras import trans
from corehq.apps.export.custom_export_helpers import make_custom_export_helper
from corehq.apps.export.exceptions import ExportNotFound, ExportAppException
from corehq.apps.export.forms import CreateFormExportForm, CreateCaseExportForm
from corehq.apps.reports.dbaccessors import touch_exports
from corehq.apps.reports.display import xmlns_to_name
from corehq.apps.reports.standard.export import (
    CaseExportInterface,
    CaseExportReport,
    FormExportInterface,
    ExcelExportReport,
)
from corehq.apps.settings.views import BaseProjectDataView
from corehq.apps.users.decorators import require_permission
from corehq.apps.users.models import Permissions
from couchexport.models import SavedExportSchema, ExportSchema
from couchexport.schema import build_latest_schema
from dimagi.utils.decorators.memoized import memoized
from django.utils.translation import ugettext as _, ugettext_noop, ugettext_lazy
from dimagi.utils.logging import notify_exception
from dimagi.utils.parsing import json_format_date
from dimagi.utils.web import json_response
from toggle import toggle_enabled

require_form_export_permission = require_permission(
    Permissions.view_report,
    'corehq.apps.reports.standard.export.ExcelExportReport',
    login_decorator=None
)

class BaseExportView(BaseProjectDataView):
    template_name = 'export/customize_export.html'
    export_type = None
    is_async = True

    @property
    def parent_pages(self):
        return [{
            'title': self.report_class.name,
            'url': self.export_home_url,
        }]

    @method_decorator(require_form_export_permission)
    def dispatch(self, request, *args, **kwargs):
        return super(BaseExportView, self).dispatch(request, *args, **kwargs)

    @property
    def export_helper(self):
        raise NotImplementedError("You must implement export_helper!")

    def redirect_url(self, export_id):
        if self.request.body:
            preview = json.loads(self.request.body).get('preview')
            if preview:
                return reverse(
                    'export_custom_data',
                    args=[self.domain, export_id],
                ) + '?format=html&limit=50&type=%(type)s' % {
                    'type': self.export_type,
                }
        return self.export_home_url

    @property
    def export_home_url(self):
        return self.report_class.get_url(domain=self.domain)

    @property
    @memoized
    def report_class(self):
        try:
            if toggle_enabled(self.request, toggles.REVAMPED_EXPORTS):
                return {
                    'form': FormExportInterface,
                    'case': CaseExportInterface,
                }[self.export_type]
            return {
                'form': ExcelExportReport,
                'case': CaseExportReport
            }[self.export_type]
        except KeyError:
            raise SuspiciousOperation

    @property
    def page_context(self):
        return self.export_helper.get_context()

    def commit(self, request):
        raise NotImplementedError('Subclasses must implement a commit method.')

    def post(self, request, *args, **kwargs):
        try:
            export_id = self.commit(request)
        except Exception, e:
            if self.is_async:
                # todo: this can probably be removed as soon as
                # http://manage.dimagi.com/default.asp?157713 is resolved
                notify_exception(request, 'problem saving an export! {}'.format(str(e)))
                response = json_response({
                    'error': str(e) or type(e).__name__
                })
                response.status_code = 500
                return response
            elif isinstance(e, ExportAppException):
                return HttpResponseRedirect(request.META['HTTP_REFERER'])
            else:
                raise
        else:
            if self.is_async:
                return json_response({
                    'redirect': self.redirect_url(export_id),
                })
            return HttpResponseRedirect(self.redirect_url(export_id))

class BaseCreateCustomExportView(BaseExportView):
    # this view likely needs a lot more cleanup. will leave that for a later time...

    @property
    @memoized
    def export_helper(self):
        return make_custom_export_helper(self.request, self.export_type, domain=self.domain)

    def commit(self, request):
        export_id = self.export_helper.update_custom_export()
        messages.success(request, _("Custom export created!"))
        return export_id

    def get(self, request, *args, **kwargs):
        # just copying what was in the old django view here. don't want to mess too much with exports just yet.
        try:
            export_tag = [self.domain, json.loads(request.GET.get("export_tag", "null") or "null")]
        except ValueError:
            return HttpResponseBadRequest()

        if self.export_helper.export_type == "form" and not export_tag[1]:
            return HttpResponseRedirect(ExcelExportReport.get_url(domain=self.domain))

        schema = build_latest_schema(export_tag)

        if not schema and self.export_helper.export_type == "form":
            schema = create_basic_form_checkpoint(export_tag)

        if request.GET.get('minimal', False):
            messages.warning(request,
                _("Warning you are using minimal mode, some things may not be functional"))

        if schema:
            app_id = request.GET.get('app_id')
            self.export_helper.custom_export = self.export_helper.ExportSchemaClass.default(
                schema=schema,
                name="%s: %s" % (
                    xmlns_to_name(self.domain, export_tag[1], app_id=app_id)
                        if self.export_helper.export_type == "form" else export_tag[1],
                    json_format_date(datetime.utcnow())
                ),
                type=self.export_helper.export_type
            )
            if self.export_helper.export_type in ['form', 'case']:
                self.export_helper.custom_export.app_id = app_id
            if self.export_helper.export_type == 'form':
                self.export_helper.custom_export.update_question_schema()

            return super(BaseCreateCustomExportView, self).get(request, *args, **kwargs)

        messages.warning(request, _("No data found for that form "
                                    "(%s). Submit some data before creating an export!") %
                         xmlns_to_name(self.domain, export_tag[1], app_id=None), extra_tags="html")
        return HttpResponseRedirect(ExcelExportReport.get_url(domain=self.domain))

class CreateCustomFormExportView(BaseCreateCustomExportView):
    urlname = 'custom_export_form'
    page_title = ugettext_lazy("Create Custom Form Export")
    export_type = 'form'

class CreateCustomCaseExportView(BaseCreateCustomExportView):
    urlname = 'custom_export_case'
    page_title = ugettext_lazy("Create Custom Case Export")
    export_type = 'case'

class BaseModifyCustomExportView(BaseExportView):

    @property
    def page_url(self):
        return reverse(self.urlname, args=[self.domain, self.export_id])

    @property
    def export_id(self):
        return self.kwargs.get('export_id')

    @property
    @memoized
    def export_helper(self):
        try:
            return make_custom_export_helper(self.request, self.export_type, self.domain, self.export_id)
        except ResourceNotFound:
            raise Http404()

class BaseEditCustomExportView(BaseModifyCustomExportView):

    def commit(self, request):
        export_id = self.export_helper.update_custom_export()
        messages.success(request, _("Custom export saved!"))
        return export_id

class EditCustomFormExportView(BaseEditCustomExportView):
    urlname = 'edit_custom_export_form'
    page_title = ugettext_noop("Edit Form Custom Export")
    export_type = 'form'

class EditCustomCaseExportView(BaseEditCustomExportView):
    urlname = 'edit_custom_export_case'
    page_title = ugettext_noop("Edit Case Custom Export")
    export_type = 'case'

class DeleteCustomExportView(BaseModifyCustomExportView):
    urlname = 'delete_custom_export'
    http_method_names = ['post']
    is_async = False

    def commit(self, request):
        try:
            saved_export = SavedExportSchema.get(self.export_id)
        except ResourceNotFound:
            raise ExportNotFound()
        self.export_type = saved_export.type
        saved_export.delete()
        touch_exports(self.domain)
        messages.success(request, _("Custom export was deleted."))

BASIC_FORM_SCHEMA = {
    "doc_type": "string",
    "domain": "string",
    "xmlns": "string",
    "form": {
        "@xmlns": "string",
        "@uiVersion": "string",
        "@name": "string",
        "#type": "string",
        "meta": {
            "@xmlns": "string",
            "username": "string",
            "instanceID": "string",
            "userID": "string",
            "timeEnd": "string",
            "appVersion": {
                "@xmlns": "string",
                "#text": "string"
            },
            "timeStart": "string",
            "deviceID": "string"
        },
        "@version": "string"
    },
    "partial_submission": "string",
    "_rev": "string",
    "#export_tag": [
       "string"
    ],
    "received_on": "string",
    "app_id": "string",
    "last_sync_token": "string",
    "submit_ip": "string",
    "computed_": {
    },
    "openrosa_headers": {
       "HTTP_DATE": "string",
       "HTTP_ACCEPT_LANGUAGE": "string",
       "HTTP_X_OPENROSA_VERSION": "string"
    },
    "date_header": "string",
    "path": "string",
    "computed_modified_on_": "string",
    "_id": "string"
}

def create_basic_form_checkpoint(index):
    checkpoint = ExportSchema(
        schema=BASIC_FORM_SCHEMA,
        timestamp=datetime(1970, 1, 1),
        index=index,
    )
    checkpoint.save()
    return checkpoint

class CreateFormExportView(BaseProjectDataView):
    urlname = 'create_form_export'
    page_title = ugettext_noop("Create Form Export")
    template_name = 'export/create_form_export.html'

    @property
    def main_context(self):
        context = super(CreateFormExportView, self).main_context
        context.update({
            'create_export_form': self.create_export_form,
            'app_to_module_options': self.app_to_module_options,
            'module_to_form_options': self.module_to_form_options,
            'module_prompt': _('Select Module...'),
            'form_prompt': _('Select Form...'),
        })
        return context

    def post(self, request, *args, **kwargs):
        if self.create_export_form.is_valid():
            app_id = self.create_export_form.cleaned_data['application']
            form_unique_id = self.create_export_form.cleaned_data['form']
            return HttpResponseRedirect(
                reverse(
                    CreateCustomFormExportView.urlname,
                    args=[self.domain],
                ) + ('?export_tag="%(export_tag)s"&app_id=%(app_id)s' % {
                    'app_id': app_id,
                    'export_tag': [
                        form for form in Application.get(app_id).get_forms()
                        if form.get_unique_id() == form_unique_id
                    ][0].xmlns,
                })
            )
        return self.get(self.request, *args, **kwargs)

    @property
    @memoized
    def create_export_form(self):
        if self.request.method == 'POST':
            return CreateFormExportForm(self.domain, self.request.POST)
        return CreateFormExportForm(self.domain)

    @property
    def app_to_module_options(self):
        return {
            app._id: [{
                'text': trans(module.name, app.langs),
                'value': module.unique_id,
            } for module in app.modules]
            for app in get_apps_in_domain(self.domain)
        }

    @property
    def module_to_form_options(self):
        return {
            module.unique_id: [{
                'text': trans(form.name, app.langs),
                'value': form.unique_id,
            } for form in module.get_forms()]
            for app in get_apps_in_domain(self.domain)
            for module in app.modules
        }

class CreateCaseExportView(BaseProjectDataView):
    urlname = 'create_case_export'
    page_title = ugettext_noop('Create Case Export')
    template_name = 'export/create_case_export.html'

    @property
    def main_context(self):
        context = super(CreateCaseExportView, self).main_context
        context.update({
            'create_export_form': self.create_export_form,
            'app_to_case_type_options': self.app_to_case_type_options,
            'case_type_prompt': _('Select Case Type...'),
        })
        return context

    def post(self, request, *args, **kwargs):
        if self.create_export_form.is_valid():
            case_type = self.create_export_form.cleaned_data['case_type']
            return HttpResponseRedirect(
                reverse(
                    CreateCustomCaseExportView.urlname,
                    args=[self.domain],
                ) + ('?export_tag="%(export_tag)s"' % {
                    'export_tag': case_type,
                })
            )
        return self.get(self.request, *args, **kwargs)

    @property
    @memoized
    def create_export_form(self):
        if self.request.method == 'POST':
            return CreateCaseExportForm(self.domain, self.request.POST)
        return CreateCaseExportForm(self.domain)

    @property
    def app_to_case_type_options(self):
        return {
            app._id: [{
                'text': case_type,
                'value': case_type,
            } for case_type in set(
                module.case_type for module in app.modules if module.case_type
            )]
            for app in get_apps_in_domain(self.domain)
        }
<commit_msg>remove buggy imports and views (I think a recent imports fix may have contributed to this?)<commit_after>from datetime import datetime
from couchdbkit import ResourceNotFound
from django.contrib import messages
from django.core.exceptions import SuspiciousOperation
from django.core.urlresolvers import reverse
from django.http import HttpResponseRedirect, HttpResponseBadRequest, Http404
from django.utils.decorators import method_decorator
import json
from corehq import toggles
from corehq.apps.app_manager.dbaccessors import get_apps_in_domain
from corehq.apps.app_manager.models import Application
from corehq.apps.app_manager.templatetags.xforms_extras import trans
from corehq.apps.export.custom_export_helpers import make_custom_export_helper
from corehq.apps.export.exceptions import ExportNotFound, ExportAppException
from corehq.apps.export.forms import CreateFormExportForm, CreateCaseExportForm
from corehq.apps.hqwebapp.templatetags.hq_shared_tags import toggle_enabled
from corehq.apps.reports.dbaccessors import touch_exports
from corehq.apps.reports.display import xmlns_to_name
from corehq.apps.reports.standard.export import (
    CaseExportReport,
    ExcelExportReport,
)
from corehq.apps.settings.views import BaseProjectDataView
from corehq.apps.users.decorators import require_permission
from corehq.apps.users.models import Permissions
from couchexport.models import SavedExportSchema, ExportSchema
from couchexport.schema import build_latest_schema
from dimagi.utils.decorators.memoized import memoized
from django.utils.translation import ugettext as _, ugettext_noop, ugettext_lazy
from dimagi.utils.logging import notify_exception
from dimagi.utils.parsing import json_format_date
from dimagi.utils.web import json_response

require_form_export_permission = require_permission(
    Permissions.view_report,
    'corehq.apps.reports.standard.export.ExcelExportReport',
    login_decorator=None
)

class BaseExportView(BaseProjectDataView):
    template_name = 'export/customize_export.html'
    export_type = None
    is_async = True

    @property
    def parent_pages(self):
        return [{
            'title': self.report_class.name,
            'url': self.export_home_url,
        }]

    @method_decorator(require_form_export_permission)
    def dispatch(self, request, *args, **kwargs):
        return super(BaseExportView, self).dispatch(request, *args, **kwargs)

    @property
    def export_helper(self):
        raise NotImplementedError("You must implement export_helper!")

    def redirect_url(self, export_id):
        if self.request.body:
            preview = json.loads(self.request.body).get('preview')
            if preview:
                return reverse(
                    'export_custom_data',
                    args=[self.domain, export_id],
                ) + '?format=html&limit=50&type=%(type)s' % {
                    'type': self.export_type,
                }
        return self.export_home_url

    @property
    def export_home_url(self):
        return self.report_class.get_url(domain=self.domain)

    @property
    @memoized
    def report_class(self):
        try:
            return {
                'form': ExcelExportReport,
                'case': CaseExportReport
            }[self.export_type]
        except KeyError:
            raise SuspiciousOperation

    @property
    def page_context(self):
        return self.export_helper.get_context()

    def commit(self, request):
        raise NotImplementedError('Subclasses must implement a commit method.')

    def post(self, request, *args, **kwargs):
        try:
            export_id = self.commit(request)
        except Exception, e:
            if self.is_async:
                # todo: this can probably be removed as soon as
                # http://manage.dimagi.com/default.asp?157713 is resolved
                notify_exception(request, 'problem saving an export! {}'.format(str(e)))
                response = json_response({
                    'error': str(e) or type(e).__name__
                })
                response.status_code = 500
                return response
            elif isinstance(e, ExportAppException):
                return HttpResponseRedirect(request.META['HTTP_REFERER'])
            else:
                raise
        else:
            if self.is_async:
                return json_response({
                    'redirect': self.redirect_url(export_id),
                })
            return HttpResponseRedirect(self.redirect_url(export_id))

class BaseCreateCustomExportView(BaseExportView):
    # this view likely needs a lot more cleanup. will leave that for a later time...

    @property
    @memoized
    def export_helper(self):
        return make_custom_export_helper(self.request, self.export_type, domain=self.domain)

    def commit(self, request):
        export_id = self.export_helper.update_custom_export()
        messages.success(request, _("Custom export created!"))
        return export_id

    def get(self, request, *args, **kwargs):
        # just copying what was in the old django view here. don't want to mess too much with exports just yet.
        try:
            export_tag = [self.domain, json.loads(request.GET.get("export_tag", "null") or "null")]
        except ValueError:
            return HttpResponseBadRequest()

        if self.export_helper.export_type == "form" and not export_tag[1]:
            return HttpResponseRedirect(ExcelExportReport.get_url(domain=self.domain))

        schema = build_latest_schema(export_tag)

        if not schema and self.export_helper.export_type == "form":
            schema = create_basic_form_checkpoint(export_tag)

        if request.GET.get('minimal', False):
            messages.warning(request,
                _("Warning you are using minimal mode, some things may not be functional"))

        if schema:
            app_id = request.GET.get('app_id')
            self.export_helper.custom_export = self.export_helper.ExportSchemaClass.default(
                schema=schema,
                name="%s: %s" % (
                    xmlns_to_name(self.domain, export_tag[1], app_id=app_id)
                        if self.export_helper.export_type == "form" else export_tag[1],
                    json_format_date(datetime.utcnow())
                ),
                type=self.export_helper.export_type
            )
            if self.export_helper.export_type in ['form', 'case']:
                self.export_helper.custom_export.app_id = app_id
            if self.export_helper.export_type == 'form':
                self.export_helper.custom_export.update_question_schema()

            return super(BaseCreateCustomExportView, self).get(request, *args, **kwargs)

        messages.warning(request, _("No data found for that form "
                                    "(%s). Submit some data before creating an export!") %
                         xmlns_to_name(self.domain, export_tag[1], app_id=None), extra_tags="html")
        return HttpResponseRedirect(ExcelExportReport.get_url(domain=self.domain))

class CreateCustomFormExportView(BaseCreateCustomExportView):
    urlname = 'custom_export_form'
    page_title = ugettext_lazy("Create Custom Form Export")
    export_type = 'form'

class CreateCustomCaseExportView(BaseCreateCustomExportView):
    urlname = 'custom_export_case'
    page_title = ugettext_lazy("Create Custom Case Export")
    export_type = 'case'

class BaseModifyCustomExportView(BaseExportView):

    @property
    def page_url(self):
        return reverse(self.urlname, args=[self.domain, self.export_id])

    @property
    def export_id(self):
        return self.kwargs.get('export_id')

    @property
    @memoized
    def export_helper(self):
        try:
            return make_custom_export_helper(self.request, self.export_type, self.domain, self.export_id)
        except ResourceNotFound:
            raise Http404()

class BaseEditCustomExportView(BaseModifyCustomExportView):

    def commit(self, request):
        export_id = self.export_helper.update_custom_export()
        messages.success(request, _("Custom export saved!"))
        return export_id

class EditCustomFormExportView(BaseEditCustomExportView):
    urlname = 'edit_custom_export_form'
    page_title = ugettext_noop("Edit Form Custom Export")
    export_type = 'form'

class EditCustomCaseExportView(BaseEditCustomExportView):
    urlname = 'edit_custom_export_case'
    page_title = ugettext_noop("Edit Case Custom Export")
    export_type = 'case'

class DeleteCustomExportView(BaseModifyCustomExportView):
    urlname = 'delete_custom_export'
    http_method_names = ['post']
    is_async = False

    def commit(self, request):
        try:
            saved_export = SavedExportSchema.get(self.export_id)
        except ResourceNotFound:
            raise ExportNotFound()
        self.export_type = saved_export.type
        saved_export.delete()
        touch_exports(self.domain)
        messages.success(request, _("Custom export was deleted."))

BASIC_FORM_SCHEMA = {
    "doc_type": "string",
    "domain": "string",
    "xmlns": "string",
    "form": {
        "@xmlns": "string",
        "@uiVersion": "string",
        "@name": "string",
        "#type": "string",
        "meta": {
            "@xmlns": "string",
            "username": "string",
            "instanceID": "string",
            "userID": "string",
            "timeEnd": "string",
            "appVersion": {
                "@xmlns": "string",
                "#text": "string"
            },
            "timeStart": "string",
            "deviceID": "string"
        },
        "@version": "string"
    },
    "partial_submission": "string",
    "_rev": "string",
    "#export_tag": [
       "string"
    ],
    "received_on": "string",
    "app_id": "string",
    "last_sync_token": "string",
    "submit_ip": "string",
    "computed_": {
    },
    "openrosa_headers": {
       "HTTP_DATE": "string",
       "HTTP_ACCEPT_LANGUAGE": "string",
       "HTTP_X_OPENROSA_VERSION": "string"
    },
    "date_header": "string",
    "path": "string",
    "computed_modified_on_": "string",
    "_id": "string"
}

def create_basic_form_checkpoint(index):
    checkpoint = ExportSchema(
        schema=BASIC_FORM_SCHEMA,
        timestamp=datetime(1970, 1, 1),
        index=index,
    )
    checkpoint.save()
    return checkpoint

class CreateFormExportView(BaseProjectDataView):
    urlname = 'create_form_export'
    page_title = ugettext_noop("Create Form Export")
    template_name = 'export/create_form_export.html'

    @property
    def main_context(self):
        context = super(CreateFormExportView, self).main_context
        context.update({
            'create_export_form': self.create_export_form,
            'app_to_module_options': self.app_to_module_options,
            'module_to_form_options': self.module_to_form_options,
            'module_prompt': _('Select Module...'),
            'form_prompt': _('Select Form...'),
        })
        return context

    def post(self, request, *args, **kwargs):
        if self.create_export_form.is_valid():
            app_id = self.create_export_form.cleaned_data['application']
            form_unique_id = self.create_export_form.cleaned_data['form']
            return HttpResponseRedirect(
                reverse(
                    CreateCustomFormExportView.urlname,
                    args=[self.domain],
                ) + ('?export_tag="%(export_tag)s"&app_id=%(app_id)s' % {
                    'app_id': app_id,
                    'export_tag': [
                        form for form in Application.get(app_id).get_forms()
                        if form.get_unique_id() == form_unique_id
                    ][0].xmlns,
                })
            )
        return self.get(self.request, *args, **kwargs)

    @property
    @memoized
    def create_export_form(self):
        if self.request.method == 'POST':
            return CreateFormExportForm(self.domain, self.request.POST)
        return CreateFormExportForm(self.domain)

    @property
    def app_to_module_options(self):
        return {
            app._id: [{
                'text': trans(module.name, app.langs),
                'value': module.unique_id,
            } for module in app.modules]
            for app in get_apps_in_domain(self.domain)
        }

    @property
    def module_to_form_options(self):
        return {
            module.unique_id: [{
                'text': trans(form.name, app.langs),
                'value': form.unique_id,
            } for form in module.get_forms()]
            for app in get_apps_in_domain(self.domain)
            for module in app.modules
        }

class CreateCaseExportView(BaseProjectDataView):
    urlname = 'create_case_export'
    page_title = ugettext_noop('Create Case Export')
    template_name = 'export/create_case_export.html'

    @property
    def main_context(self):
        context = super(CreateCaseExportView, self).main_context
        context.update({
            'create_export_form': self.create_export_form,
            'app_to_case_type_options': self.app_to_case_type_options,
            'case_type_prompt': _('Select Case Type...'),
        })
        return context

    def post(self, request, *args, **kwargs):
        if self.create_export_form.is_valid():
            case_type = self.create_export_form.cleaned_data['case_type']
            return HttpResponseRedirect(
                reverse(
                    CreateCustomCaseExportView.urlname,
                    args=[self.domain],
                ) + ('?export_tag="%(export_tag)s"' % {
                    'export_tag': case_type,
                })
            )
        return self.get(self.request, *args, **kwargs)

    @property
    @memoized
    def create_export_form(self):
        if self.request.method == 'POST':
            return CreateCaseExportForm(self.domain, self.request.POST)
        return CreateCaseExportForm(self.domain)

    @property
    def app_to_case_type_options(self):
        return {
            app._id: [{
                'text': case_type,
                'value': case_type,
            } for case_type in set(
                module.case_type for module in app.modules if module.case_type
            )]
            for app in get_apps_in_domain(self.domain)
        }
<|endoftext|>


<commit_before># Copyright 2015 OpenStack Foundation
# Copyright (c) 2015 Hewlett-Packard Development Company, L.P.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

from neutron.db import models_v2
from neutron.plugins.ml2 import driver_api as api

from neutron_lib import exceptions
import sqlalchemy as sa
from sqlalchemy.orm import exc

class L2GatewayCommonDbMixin(object):

    def _apply_filters_to_query(self, query, model, filters):
        """Apply filters to query for the models."""
        if filters:
            for key, value in filters.items():
                column = getattr(model, key, None)
                if column:
                    query = query.filter(column.in_(value))
        return query

    def _model_query(self, context, model):
        """Query model based on filter."""
        query = context.session.query(model)
        query_filter = None
        if not context.is_admin and hasattr(model, 'tenant_id'):
            if hasattr(model, 'shared'):
                query_filter = ((model.tenant_id == context.tenant_id) |
                                (model.shared == sa.true()))
            else:
                query_filter = (model.tenant_id == context.tenant_id)
        if query_filter is not None:
            query = query.filter(query_filter)
        return query

    def _get_collection_query(self, context, model, filters=None,
                              sorts=None, limit=None, marker_obj=None,
                              page_reverse=False):
        """Get collection query for the models."""
        collection = self._model_query(context, model)
        collection = self._apply_filters_to_query(collection, model, filters)
        return collection

    def _get_marker_obj(self, context, resource, limit, marker):
        """Get marker object for the resource."""
        if limit and marker:
            return getattr(self, '_get_%s' % resource)(context, marker)
        return None

    def _fields(self, resource, fields):
        """Get fields for the resource for get query."""
        if fields:
            return dict(((key, item) for key, item in resource.items()
                         if key in fields))
        return resource

    def _get_tenant_id_for_create(self, context, resource):
        """Get tenant id for creation of resources."""
        if context.is_admin and 'tenant_id' in resource:
            tenant_id = resource['tenant_id']
        elif ('tenant_id' in resource and
              resource['tenant_id'] != context.tenant_id):
            reason = _('Cannot create resource for another tenant')
            raise exceptions.AdminRequired(reason=reason)
        else:
            tenant_id = context.tenant_id
        return tenant_id

    def _get_collection(self, context, model, dict_func, filters=None,
                        fields=None, sorts=None, limit=None, marker_obj=None,
                        page_reverse=False):
        """Get collection object based on query for resources."""
        query = self._get_collection_query(context, model, filters=filters,
                                           sorts=sorts,
                                           limit=limit,
                                           marker_obj=marker_obj,
                                           page_reverse=page_reverse)
        items = [dict_func(c, fields) for c in query]
        if limit and page_reverse:
            items.reverse()
        return items

    def _make_segment_dict(self, record):
        """Make a segment dictionary out of a DB record."""
        return {api.ID: record.id,
                api.NETWORK_TYPE: record.network_type,
                api.PHYSICAL_NETWORK: record.physical_network,
                api.SEGMENTATION_ID: record.segmentation_id}

    def _get_network(self, context, id):
        try:
            network = self._get_by_id(context, models_v2.Network, id)
        except exc.NoResultFound:
            raise exceptions.NetworkNotFound(net_id=id)
        return network

    def _get_by_id(self, context, model, id):
        query = self._model_query(context, model)
        return query.filter(model.id == id).one()
<commit_msg>use MechanismDriver from neutron-lib<commit_after># Copyright 2015 OpenStack Foundation
# Copyright (c) 2015 Hewlett-Packard Development Company, L.P.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

from neutron.db import models_v2

from neutron_lib import exceptions
from neutron_lib.plugins.ml2 import api
import sqlalchemy as sa
from sqlalchemy.orm import exc

class L2GatewayCommonDbMixin(object):

    def _apply_filters_to_query(self, query, model, filters):
        """Apply filters to query for the models."""
        if filters:
            for key, value in filters.items():
                column = getattr(model, key, None)
                if column:
                    query = query.filter(column.in_(value))
        return query

    def _model_query(self, context, model):
        """Query model based on filter."""
        query = context.session.query(model)
        query_filter = None
        if not context.is_admin and hasattr(model, 'tenant_id'):
            if hasattr(model, 'shared'):
                query_filter = ((model.tenant_id == context.tenant_id) |
                                (model.shared == sa.true()))
            else:
                query_filter = (model.tenant_id == context.tenant_id)
        if query_filter is not None:
            query = query.filter(query_filter)
        return query

    def _get_collection_query(self, context, model, filters=None,
                              sorts=None, limit=None, marker_obj=None,
                              page_reverse=False):
        """Get collection query for the models."""
        collection = self._model_query(context, model)
        collection = self._apply_filters_to_query(collection, model, filters)
        return collection

    def _get_marker_obj(self, context, resource, limit, marker):
        """Get marker object for the resource."""
        if limit and marker:
            return getattr(self, '_get_%s' % resource)(context, marker)
        return None

    def _fields(self, resource, fields):
        """Get fields for the resource for get query."""
        if fields:
            return dict(((key, item) for key, item in resource.items()
                         if key in fields))
        return resource

    def _get_tenant_id_for_create(self, context, resource):
        """Get tenant id for creation of resources."""
        if context.is_admin and 'tenant_id' in resource:
            tenant_id = resource['tenant_id']
        elif ('tenant_id' in resource and
              resource['tenant_id'] != context.tenant_id):
            reason = _('Cannot create resource for another tenant')
            raise exceptions.AdminRequired(reason=reason)
        else:
            tenant_id = context.tenant_id
        return tenant_id

    def _get_collection(self, context, model, dict_func, filters=None,
                        fields=None, sorts=None, limit=None, marker_obj=None,
                        page_reverse=False):
        """Get collection object based on query for resources."""
        query = self._get_collection_query(context, model, filters=filters,
                                           sorts=sorts,
                                           limit=limit,
                                           marker_obj=marker_obj,
                                           page_reverse=page_reverse)
        items = [dict_func(c, fields) for c in query]
        if limit and page_reverse:
            items.reverse()
        return items

    def _make_segment_dict(self, record):
        """Make a segment dictionary out of a DB record."""
        return {api.ID: record.id,
                api.NETWORK_TYPE: record.network_type,
                api.PHYSICAL_NETWORK: record.physical_network,
                api.SEGMENTATION_ID: record.segmentation_id}

    def _get_network(self, context, id):
        try:
            network = self._get_by_id(context, models_v2.Network, id)
        except exc.NoResultFound:
            raise exceptions.NetworkNotFound(net_id=id)
        return network

    def _get_by_id(self, context, model, id):
        query = self._model_query(context, model)
        return query.filter(model.id == id).one()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# Created: 06.2020
# Copyright (c) 2020, Matthew Broadway
# License: MIT License
import argparse
import signal
import sys
from functools import partial
from typing import Optional

from PyQt5 import QtWidgets as qw, QtCore as qc, QtGui as qg

import ezdxf
from ezdxf.addons.drawing import Frontend, RenderContext
from ezdxf.addons.drawing.properties import is_dark_color
from ezdxf.addons.drawing.pyqt import _get_x_scale, PyQtBackend, CorrespondingDXFEntity, \
    CorrespondingDXFEntityStack
from ezdxf.drawing import Drawing
from ezdxf.lldxf.const import DXFStructureError

class CADGraphicsView(qw.QGraphicsView):
    def __init__(self, view_buffer: float = 0.2):
        super().__init__()
        self._zoom = 1
        self._default_zoom = 1
        self._zoom_limits = (0.5, 100)
        self._view_buffer = view_buffer

        self.setTransformationAnchor(qw.QGraphicsView.AnchorUnderMouse)
        self.setResizeAnchor(qw.QGraphicsView.AnchorUnderMouse)
        self.setVerticalScrollBarPolicy(qc.Qt.ScrollBarAlwaysOff)
        self.setHorizontalScrollBarPolicy(qc.Qt.ScrollBarAlwaysOff)
        self.setDragMode(qw.QGraphicsView.ScrollHandDrag)
        self.setFrameShape(qw.QFrame.NoFrame)
        self.setRenderHints(qg.QPainter.Antialiasing | qg.QPainter.TextAntialiasing | qg.QPainter.SmoothPixmapTransform)

    def clear(self):
        pass

    def fit_to_scene(self):
        r = self.sceneRect()
        bx, by = r.width() * self._view_buffer / 2, r.height() * self._view_buffer / 2
        self.fitInView(self.sceneRect().adjusted(-bx, -by, bx, by), qc.Qt.KeepAspectRatio)
        self._default_zoom = _get_x_scale(self.transform())
        self._zoom = 1

    def _get_zoom_amount(self) -> float:
        return _get_x_scale(self.transform()) / self._default_zoom

    def wheelEvent(self, event: qg.QWheelEvent) -> None:
        # dividing by 120 gets number of notches on a typical scroll wheel. See QWheelEvent documentation
        delta_notches = event.angleDelta().y() / 120
        zoom_per_scroll_notch = 0.2
        factor = 1 + zoom_per_scroll_notch * delta_notches
        resulting_zoom = self._zoom * factor
        if resulting_zoom < self._zoom_limits[0]:
            factor = self._zoom_limits[0] / self._zoom
        elif resulting_zoom > self._zoom_limits[1]:
            factor = self._zoom_limits[1] / self._zoom
        self.scale(factor, factor)
        self._zoom *= factor

class CADGraphicsViewWithOverlay(CADGraphicsView):
    element_selected = qc.pyqtSignal(object, qc.QPointF)

    def __init__(self):
        super().__init__()
        self._current_item: Optional[qw.QGraphicsItem] = None

    def clear(self):
        super().clear()
        self._current_item = None

    def drawForeground(self, painter: qg.QPainter, rect: qc.QRectF) -> None:
        if self._current_item is not None:
            r = self._current_item.boundingRect()
            r = self._current_item.sceneTransform().mapRect(r)
            painter.fillRect(r, qg.QColor(0, 255, 0, 100))

    def mouseMoveEvent(self, event: qg.QMouseEvent) -> None:
        pos = self.mapToScene(event.pos())
        self._current_item = self.scene().itemAt(pos, qg.QTransform())
        self.element_selected.emit(self._current_item, pos)
        self.scene().invalidate(self.sceneRect(), qw.QGraphicsScene.ForegroundLayer)
        super().mouseMoveEvent(event)

class CadViewer(qw.QMainWindow):
    def __init__(self):
        super().__init__()
        self.doc = None
        self._render_context = None
        self._visible_layers = None
        self._current_layout = None

        self.scene = qw.QGraphicsScene()

        self.view = CADGraphicsViewWithOverlay()
        self.view.setScene(self.scene)
        self.view.scale(1, -1)  # so that +y is up
        self.view.element_selected.connect(self._on_element_selected)

        self.renderer = PyQtBackend(self.scene)

        menu = self.menuBar()
        select_doc_action = qw.QAction('Select Document', self)
        select_doc_action.triggered.connect(self._select_doc)
        menu.addAction(select_doc_action)
        self.select_layout_menu = menu.addMenu('Select Layout')

        toggle_sidebar_action = qw.QAction('Toggle Sidebar', self)
        toggle_sidebar_action.triggered.connect(self._toggle_sidebar)
        menu.addAction(toggle_sidebar_action)

        self.sidebar = qw.QSplitter(qc.Qt.Vertical)
        self.layers = qw.QListWidget()
        self.layers.setStyleSheet('font-size: 12pt')
        self.layers.itemChanged.connect(self._layers_updated)
        self.sidebar.addWidget(self.layers)
        self.info = qw.QPlainTextEdit()
        self.info.setReadOnly(True)
        self.sidebar.addWidget(self.info)

        container = qw.QSplitter()
        self.setCentralWidget(container)
        container.addWidget(self.view)
        container.addWidget(self.sidebar)
        container.setCollapsible(0, False)
        container.setCollapsible(1, True)
        w = container.width()
        container.setSizes([int(3 * w / 4), int(w / 4)])

        self.setWindowTitle('CAD Viewer')
        self.resize(1600, 900)
        self.show()

    def _select_doc(self):
        path, _ = qw.QFileDialog.getOpenFileName(self, caption='Select CAD Document', filter='DXF Documents (*.dxf)')
        if path:
            try:
                self.set_document(ezdxf.readfile(path))
            except IOError as e:
                qw.QMessageBox.critical(self, 'Loading Error', str(e))
            except DXFStructureError as e:
                qw.QMessageBox.critical(self, 'DXF Structure Error', f'Invalid DXF file "{path}": {str(e)}')

    def set_document(self, document: Drawing):
        auditor = document.audit()
        error_count = len(auditor.errors)
        if error_count > 0:
            ret = qw.QMessageBox.question(
                self, 'Found DXF Errors',
                f'Found {error_count} errors in file "{document.filename}"\nLoad file anyway? '
            )
            if ret == qw.QMessageBox.No:
                auditor.print_error_report(auditor.errors)
                return
        self.doc = document
        self._render_context = RenderContext(document)
        self._visible_layers = None
        self._current_layout = None
        self._populate_layouts()
        self._populate_layer_list()
        self.draw_layout('Model')
        self.setWindowTitle('CAD Viewer - ' + document.filename)

    def _populate_layer_list(self):
        self.layers.blockSignals(True)
        self.layers.clear()
        for layer in self._render_context.layers.values():
            name = layer.layer
            item = qw.QListWidgetItem(name)
            item.setCheckState(qc.Qt.Checked)
            item.setBackground(qg.QColor(layer.color))
            if is_dark_color(layer.color, 0.4):
                item.setForeground(qg.QColor('#FFFFFF'))
            else:
                item.setForeground(qg.QColor(qg.QColor('#000000')))
            self.layers.addItem(item)
        self.layers.blockSignals(False)

    def _populate_layouts(self):
        self.select_layout_menu.clear()
        for layout_name in self.doc.layout_names_in_taborder():
            action = qw.QAction(layout_name, self)
            action.triggered.connect(partial(self.draw_layout, layout_name))
            self.select_layout_menu.addAction(action)

    def draw_layout(self, layout_name: str):
        print(f'drawing {layout_name}')
        self._current_layout = layout_name
        self.renderer.clear()
        self.view.clear()
        layout = self.doc.layout(layout_name)
        self._update_render_context(layout)
        try:
            Frontend(self._render_context, self.renderer).draw_layout(layout)
        except DXFStructureError as e:
            qw.QMessageBox.critical(self, 'DXF Structure Error', f'Abort rendering of layout "{layout_name}": {str(e)}')
        finally:
            self.renderer.finalize()
        self.view.fit_to_scene()

    def _update_render_context(self, layout):
        assert self._render_context
        self._render_context.set_current_layout(layout)
        # Direct modification of RenderContext.layers would be more flexible, but would also expose the internals.
        if self._visible_layers is not None:
            self._render_context.set_layers_state(self._visible_layers, state=True)

    def resizeEvent(self, event: qg.QResizeEvent) -> None:
        self.view.fit_to_scene()

    @qc.pyqtSlot(qw.QListWidgetItem)
    def _layers_updated(self, _item: qw.QListWidgetItem):
        self._visible_layers = set()
        for i in range(self.layers.count()):
            layer = self.layers.item(i)
            if layer.checkState() == qc.Qt.Checked:
                self._visible_layers.add(layer.text())
        self.draw_layout(self._current_layout)

    @qc.pyqtSlot()
    def _toggle_sidebar(self):
        self.sidebar.setHidden(not self.sidebar.isHidden())

    @qc.pyqtSlot(object, qc.QPointF)
    def _on_element_selected(self, element: Optional[qw.QGraphicsItem], mouse_pos: qc.QPointF):
        text = f'mouse position: {mouse_pos.x():.4f}, {mouse_pos.y():.4f}\n'
        if element is None:
            text += 'No element selected'
        else:
            dxf_entity = element.data(CorrespondingDXFEntity)
            if dxf_entity is None:
                text += 'No data'
            else:
                text += f'Current Entity: {dxf_entity}\nLayer: {dxf_entity.dxf.layer}\n\nDXF Attributes:\n'
                for key, value in dxf_entity.dxf.all_existing_dxf_attribs().items():
                    text += f'- {key}: {value}\n'

                dxf_entity_stack = element.data(CorrespondingDXFEntityStack)
                if dxf_entity_stack:
                    text += '\nParents:\n'
                    for entity in reversed(dxf_entity_stack):
                        text += f'- {entity}\n'

        self.info.setPlainText(text)

def _main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--cad_file')
    parser.add_argument('--layout', default='Model')
    args = parser.parse_args()

    signal.signal(signal.SIGINT, signal.SIG_DFL)  # handle Ctrl+C properly
    app = qw.QApplication(sys.argv)

    v = CadViewer()
    if args.cad_file is not None:
        v.set_document(ezdxf.readfile(args.cad_file))
        try:
            v.draw_layout(args.layout)
        except KeyError:
            print(f'could not find layout "{args.layout}". Valid layouts: {[l.name for l in v.doc.layouts]}')
            sys.exit(1)
    sys.exit(app.exec_())

if __name__ == '__main__':
    _main()
<commit_msg>added dwg support to cad_viewer using the odafc addon<commit_after>#!/usr/bin/env python3
# Created: 06.2020
# Copyright (c) 2020, Matthew Broadway
# License: MIT License
import argparse
import os
import signal
import sys
from functools import partial
from typing import Optional

from PyQt5 import QtWidgets as qw, QtCore as qc, QtGui as qg

import ezdxf
from ezdxf.addons import odafc
from ezdxf.addons.drawing import Frontend, RenderContext
from ezdxf.addons.drawing.properties import is_dark_color
from ezdxf.addons.drawing.pyqt import _get_x_scale, PyQtBackend, CorrespondingDXFEntity, \
    CorrespondingDXFEntityStack
from ezdxf.drawing import Drawing
from ezdxf.lldxf.const import DXFStructureError

class CADGraphicsView(qw.QGraphicsView):
    def __init__(self, view_buffer: float = 0.2):
        super().__init__()
        self._zoom = 1
        self._default_zoom = 1
        self._zoom_limits = (0.5, 100)
        self._view_buffer = view_buffer

        self.setTransformationAnchor(qw.QGraphicsView.AnchorUnderMouse)
        self.setResizeAnchor(qw.QGraphicsView.AnchorUnderMouse)
        self.setVerticalScrollBarPolicy(qc.Qt.ScrollBarAlwaysOff)
        self.setHorizontalScrollBarPolicy(qc.Qt.ScrollBarAlwaysOff)
        self.setDragMode(qw.QGraphicsView.ScrollHandDrag)
        self.setFrameShape(qw.QFrame.NoFrame)
        self.setRenderHints(qg.QPainter.Antialiasing | qg.QPainter.TextAntialiasing | qg.QPainter.SmoothPixmapTransform)

    def clear(self):
        pass

    def fit_to_scene(self):
        r = self.sceneRect()
        bx, by = r.width() * self._view_buffer / 2, r.height() * self._view_buffer / 2
        self.fitInView(self.sceneRect().adjusted(-bx, -by, bx, by), qc.Qt.KeepAspectRatio)
        self._default_zoom = _get_x_scale(self.transform())
        self._zoom = 1

    def _get_zoom_amount(self) -> float:
        return _get_x_scale(self.transform()) / self._default_zoom

    def wheelEvent(self, event: qg.QWheelEvent) -> None:
        # dividing by 120 gets number of notches on a typical scroll wheel. See QWheelEvent documentation
        delta_notches = event.angleDelta().y() / 120
        zoom_per_scroll_notch = 0.2
        factor = 1 + zoom_per_scroll_notch * delta_notches
        resulting_zoom = self._zoom * factor
        if resulting_zoom < self._zoom_limits[0]:
            factor = self._zoom_limits[0] / self._zoom
        elif resulting_zoom > self._zoom_limits[1]:
            factor = self._zoom_limits[1] / self._zoom
        self.scale(factor, factor)
        self._zoom *= factor

class CADGraphicsViewWithOverlay(CADGraphicsView):
    element_selected = qc.pyqtSignal(object, qc.QPointF)

    def __init__(self):
        super().__init__()
        self._current_item: Optional[qw.QGraphicsItem] = None

    def clear(self):
        super().clear()
        self._current_item = None

    def drawForeground(self, painter: qg.QPainter, rect: qc.QRectF) -> None:
        if self._current_item is not None:
            r = self._current_item.boundingRect()
            r = self._current_item.sceneTransform().mapRect(r)
            painter.fillRect(r, qg.QColor(0, 255, 0, 100))

    def mouseMoveEvent(self, event: qg.QMouseEvent) -> None:
        pos = self.mapToScene(event.pos())
        self._current_item = self.scene().itemAt(pos, qg.QTransform())
        self.element_selected.emit(self._current_item, pos)
        self.scene().invalidate(self.sceneRect(), qw.QGraphicsScene.ForegroundLayer)
        super().mouseMoveEvent(event)

class CadViewer(qw.QMainWindow):
    def __init__(self):
        super().__init__()
        self.doc = None
        self._render_context = None
        self._visible_layers = None
        self._current_layout = None

        self.scene = qw.QGraphicsScene()

        self.view = CADGraphicsViewWithOverlay()
        self.view.setScene(self.scene)
        self.view.scale(1, -1)  # so that +y is up
        self.view.element_selected.connect(self._on_element_selected)

        self.renderer = PyQtBackend(self.scene)

        menu = self.menuBar()
        select_doc_action = qw.QAction('Select Document', self)
        select_doc_action.triggered.connect(self._select_doc)
        menu.addAction(select_doc_action)
        self.select_layout_menu = menu.addMenu('Select Layout')

        toggle_sidebar_action = qw.QAction('Toggle Sidebar', self)
        toggle_sidebar_action.triggered.connect(self._toggle_sidebar)
        menu.addAction(toggle_sidebar_action)

        self.sidebar = qw.QSplitter(qc.Qt.Vertical)
        self.layers = qw.QListWidget()
        self.layers.setStyleSheet('font-size: 12pt')
        self.layers.itemChanged.connect(self._layers_updated)
        self.sidebar.addWidget(self.layers)
        self.info = qw.QPlainTextEdit()
        self.info.setReadOnly(True)
        self.sidebar.addWidget(self.info)

        container = qw.QSplitter()
        self.setCentralWidget(container)
        container.addWidget(self.view)
        container.addWidget(self.sidebar)
        container.setCollapsible(0, False)
        container.setCollapsible(1, True)
        w = container.width()
        container.setSizes([int(3 * w / 4), int(w / 4)])

        self.setWindowTitle('CAD Viewer')
        self.resize(1600, 900)
        self.show()

    def _select_doc(self):
        path, _ = qw.QFileDialog.getOpenFileName(self, caption='Select CAD Document',
                                                 filter='DXF Documents(*.dxf);;DWG Documents(*.dwg)')
        if path:
            try:
                if os.path.splitext(path)[1].lower() == '.dwg':
                    doc = odafc.readfile(path)
                else:
                    doc = ezdxf.readfile(path)
                self.set_document(doc)
            except IOError as e:
                qw.QMessageBox.critical(self, 'Loading Error', str(e))
            except DXFStructureError as e:
                qw.QMessageBox.critical(self, 'DXF Structure Error', f'Invalid DXF file "{path}": {str(e)}')

    def set_document(self, document: Drawing):
        auditor = document.audit()
        error_count = len(auditor.errors)
        if error_count > 0:
            ret = qw.QMessageBox.question(
                self, 'Found DXF Errors',
                f'Found {error_count} errors in file "{document.filename}"\nLoad file anyway? '
            )
            if ret == qw.QMessageBox.No:
                auditor.print_error_report(auditor.errors)
                return
        self.doc = document
        self._render_context = RenderContext(document)
        self._visible_layers = None
        self._current_layout = None
        self._populate_layouts()
        self._populate_layer_list()
        self.draw_layout('Model')
        self.setWindowTitle('CAD Viewer - ' + document.filename)

    def _populate_layer_list(self):
        self.layers.blockSignals(True)
        self.layers.clear()
        for layer in self._render_context.layers.values():
            name = layer.layer
            item = qw.QListWidgetItem(name)
            item.setCheckState(qc.Qt.Checked)
            item.setBackground(qg.QColor(layer.color))
            if is_dark_color(layer.color, 0.4):
                item.setForeground(qg.QColor('#FFFFFF'))
            else:
                item.setForeground(qg.QColor(qg.QColor('#000000')))
            self.layers.addItem(item)
        self.layers.blockSignals(False)

    def _populate_layouts(self):
        self.select_layout_menu.clear()
        for layout_name in self.doc.layout_names_in_taborder():
            action = qw.QAction(layout_name, self)
            action.triggered.connect(partial(self.draw_layout, layout_name))
            self.select_layout_menu.addAction(action)

    def draw_layout(self, layout_name: str):
        print(f'drawing {layout_name}')
        self._current_layout = layout_name
        self.renderer.clear()
        self.view.clear()
        layout = self.doc.layout(layout_name)
        self._update_render_context(layout)
        try:
            Frontend(self._render_context, self.renderer).draw_layout(layout)
        except DXFStructureError as e:
            qw.QMessageBox.critical(self, 'DXF Structure Error', f'Abort rendering of layout "{layout_name}": {str(e)}')
        finally:
            self.renderer.finalize()
        self.view.fit_to_scene()

    def _update_render_context(self, layout):
        assert self._render_context
        self._render_context.set_current_layout(layout)
        # Direct modification of RenderContext.layers would be more flexible, but would also expose the internals.
        if self._visible_layers is not None:
            self._render_context.set_layers_state(self._visible_layers, state=True)

    def resizeEvent(self, event: qg.QResizeEvent) -> None:
        self.view.fit_to_scene()

    @qc.pyqtSlot(qw.QListWidgetItem)
    def _layers_updated(self, _item: qw.QListWidgetItem):
        self._visible_layers = set()
        for i in range(self.layers.count()):
            layer = self.layers.item(i)
            if layer.checkState() == qc.Qt.Checked:
                self._visible_layers.add(layer.text())
        self.draw_layout(self._current_layout)

    @qc.pyqtSlot()
    def _toggle_sidebar(self):
        self.sidebar.setHidden(not self.sidebar.isHidden())

    @qc.pyqtSlot(object, qc.QPointF)
    def _on_element_selected(self, element: Optional[qw.QGraphicsItem], mouse_pos: qc.QPointF):
        text = f'mouse position: {mouse_pos.x():.4f}, {mouse_pos.y():.4f}\n'
        if element is None:
            text += 'No element selected'
        else:
            dxf_entity = element.data(CorrespondingDXFEntity)
            if dxf_entity is None:
                text += 'No data'
            else:
                text += f'Current Entity: {dxf_entity}\nLayer: {dxf_entity.dxf.layer}\n\nDXF Attributes:\n'
                for key, value in dxf_entity.dxf.all_existing_dxf_attribs().items():
                    text += f'- {key}: {value}\n'

                dxf_entity_stack = element.data(CorrespondingDXFEntityStack)
                if dxf_entity_stack:
                    text += '\nParents:\n'
                    for entity in reversed(dxf_entity_stack):
                        text += f'- {entity}\n'

        self.info.setPlainText(text)

def _main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--cad_file')
    parser.add_argument('--layout', default='Model')
    args = parser.parse_args()

    signal.signal(signal.SIGINT, signal.SIG_DFL)  # handle Ctrl+C properly
    app = qw.QApplication(sys.argv)

    v = CadViewer()
    if args.cad_file is not None:
        v.set_document(ezdxf.readfile(args.cad_file))
        try:
            v.draw_layout(args.layout)
        except KeyError:
            print(f'could not find layout "{args.layout}". Valid layouts: {[l.name for l in v.doc.layouts]}')
            sys.exit(1)
    sys.exit(app.exec_())

if __name__ == '__main__':
    _main()
<|endoftext|>